update: codes

Author: kozistr

backends/candle/tests/common.rs

@@ -113,6 +113,8 @@ enum ModuleType {
     Pooling,
     #[serde(rename = "sentence_transformers.models.Transformer")]
     Transformer,
+    #[serde(rename = "sentence_transformers.models.StaticEmbedding")]
+    StaticEmbedding,
 }
 
 #[derive(Deserialize)]
@@ -153,7 +155,11 @@ pub fn download_artifacts(
     };
 
     api_repo.get("config.json")?;
-    api_repo.get("tokenizer.json")?;
+
+    match api_repo.get("tokenizer.json") {
+        Ok(path) => path,
+        Err(_) => api_repo.get("0_StaticEmbedding/tokenizer.json")?,
+    };
 
     let model_files = match download_safetensors(&api_repo) {
         Ok(p) => p,
@@ -203,6 +209,17 @@ fn download_safetensors(api: &ApiRepo) -> Result<Vec<PathBuf>, ApiError> {
         Ok(p) => return Ok(vec![p]),
         Err(err) => tracing::warn!("Could not download `model.safetensors`: {}", err),
     };
+    if let Ok(path) = api.get("model.safetensors") {
+        return Ok(vec![path]);
+    }
+
+    tracing::warn!("Could not download `model.safetensors`");
+    tracing::info!("Downloading `0_StaticEmbedding/model.safetensors`");
+    if let Ok(path) = api.get("0_StaticEmbedding/model.safetensors") {
+        return Ok(vec![path]);
+    }
+
+    tracing::warn!("Could not download `model.safetensors`");
 
     // Sharded weights
     // Download and parse index file
@@ -279,7 +296,16 @@ pub fn cosine_matcher() -> YamlMatcher<SnapshotEmbeddings> {
 pub fn load_tokenizer(model_root: &Path) -> Result<Tokenizer> {
     // Load tokenizer
     let tokenizer_path = model_root.join("tokenizer.json");
-    let mut tokenizer = Tokenizer::from_file(tokenizer_path).expect("tokenizer.json not found");
+    let mut tokenizer = match Tokenizer::from_file(&tokenizer_path) {
+        Ok(t) => t,
+        Err(e) if e.to_string().contains("No such file") || e.to_string().contains("not found") => {
+            let fallback_path = model_root.join("0_StaticEmbedding").join("tokenizer.json");
+            Tokenizer::from_file(&fallback_path)
+                .expect("0_StaticEmbedding/tokenizer.json not found.")
+        }
+        Err(_) => anyhow::bail!("text-embeddings-inference only supports fast tokenizers"),
+    };
+
     // See https://github.com/huggingface/tokenizers/pull/1357
     if let Some(pre_tokenizer) = tokenizer.get_pre_tokenizer() {
         if let PreTokenizerWrapper::Metaspace(m) = pre_tokenizer {

backends/candle/tests/test_static_embedding.rs

@@ -9,8 +9,11 @@ use text_embeddings_backend_core::{Backend, ModelType, Pool};
 #[test]
 #[serial_test::serial]
 fn test_static_embedding() -> Result<()> {
-    let (model_root, _) =
-        download_artifacts("sentence-transformers/static-similarity-mrl-multilingual-v1", None, None)?;
+    let (model_root, _) = download_artifacts(
+        "sentence-transformers/static-similarity-mrl-multilingual-v1",
+        Some("refs/pr/7"),
+        None,
+    )?;
     let tokenizer = load_tokenizer(&model_root)?;
 
     let backend = CandleBackend::new(
@@ -71,8 +74,11 @@ fn test_static_embedding() -> Result<()> {
 #[test]
 #[serial_test::serial]
 fn test_static_embedding_pooled_raw() -> Result<()> {
-    let (model_root, _) =
-        download_artifacts("sentence-transformers/static-similarity-mrl-multilingual-v1", None, None)?;
+    let (model_root, _) = download_artifacts(
+        "sentence-transformers/static-similarity-mrl-multilingual-v1",
+        Some("refs/pr/7"),
+        None,
+    )?;
     let tokenizer = load_tokenizer(&model_root)?;
 
     let backend = CandleBackend::new(
@@ -99,7 +105,11 @@ fn test_static_embedding_pooled_raw() -> Result<()> {
 
     let (pooled_embeddings, raw_embeddings) = sort_embeddings(backend.embed(input_batch)?);
     let pooled_embeddings_batch = SnapshotEmbeddings::from(pooled_embeddings);
-    insta::assert_yaml_snapshot!("static_embedding_batch_pooled", pooled_embeddings_batch, &matcher);
+    insta::assert_yaml_snapshot!(
+        "static_embedding_batch_pooled",
+        pooled_embeddings_batch,
+        &matcher
+    );
 
     let raw_embeddings_batch = SnapshotEmbeddings::from(raw_embeddings);
     insta::assert_yaml_snapshot!("static_embedding_batch_raw", raw_embeddings_batch, &matcher);
@@ -118,7 +128,11 @@ fn test_static_embedding_pooled_raw() -> Result<()> {
 
     let (pooled_embeddings, _) = sort_embeddings(backend.embed(input_single)?);
     let embeddings_single = SnapshotEmbeddings::from(pooled_embeddings);
-    insta::assert_yaml_snapshot!("static_embedding_single_pooled", embeddings_single, &matcher);
+    insta::assert_yaml_snapshot!(
+        "static_embedding_single_pooled",
+        embeddings_single,
+        &matcher
+    );
 
     assert_eq!(pooled_embeddings_batch[0], embeddings_single[0]);
     assert_eq!(pooled_embeddings_batch[2], embeddings_single[0]);