@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 71d0ae1c54
Fork 0

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 12, 23:30 UTC

[Misc] Update embedding/cross encoder tests to use `mteb` v2 (#27329)

Signed-off-by: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Signed-off-by: wang.yuqi <noooop@126.com> Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: wang.yuqi <noooop@126.com> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>

Roman Solomatin committed Nov 19, 2025 at 09:28 UTC 71d0ae1c54543689ea7541aa20b9522982b0815e
6 files changed +143 -83
requirements/test.in
+1 -1
@@ -36,7 +36,7 @@ opencv-python-headless >= 4.11.0 # required for video test
36 datamodel_code_generator # required for minicpm3 test
37 # TODO: Use lm-eval[api]==0.4.10 once released
38 lm-eval[api] @ git+https://github.com/EleutherAI/lm-evaluation-harness.git@206b7722158f58c35b7ffcd53b035fdbdda5126d # required for model evaluation test
39 -mteb[bm25s]>=1.38.11, <2 # required for mteb test
39 +mteb[bm25s]>=2, <3 # required for mteb test
40 transformers==4.57.1
41 tokenizers==0.22.0
42 schemathesis>=3.39.15 # Required for openai schema test.
requirements/test.txt
+1 -3
@@ -201,8 +201,6 @@ email-validator==2.2.0
201 # via pydantic
202 encodec==0.1.1
203 # via vocos
204 -eval-type-backport==0.2.2
205 - # via mteb
204 evaluate==0.4.3
205 # via lm-eval
206 fastapi==0.116.1
@@ -490,7 +488,7 @@ msgpack==1.1.0
488 # via
489 # librosa
490 # ray
493 -mteb==1.38.11
491 +mteb==2.1.2
492 # via -r requirements/test.in
493 multidict==6.1.0
494 # via
tests/models/language/pooling_mteb_test/mteb_utils.py
+114 -65
@@ -2,12 +2,14 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import tempfile
5 -from collections.abc import Sequence
5
6 import mteb
7 import numpy as np
8 import requests
9 import torch
10 +from mteb.models import ModelMeta
11 +from mteb.types import Array
12 +from torch.utils.data import DataLoader
13
14 import tests.ci_envs as ci_envs
15 from tests.models.utils import (
@@ -27,24 +29,47 @@ MTEB_EMBED_TOL = 1e-4
29
30 # See #19344
31 MTEB_RERANK_TASKS = ["NFCorpus"]
30 -MTEB_RERANK_LANGS = ["en"]
32 +MTEB_RERANK_LANGS = ["eng"]
33 MTEB_RERANK_TOL = 2e-3
34
35 +_empty_model_meta = ModelMeta(
36 + loader=None,
37 + name="vllm/model",
38 + revision="1",
39 + release_date=None,
40 + languages=None,
41 + framework=[],
42 + similarity_fn_name=None,
43 + n_parameters=None,
44 + memory_usage_mb=None,
45 + max_tokens=None,
46 + embed_dim=None,
47 + license=None,
48 + open_weights=None,
49 + public_training_code=None,
50 + public_training_data=None,
51 + use_instructions=None,
52 + training_datasets=None,
53 + modalities=["text"], # 'image' can be added to evaluate multimodal models
54 +)
55 +
56 +
57 +class VllmMtebEncoder(mteb.EncoderProtocol):
58 + mteb_model_meta = _empty_model_meta
59
34 -class VllmMtebEncoder(mteb.Encoder):
60 def __init__(self, vllm_model):
36 - super().__init__()
61 self.llm = vllm_model
62 self.rng = np.random.default_rng(seed=42)
63
64 def encode(
65 self,
42 - sentences: Sequence[str],
66 + inputs: DataLoader[mteb.types.BatchedInput],
67 *args,
68 **kwargs,
69 ) -> np.ndarray:
70 # Hoping to discover potential scheduling
71 # issues by randomizing the order.
72 + sentences = [text for batch in inputs for text in batch["text"]]
73 r = self.rng.permutation(len(sentences))
74 sentences = [sentences[i] for i in r]
75 outputs = self.llm.embed(sentences, use_tqdm=False)
@@ -52,36 +77,70 @@ class VllmMtebEncoder(mteb.Encoder):
77 embeds = embeds[np.argsort(r)]
78 return embeds
79
80 + def similarity(
81 + self,
82 + embeddings1: np.ndarray,
83 + embeddings2: np.ndarray,
84 + ) -> np.ndarray:
85 + # Cosine similarity
86 + norm1 = np.linalg.norm(embeddings1, axis=1, keepdims=True)
87 + norm2 = np.linalg.norm(embeddings2, axis=1, keepdims=True)
88 + sim = np.dot(embeddings1, embeddings2.T) / (norm1 * norm2.T)
89 + return sim
90 +
91 + def similarity_pairwise(
92 + self,
93 + embeddings1: Array,
94 + embeddings2: Array,
95 + ) -> Array:
96 + # Cosine similarity
97 + norm1 = np.linalg.norm(embeddings1, axis=1, keepdims=True)
98 + norm2 = np.linalg.norm(embeddings2, axis=1, keepdims=True)
99 + sim = np.sum(embeddings1 * embeddings2, axis=1) / (
100 + norm1.flatten() * norm2.flatten()
101 + )
102 + return sim
103 +
104 +
105 +class VllmMtebCrossEncoder(mteb.CrossEncoderProtocol):
106 + mteb_model_meta = _empty_model_meta
107 +
108 + def __init__(self, vllm_model):
109 + self.llm = vllm_model
110 + self.rng = np.random.default_rng(seed=42)
111 +
112 def predict(
113 self,
57 - sentences: list[tuple[str, str, str | None]], # query, corpus, prompt
114 + inputs1: DataLoader[mteb.types.BatchedInput],
115 + inputs2: DataLoader[mteb.types.BatchedInput],
116 *args,
117 **kwargs,
118 ) -> np.ndarray:
61 - r = self.rng.permutation(len(sentences))
62 - sentences = [sentences[i] for i in r]
63 -
64 - queries = [s[0] for s in sentences]
65 - corpus = [s[1] for s in sentences]
119 + queries = [text for batch in inputs1 for text in batch["text"]]
120 + corpus = [text for batch in inputs2 for text in batch["text"]]
121
122 outputs = self.llm.score(
123 queries, corpus, truncate_prompt_tokens=-1, use_tqdm=False
124 )
125 scores = np.array(outputs)
71 - scores = scores[np.argsort(r)]
126 return scores
127
128
75 -class OpenAIClientMtebEncoder(mteb.Encoder):
129 +class OpenAIClientMtebEncoder(VllmMtebEncoder):
130 def __init__(self, model_name: str, client):
77 - super().__init__()
131 self.model_name = model_name
132 self.client = client
133 self.rng = np.random.default_rng(seed=42)
134
82 - def encode(self, sentences: Sequence[str], *args, **kwargs) -> np.ndarray:
135 + def encode(
136 + self,
137 + inputs: DataLoader[mteb.types.BatchedInput],
138 + *args,
139 + **kwargs,
140 + ) -> np.ndarray:
141 # Hoping to discover potential scheduling
142 # issues by randomizing the order.
143 + sentences = [text for batch in inputs for text in batch["text"]]
144 r = self.rng.permutation(len(sentences))
145 sentences = [sentences[i] for i in r]
146
@@ -94,28 +153,29 @@ class OpenAIClientMtebEncoder(mteb.Encoder):
153 return embeds
154
155
97 -class ScoreClientMtebEncoder(mteb.Encoder):
156 +class ScoreClientMtebEncoder(mteb.CrossEncoderProtocol):
157 + mteb_model_meta = _empty_model_meta
158 +
159 def __init__(self, model_name: str, url):
99 - super().__init__()
160 self.model_name = model_name
161 self.url = url
162 self.rng = np.random.default_rng(seed=42)
163
164 def predict(
165 self,
106 - sentences: list[tuple[str, str, str | None]], # query, corpus, prompt
166 + inputs1: DataLoader[mteb.types.BatchedInput],
167 + inputs2: DataLoader[mteb.types.BatchedInput],
168 *args,
169 **kwargs,
170 ) -> np.ndarray:
110 - r = self.rng.permutation(len(sentences))
111 - sentences = [sentences[i] for i in r]
171 + queries = [text for batch in inputs1 for text in batch["text"]]
172 + full_corpus = [text for batch in inputs2 for text in batch["text"]]
173
174 outputs = []
114 - for query, corpus, prompt in sentences:
175 + for query, corpus in zip(queries, full_corpus):
176 outputs.append(self.get_score(query, corpus))
177
178 scores = np.array(outputs)
118 - scores = scores[np.argsort(r)]
179 return scores
180
181 def get_score(self, query, corpus):
@@ -145,16 +205,13 @@ class RerankClientMtebEncoder(ScoreClientMtebEncoder):
205 return response["results"][0]["relevance_score"]
206
207
148 -def run_mteb_embed_task(encoder, tasks):
208 +def run_mteb_embed_task(encoder: mteb.EncoderProtocol, tasks):
209 tasks = mteb.get_tasks(tasks=tasks)
150 - evaluation = mteb.MTEB(tasks=tasks)
151 - results = evaluation.run(
210 + results = mteb.evaluate(
211 encoder,
153 - verbosity=0,
154 - output_folder=None,
155 - encode_kwargs={
156 - "show_progress_bar": False,
157 - },
212 + tasks,
213 + cache=None,
214 + show_progress_bar=False,
215 )
216
217 main_score = results[0].scores["test"][0]["main_score"]
@@ -244,33 +301,39 @@ def mteb_test_embed_models(
301 assert st_main_score - vllm_main_score < atol
302
303
247 -def run_mteb_rerank(cross_encoder, tasks, languages):
248 - with tempfile.TemporaryDirectory() as results_folder:
304 +def run_mteb_rerank(cross_encoder: mteb.CrossEncoderProtocol, tasks, languages):
305 + with tempfile.TemporaryDirectory() as prediction_folder:
306 bm25s = mteb.get_model("bm25s")
250 - tasks = mteb.get_tasks(tasks=tasks, languages=languages)
251 -
252 - subset = "default"
307 eval_splits = ["test"]
308
255 - evaluation = mteb.MTEB(tasks=tasks)
256 - evaluation.run(
309 + mteb_tasks: list[mteb.abstasks.AbsTaskRetrieval] = mteb.get_tasks(
310 + tasks=tasks, languages=languages, eval_splits=eval_splits
311 + )
312 +
313 + mteb.evaluate(
314 bm25s,
258 - verbosity=0,
259 - eval_splits=eval_splits,
260 - save_predictions=True,
261 - output_folder=f"{results_folder}/stage1",
262 - encode_kwargs={"show_progress_bar": False},
315 + mteb_tasks,
316 + prediction_folder=prediction_folder,
317 + show_progress_bar=False,
318 + # don't save results for test runs
319 + cache=None,
320 + overwrite_strategy="always",
321 )
322
265 - results = evaluation.run(
323 + second_stage_tasks = []
324 + for task in mteb_tasks:
325 + second_stage_tasks.append(
326 + task.convert_to_reranking(
327 + prediction_folder,
328 + top_k=10,
329 + )
330 + )
331 +
332 + results = mteb.evaluate(
333 cross_encoder,
267 - verbosity=0,
268 - eval_splits=eval_splits,
269 - top_k=10,
270 - save_predictions=True,
271 - output_folder=f"{results_folder}/stage2",
272 - previous_results=f"{results_folder}/stage1/NFCorpus_{subset}_predictions.json",
273 - encode_kwargs={"show_progress_bar": False},
334 + second_stage_tasks,
335 + show_progress_bar=False,
336 + cache=None,
337 )
338 main_score = results[0].scores["test"][0]["main_score"]
339 return main_score
@@ -280,20 +343,6 @@ def mteb_test_rerank_models_hf(
343 hf_runner, model_name, hf_dtype="float32", hf_model_callback=None
344 ):
345 with hf_runner(model_name, is_cross_encoder=True, dtype=hf_dtype) as hf_model:
283 - original_predict = hf_model.predict
284 -
285 - def _predict(
286 - sentences: list[tuple[str, str, str | None]], # query, corpus, prompt
287 - *args,
288 - **kwargs,
289 - ):
290 - # vllm and st both remove the prompt, fair comparison.
291 - prompts = [(s[0], s[1]) for s in sentences]
292 - return original_predict(prompts, *args, **kwargs, batch_size=8)
293 -
294 - hf_model.predict = _predict
295 - hf_model.original_predict = original_predict
296 -
346 if hf_model_callback is not None:
347 hf_model_callback(hf_model)
348
@@ -310,7 +359,7 @@ def mteb_test_rerank_models(
359 model_info: RerankModelInfo,
360 vllm_extra_kwargs=None,
361 hf_model_callback=None,
313 - vllm_mteb_encoder=VllmMtebEncoder,
362 + vllm_mteb_encoder=VllmMtebCrossEncoder,
363 atol=MTEB_RERANK_TOL,
364 ):
365 vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)
tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py
+21 -10
@@ -2,13 +2,15 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from typing import Any
4
5 +import mteb
6 import numpy as np
7 import pytest
8 import torch
9 +from torch.utils.data import DataLoader
10
11 from tests.conftest import HfRunner
12 from tests.models.language.pooling_mteb_test.mteb_utils import (
11 - VllmMtebEncoder,
13 + VllmMtebCrossEncoder,
14 mteb_test_rerank_models,
15 )
16 from tests.models.utils import LASTPoolingRerankModelInfo, RerankModelInfo
@@ -103,7 +105,7 @@ class GemmaRerankerHfRunner(HfRunner):
105 return torch.Tensor(scores)
106
107
106 -class GemmaMtebEncoder(VllmMtebEncoder):
108 +class GemmaMtebEncoder(VllmMtebCrossEncoder):
109 def __init__(self, *args, **kwargs):
110 super().__init__(*args, **kwargs)
111 self.query_template = "A: {query}\n"
@@ -111,17 +113,26 @@ class GemmaMtebEncoder(VllmMtebEncoder):
113
114 def predict(
115 self,
114 - sentences: list[tuple[str, str, str | None]], # query, corpus, prompt
116 + inputs1: DataLoader[mteb.types.BatchedInput],
117 + inputs2: DataLoader[mteb.types.BatchedInput],
118 *args,
119 **kwargs,
120 ) -> np.ndarray:
118 - _sentences = []
119 - for query, corpus, prompt in sentences:
120 - query = self.query_template.format(query=query)
121 - corpus = self.document_template.format(doc=corpus, prompt=PROMPT)
122 - _sentences.append((query, corpus, prompt))
123 -
124 - return super().predict(_sentences, *args, **kwargs)
121 + queries = [
122 + self.query_template.format(query=text)
123 + for batch in inputs1
124 + for text in batch["text"]
125 + ]
126 + corpus = [
127 + self.document_template.format(doc=text, prompt=PROMPT)
128 + for batch in inputs2
129 + for text in batch["text"]
130 + ]
131 + outputs = self.llm.score(
132 + queries, corpus, truncate_prompt_tokens=-1, use_tqdm=False
133 + )
134 + scores = np.array(outputs)
135 + return scores
136
137
138 @pytest.mark.parametrize("model_info", RERANK_MODELS)
tests/models/language/pooling_mteb_test/test_mxbai_rerank.py
+3 -2
@@ -70,8 +70,9 @@ class MxbaiRerankerHfRunner(HfRunner):
70 return scores
71
72 scores = []
73 - for prompt in prompts:
74 - inputs = process_inputs([prompt])
73 + for query, doc, *_ in prompts:
74 + pairs = [(query, doc)]
75 + inputs = process_inputs(pairs)
76 score = compute_logits(inputs)
77 scores.append(score[0].item())
78 return torch.Tensor(scores)
tests/models/language/pooling_mteb_test/test_qwen3_reranker.py
+3 -2
@@ -72,8 +72,9 @@ class Qwen3RerankerHfRunner(HfRunner):
72 return scores
73
74 scores = []
75 - for prompt in prompts:
76 - inputs = process_inputs([prompt])
75 + for query, doc, *_ in prompts:
76 + pairs = [(query, doc)]
77 + inputs = process_inputs(pairs)
78 score = compute_logits(inputs)
79 scores.append(score[0].item())
80 return torch.Tensor(scores)