Add the code-gemma models. (#2038)

* Add the code-gemma models.

* Tweak to the gemma config.
This commit is contained in:
Laurent Mazare
2024-04-10 21:19:21 +02:00
committed by GitHub
parent b81ecf712d
commit a0460cd2b1
2 changed files with 27 additions and 4 deletions

View File

@ -30,6 +30,14 @@ enum Which {
InstructV1_1_2B,
#[value(name = "1.1-7b-it")]
InstructV1_1_7B,
#[value(name = "code-2b")]
CodeBase2B,
#[value(name = "code-7b")]
CodeBase7B,
#[value(name = "code-2b-it")]
CodeInstruct2B,
#[value(name = "code-7b-it")]
CodeInstruct7B,
}
struct TextGeneration {
@ -224,6 +232,10 @@ fn main() -> Result<()> {
Which::Base7B => "google/gemma-7b".to_string(),
Which::Instruct2B => "google/gemma-2b-it".to_string(),
Which::Instruct7B => "google/gemma-7b-it".to_string(),
Which::CodeBase2B => "google/codegemma-2b".to_string(),
Which::CodeBase7B => "google/codegemma-7b".to_string(),
Which::CodeInstruct2B => "google/codegemma-2b-it".to_string(),
Which::CodeInstruct7B => "google/codegemma-7b-it".to_string(),
},
};
let repo = api.repo(Repo::with_revision(