From a1eccc762a69b2b4feea75e8c51e11e24b0b4273 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 3 Apr 2024 11:05:59 +0200 Subject: [PATCH] Prefer safetensors to pytorch when both are available --- milli/src/vector/hf.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/milli/src/vector/hf.rs b/milli/src/vector/hf.rs index e341a553e..4c3c0670c 100644 --- a/milli/src/vector/hf.rs +++ b/milli/src/vector/hf.rs @@ -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)? };