Add test for gte-multilingual-reranker-base in test_gte.rs

Author: alvarobartt

backends/candle/tests/test_gte.rs

@@ -8,43 +8,51 @@ use text_embeddings_backend_core::{Backend, ModelType, Pool};
 
 #[test]
 fn test_gte() -> Result<()> {
-    let model_root = download_artifacts("Alibaba-NLP/gte-base-en-v1.5", None)?;
-    let tokenizer = load_tokenizer(&model_root)?;
-
-    let backend = CandleBackend::new(
-        &model_root,
-        "float32".to_string(),
-        ModelType::Embedding(Pool::Cls),
-    )?;
-
-    let input_batch = batch(
-        vec![
-            tokenizer.encode("What is Deep Learning?", true).unwrap(),
-            tokenizer.encode("Deep Learning is...", true).unwrap(),
-            tokenizer.encode("What is Deep Learning?", true).unwrap(),
-        ],
-        [0, 1, 2].to_vec(),
-        vec![],
-    );
-
-    let matcher = cosine_matcher();
-
-    let (pooled_embeddings, _) = sort_embeddings(backend.embed(input_batch)?);
-    let embeddings_batch = SnapshotEmbeddings::from(pooled_embeddings);
-    insta::assert_yaml_snapshot!("gte_batch", embeddings_batch, &matcher);
-
-    let input_single = batch(
-        vec![tokenizer.encode("What is Deep Learning?", true).unwrap()],
-        [0].to_vec(),
-        vec![],
-    );
-
-    let (pooled_embeddings, _) = sort_embeddings(backend.embed(input_single)?);
-    let embeddings_single = SnapshotEmbeddings::from(pooled_embeddings);
-
-    insta::assert_yaml_snapshot!("gte_single", embeddings_single, &matcher);
-    assert_eq!(embeddings_batch[0], embeddings_single[0]);
-    assert_eq!(embeddings_batch[2], embeddings_single[0]);
+    let model_ids = vec![
+        "Alibaba-NLP/gte-base-en-v1.5",
+        "Alibaba-NLP/gte-multilingual-base", // Included in test due to different safetensors
+                                             // format as it comes with the "new." prefix
+    ];
+
+    for model_id in model_ids {
+        let model_root = download_artifacts(model_id, None)?;
+        let tokenizer = load_tokenizer(&model_root)?;
+
+        let backend = CandleBackend::new(
+            &model_root,
+            "float32".to_string(),
+            ModelType::Embedding(Pool::Cls),
+        )?;
+
+        let input_batch = batch(
+            vec![
+                tokenizer.encode("What is Deep Learning?", true).unwrap(),
+                tokenizer.encode("Deep Learning is...", true).unwrap(),
+                tokenizer.encode("What is Deep Learning?", true).unwrap(),
+            ],
+            [0, 1, 2].to_vec(),
+            vec![],
+        );
+
+        let matcher = cosine_matcher();
+
+        let (pooled_embeddings, _) = sort_embeddings(backend.embed(input_batch)?);
+        let embeddings_batch = SnapshotEmbeddings::from(pooled_embeddings);
+        insta::assert_yaml_snapshot!("gte_batch", embeddings_batch, &matcher);
+
+        let input_single = batch(
+            vec![tokenizer.encode("What is Deep Learning?", true).unwrap()],
+            [0].to_vec(),
+            vec![],
+        );
+
+        let (pooled_embeddings, _) = sort_embeddings(backend.embed(input_single)?);
+        let embeddings_single = SnapshotEmbeddings::from(pooled_embeddings);
+
+        insta::assert_yaml_snapshot!("gte_single", embeddings_single, &matcher);
+        assert_eq!(embeddings_batch[0], embeddings_single[0]);
+        assert_eq!(embeddings_batch[2], embeddings_single[0]);
+    }
 
     Ok(())
 }