Add the model argument. (#471)

This commit is contained in:
Laurent Mazare
2023-08-16 16:41:06 +01:00
committed by GitHub
parent 575e88a999
commit 2e206e269d

View File

@ -248,7 +248,7 @@ impl ModelWeights {
struct Args { struct Args {
/// GGML file to load, typically a .bin file generated by the quantize command from llama.cpp /// GGML file to load, typically a .bin file generated by the quantize command from llama.cpp
#[arg(long)] #[arg(long)]
model: String, model: Option<String>,
/// The initial prompt. /// The initial prompt.
#[arg(long)] #[arg(long)]
@ -283,12 +283,24 @@ impl Args {
}; };
Tokenizer::from_file(tokenizer_path).map_err(anyhow::Error::msg) Tokenizer::from_file(tokenizer_path).map_err(anyhow::Error::msg)
} }
fn model(&self) -> anyhow::Result<std::path::PathBuf> {
let model_path = match &self.model {
Some(config) => std::path::PathBuf::from(config),
None => {
let api = hf_hub::api::sync::Api::new()?;
let api = api.model("TheBloke/Llama-2-7B-GGML".to_string());
api.get("llama-2-7b.ggmlv3.q4_0.bin")?
}
};
Ok(model_path)
}
} }
fn main() -> anyhow::Result<()> { fn main() -> anyhow::Result<()> {
let args = Args::parse(); let args = Args::parse();
let mut file = std::fs::File::open(&args.model)?; let mut file = std::fs::File::open(&args.model()?)?;
let start = std::time::Instant::now(); let start = std::time::Instant::now();
let model = Content::read(&mut file)?; let model = Content::read(&mut file)?;