Fix canonical vector lookup in test_text_onnx_embeddings.py

qdrant · Mar 12, 2024 · e01c4f1 · e01c4f1
1 parent bc09ed9
commit e01c4f1
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py
@@ -44,8 +44,12 @@ def test_embedding():
         embeddings = np.stack(embeddings, axis=0)
         assert embeddings.shape == (2, dim)
 
-        canonical_vector = CANONICAL_VECTOR_VALUES[model_desc["model"]]
+        canonical_vector = CANONICAL_VECTOR_VALUES.get(model_desc["model"], None)
+        if canonical_vector is None:
+            print(f"Canonical vector not found for {model_desc['model']}")
+            continue
         assert np.allclose(embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3), model_desc["model"]
+
 
 
 @pytest.mark.parametrize(