Prefer safetensors to pytorch when both are available

This commit is contained in:
Louis Dureuil 2024-04-03 11:05:59 +02:00
parent 75f81a0bab
commit a1eccc762a
No known key found for this signature in database
1 changed files with 4 additions and 4 deletions

View File

@ -87,11 +87,11 @@ impl Embedder {
let config = api.get("config.json").map_err(NewEmbedderError::api_get)?;
let tokenizer = api.get("tokenizer.json").map_err(NewEmbedderError::api_get)?;
let (weights, source) = {
api.get("pytorch_model.bin")
.map(|filename| (filename, WeightSource::Pytorch))
api.get("model.safetensors")
.map(|filename| (filename, WeightSource::Safetensors))
.or_else(|_| {
api.get("model.safetensors")
.map(|filename| (filename, WeightSource::Safetensors))
api.get("pytorch_model.bin")
.map(|filename| (filename, WeightSource::Pytorch))
})
.map_err(NewEmbedderError::api_get)?
};