bert_with_rope

Author: noooop

tests/models/language/pooling/test_nomic.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from tests.models.utils import EmbedModelInfo
+from ...utils import EmbedModelInfo, run_embedding_correctness_test
 
 MODELS = [
     EmbedModelInfo("nomic-ai/nomic-embed-text-v1",
@@ -21,11 +21,26 @@
 
 
 @pytest.mark.parametrize("model_info", MODELS)
-def test_models(
-    hf_runner,
-    vllm_runner,
-    model_info: EmbedModelInfo,
-    monkeypatch,
-) -> None:
+def test_models(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     from .mteb_utils import mteb_test_embed_models
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)
+
+
+@pytest.mark.parametrize("model_info", MODELS)
+def test_models_correctness(hf_runner, vllm_runner, model_info: EmbedModelInfo,
+                            example_prompts) -> None:
+    if not model_info.enable_test:
+        pytest.skip("Skipping test.")
+
+    with vllm_runner(model_info.name,
+                     task="embed",
+                     dtype=model_info.dtype,
+                     max_model_len=None) as vllm_model:
+        vllm_outputs = vllm_model.encode(example_prompts)
+
+    with hf_runner(
+            model_info.name,
+            dtype=model_info.dtype,
+            is_sentence_transformer=True,
+    ) as hf_model:
+        run_embedding_correctness_test(hf_model, example_prompts, vllm_outputs)

tests/models/language/pooling/test_snowflake_arctic_embed.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from tests.models.utils import EmbedModelInfo
+from ...utils import EmbedModelInfo, run_embedding_correctness_test
 
 MODELS = [
     EmbedModelInfo("Snowflake/snowflake-arctic-embed-xs",
@@ -41,11 +41,34 @@
 
 
 @pytest.mark.parametrize("model_info", MODELS)
-def test_models(
+def test_models_mteb(
     hf_runner,
     vllm_runner,
     model_info: EmbedModelInfo,
-    monkeypatch,
 ) -> None:
     from .mteb_utils import mteb_test_embed_models
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)
+
+
+@pytest.mark.parametrize("model_info", MODELS)
+def test_models_correctness(
+    hf_runner,
+    vllm_runner,
+    model_info: EmbedModelInfo,
+    example_prompts,
+) -> None:
+    if not model_info.enable_test:
+        pytest.skip("Skipping test.")
+
+    with vllm_runner(model_info.name,
+                     task="embed",
+                     dtype=model_info.dtype,
+                     max_model_len=None) as vllm_model:
+        vllm_outputs = vllm_model.encode(example_prompts)
+
+    with hf_runner(
+            model_info.name,
+            dtype=model_info.dtype,
+            is_sentence_transformer=True,
+    ) as hf_model:
+        run_embedding_correctness_test(hf_model, example_prompts, vllm_outputs)