Helper function to load sharded safetensors files (#1481)

* Fix the quantized mistral example.

* Add a helper function to load sharded safetensors weights.

* Use the sharded loader.
This commit is contained in:
Laurent Mazare
2023-12-25 21:49:21 +01:00
committed by GitHub
parent eae3a20d43
commit 37c539f2b7
7 changed files with 40 additions and 67 deletions

View File

@ -278,10 +278,10 @@ fn main() -> Result<()> {
} else {
match args.model {
WhichModel::V1 | WhichModel::V1_5 => vec![repo.get("model.safetensors")?],
WhichModel::V2 => vec![
repo.get("model-00001-of-00002.safetensors")?,
repo.get("model-00002-of-00002.safetensors")?,
],
WhichModel::V2 => candle_examples::hub_load_safetensors(
&repo,
"model.safetensors.index.json",
)?,
WhichModel::PuffinPhiV2 => vec![repo.get("model-puffin-phi-v2.safetensors")?],
WhichModel::PhiHermes => vec![repo.get("model-phi-hermes-1_3B.safetensors")?],
}