A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 09:45 UTC
[CI] Add mteb testing to test the accuracy of the embedding model (#17175)
wang.yuqi committed
May 20, 2025 at 21:51 UTC
86847700d74b227b3717aeb9a8bcfcdc352697bf
6 files changed
+64
-5
requirements/test.in
+1
@@ -33,6 +33,7 @@ num2words # required for smolvlm test
33
opencv-python-headless >= 4.11.0 # required for video test
34
datamodel_code_generator # required for minicpm3 test
35
lm-eval[api]==0.4.8 # required for model evaluation test
36
+mteb>=1.38.11, <2 # required for mteb test
37
transformers==4.51.3
38
tokenizers==0.21.1
39
huggingface-hub[hf_xet]>=0.30.0 # Required for Xet downloads.
requirements/test.txt
+21
-1
@@ -99,6 +99,7 @@ datasets==3.0.2
99
# via
100
# evaluate
101
# lm-eval
102
+ # mteb
103
decorator==5.1.1
104
# via librosa
105
dill==0.3.8
@@ -124,6 +125,8 @@ email-validator==2.2.0
125
# via pydantic
126
encodec==0.1.1
127
# via vocos
128
+eval-type-backport==0.2.2
129
+ # via mteb
130
evaluate==0.4.3
131
# via lm-eval
132
fastparquet==2024.11.0
@@ -291,6 +294,8 @@ msgpack==1.1.0
294
# via
295
# librosa
296
# ray
297
+mteb==1.38.11
298
+ # via -r requirements/test.in
299
multidict==6.1.0
300
# via
301
# aiohttp
@@ -331,6 +336,7 @@ numpy==1.26.4
336
# librosa
337
# matplotlib
338
# mistral-common
339
+ # mteb
340
# numba
341
# numexpr
342
# opencv-python-headless
@@ -443,6 +449,8 @@ plotly==5.24.1
449
# via genai-perf
450
pluggy==1.5.0
451
# via pytest
452
+polars==1.29.0
453
+ # via mteb
454
pooch==1.8.2
455
# via librosa
456
portalocker==2.10.1
@@ -476,6 +484,7 @@ pydantic==2.9.2
484
# via
485
# datamodel-code-generator
486
# mistral-common
487
+ # mteb
488
pydantic-core==2.23.4
489
# via pydantic
490
pygments==2.18.0
@@ -522,6 +531,8 @@ python-dateutil==2.9.0.post0
531
# typepy
532
python-rapidjson==1.20
533
# via tritonclient
534
+pytrec-eval-terrier==0.5.7
535
+ # via mteb
536
pytz==2024.2
537
# via
538
# pandas
@@ -564,6 +575,7 @@ requests==2.32.3
575
# huggingface-hub
576
# lm-eval
577
# mistral-common
578
+ # mteb
579
# pooch
580
# ray
581
# responses
@@ -580,6 +592,7 @@ rfc3987==1.3.8
592
rich==13.9.4
593
# via
594
# genai-perf
595
+ # mteb
596
# typer
597
rouge-score==0.1.2
598
# via lm-eval
@@ -607,16 +620,20 @@ scikit-learn==1.5.2
620
# via
621
# librosa
622
# lm-eval
623
+ # mteb
624
# sentence-transformers
625
scipy==1.13.1
626
# via
627
# librosa
628
+ # mteb
629
# scikit-learn
630
# sentence-transformers
631
# statsmodels
632
# vocos
633
sentence-transformers==3.2.1
619
- # via -r requirements/test.in
634
+ # via
635
+ # -r requirements/test.in
636
+ # mteb
637
sentencepiece==0.2.0
638
# via mistral-common
639
setuptools==77.0.3
@@ -696,6 +713,7 @@ torch==2.7.0+cu128
713
# fastsafetensors
714
# lm-eval
715
# mamba-ssm
716
+ # mteb
717
# peft
718
# runai-model-streamer
719
# sentence-transformers
@@ -720,6 +738,7 @@ tqdm==4.66.6
738
# evaluate
739
# huggingface-hub
740
# lm-eval
741
+ # mteb
742
# nltk
743
# peft
744
# pqdm
@@ -759,6 +778,7 @@ typing-extensions==4.12.2
778
# huggingface-hub
779
# librosa
780
# mistral-common
781
+ # mteb
782
# pqdm
783
# pydantic
784
# pydantic-core
tests/entrypoints/openai/correctness/test_mteb.py
new
+42
@@ -0,0 +1,42 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import math
3
+import os
4
+
5
+import pytest
6
+
7
+from tests.models.language.pooling.mteb_utils import (MTEB_EMBED_TASKS,
8
+ OpenAIClientMtebEncoder,
9
+ run_mteb_embed_task,
10
+ run_mteb_embed_task_st)
11
+from tests.utils import RemoteOpenAIServer
12
+
13
+os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
14
+
15
+MODEL_NAME = "BAAI/bge-m3"
16
+DTYPE = "float16"
17
+MAIN_SCORE = 0.7873427091972599
18
+
19
+
20
+@pytest.fixture(scope="module")
21
+def server():
22
+ args = [
23
+ "--task", "embed", "--dtype", DTYPE, "--enforce-eager",
24
+ "--max-model-len", "512"
25
+ ]
26
+
27
+ with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
28
+ yield remote_server
29
+
30
+
31
+def test_mteb(server):
32
+ client = server.get_client()
33
+ encoder = OpenAIClientMtebEncoder(MODEL_NAME, client)
34
+ vllm_main_score = run_mteb_embed_task(encoder, MTEB_EMBED_TASKS)
35
+ st_main_score = MAIN_SCORE or run_mteb_embed_task_st(
36
+ MODEL_NAME, MTEB_EMBED_TASKS)
37
+
38
+ print("VLLM main score: ", vllm_main_score)
39
+ print("SentenceTransformer main score: ", st_main_score)
40
+ print("Difference: ", st_main_score - vllm_main_score)
41
+
42
+ assert math.isclose(st_main_score, vllm_main_score, rel_tol=1e-4)
tests/models/language/pooling/test_gte.py
-2
@@ -58,8 +58,6 @@ MODELS = [
58
@pytest.mark.parametrize("model_info", MODELS)
59
def test_models_mteb(hf_runner, vllm_runner,
60
model_info: EmbedModelInfo) -> None:
61
- pytest.skip("Skipping mteb test.")
62
-
61
from .mteb_utils import mteb_test_embed_models
62
63
vllm_extra_kwargs: dict[str, Any] = {}
tests/models/language/pooling/test_nomic.py
-1
@@ -23,7 +23,6 @@ MODELS = [
23
@pytest.mark.parametrize("model_info", MODELS)
24
def test_models_mteb(hf_runner, vllm_runner,
25
model_info: EmbedModelInfo) -> None:
26
- pytest.skip("Skipping mteb test.")
26
from .mteb_utils import mteb_test_embed_models
27
mteb_test_embed_models(hf_runner, vllm_runner, model_info)
28
tests/models/language/pooling/test_snowflake_arctic_embed.py
-1
@@ -46,7 +46,6 @@ def test_models_mteb(
46
vllm_runner,
47
model_info: EmbedModelInfo,
48
) -> None:
49
- pytest.skip("Skipping mteb test.")
49
from .mteb_utils import mteb_test_embed_models
50
mteb_test_embed_models(hf_runner, vllm_runner, model_info)
51