@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1ff67df182
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 10, 22:45 UTC

[CI] Reorganization pooling_mteb_test (#31265)

Signed-off-by: wang.yuqi <noooop@126.com>

wang.yuqi committed Dec 24, 2025 at 23:36 UTC 1ff67df1826a2e6d2d9be684151604c46f04947d
18 files changed +480 -347
tests/entrypoints/pooling/embed/test_correctness_mteb.py
+1 -1
@@ -4,7 +4,7 @@ import os
4
5 import pytest
6
7 -from tests.models.language.pooling_mteb_test.mteb_utils import (
7 +from tests.models.language.pooling_mteb_test.mteb_embed_utils import (
8 MTEB_EMBED_TASKS,
9 MTEB_EMBED_TOL,
10 OpenAIClientMtebEncoder,
tests/entrypoints/pooling/score/test_correctness_mteb.py
+1 -1
@@ -4,7 +4,7 @@ import os
4
5 import pytest
6
7 -from tests.models.language.pooling_mteb_test.mteb_utils import (
7 +from tests.models.language.pooling_mteb_test.mteb_score_utils import (
8 MTEB_RERANK_LANGS,
9 MTEB_RERANK_TASKS,
10 MTEB_RERANK_TOL,
tests/entrypoints/pooling/score/test_utils.py
+4 -5
@@ -202,11 +202,10 @@ class TestGetScorePrompt:
202 tokenization_kwargs,
203 mock_model_no_score_template,
204 ):
205 - # FIXME: Models implementing SupportsScoreTemplate must use their custom
206 - # template implementation by default to preserve existing functionality.
207 - # Attempting to use tokenizer_config.json templates would most likely break
208 - # these models, as often they just inherit the template from the original LLM.
209 - # CLI --chat-template overrides are still supported.
205 + # FIXME: For now, we only apply a template when one is explicitly provided.
206 + # We cannot rely on the tokenizer's chat template because many models
207 + # inherit junk templates from their base LLM, which breaks both the models
208 + # and the tests that use them.
209 with (
210 patch(
211 "vllm.model_executor.model_loader.get_model_cls",
tests/models/language/pooling_mteb_test/mteb_embed_utils.py new
+228
@@ -0,0 +1,228 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import mteb
5 +import numpy as np
6 +import torch
7 +from mteb.models import ModelMeta
8 +from mteb.types import Array
9 +from torch.utils.data import DataLoader
10 +
11 +import tests.ci_envs as ci_envs
12 +from tests.models.utils import (
13 + EmbedModelInfo,
14 + check_embeddings_close,
15 + get_vllm_extra_kwargs,
16 +)
17 +
18 +# Most embedding models on the STS12 task (See #17175):
19 +# - Model implementation and minor changes in tensor dtype
20 +# results in differences less than 1e-4
21 +# - Different model results in differences more than 1e-3
22 +# 1e-4 is a good tolerance threshold
23 +MTEB_EMBED_TASKS = ["STS12"]
24 +MTEB_EMBED_TOL = 1e-4
25 +
26 +
27 +_empty_model_meta = ModelMeta(
28 + loader=None,
29 + name="vllm/model",
30 + revision="1",
31 + release_date=None,
32 + languages=None,
33 + framework=[],
34 + similarity_fn_name=None,
35 + n_parameters=None,
36 + memory_usage_mb=None,
37 + max_tokens=None,
38 + embed_dim=None,
39 + license=None,
40 + open_weights=None,
41 + public_training_code=None,
42 + public_training_data=None,
43 + use_instructions=None,
44 + training_datasets=None,
45 + modalities=["text"], # 'image' can be added to evaluate multimodal models
46 +)
47 +
48 +
49 +class MtebEmbedMixin(mteb.EncoderProtocol):
50 + mteb_model_meta = _empty_model_meta
51 +
52 + def similarity(
53 + self,
54 + embeddings1: np.ndarray,
55 + embeddings2: np.ndarray,
56 + ) -> np.ndarray:
57 + # Cosine similarity
58 + norm1 = np.linalg.norm(embeddings1, axis=1, keepdims=True)
59 + norm2 = np.linalg.norm(embeddings2, axis=1, keepdims=True)
60 + sim = np.dot(embeddings1, embeddings2.T) / (norm1 * norm2.T)
61 + return sim
62 +
63 + def similarity_pairwise(
64 + self,
65 + embeddings1: Array,
66 + embeddings2: Array,
67 + ) -> Array:
68 + # Cosine similarity
69 + norm1 = np.linalg.norm(embeddings1, axis=1, keepdims=True)
70 + norm2 = np.linalg.norm(embeddings2, axis=1, keepdims=True)
71 + sim = np.sum(embeddings1 * embeddings2, axis=1) / (
72 + norm1.flatten() * norm2.flatten()
73 + )
74 + return sim
75 +
76 +
77 +class VllmMtebEncoder(MtebEmbedMixin):
78 + def __init__(self, vllm_model):
79 + self.llm = vllm_model
80 + self.rng = np.random.default_rng(seed=42)
81 +
82 + def encode(
83 + self,
84 + inputs: DataLoader[mteb.types.BatchedInput],
85 + *args,
86 + **kwargs,
87 + ) -> np.ndarray:
88 + # Hoping to discover potential scheduling
89 + # issues by randomizing the order.
90 + sentences = [text for batch in inputs for text in batch["text"]]
91 + r = self.rng.permutation(len(sentences))
92 + sentences = [sentences[i] for i in r]
93 + outputs = self.llm.embed(sentences, use_tqdm=False)
94 + embeds = np.array(outputs)
95 + embeds = embeds[np.argsort(r)]
96 + return embeds
97 +
98 +
99 +class OpenAIClientMtebEncoder(MtebEmbedMixin):
100 + def __init__(self, model_name: str, client):
101 + self.model_name = model_name
102 + self.client = client
103 + self.rng = np.random.default_rng(seed=42)
104 +
105 + def encode(
106 + self,
107 + inputs: DataLoader[mteb.types.BatchedInput],
108 + *args,
109 + **kwargs,
110 + ) -> np.ndarray:
111 + # Hoping to discover potential scheduling
112 + # issues by randomizing the order.
113 + sentences = [text for batch in inputs for text in batch["text"]]
114 + r = self.rng.permutation(len(sentences))
115 + sentences = [sentences[i] for i in r]
116 +
117 + embeddings = self.client.embeddings.create(
118 + model=self.model_name, input=sentences
119 + )
120 + outputs = [d.embedding for d in embeddings.data]
121 + embeds = np.array(outputs)
122 + embeds = embeds[np.argsort(r)]
123 + return embeds
124 +
125 +
126 +def run_mteb_embed_task(encoder: mteb.EncoderProtocol, tasks):
127 + tasks = mteb.get_tasks(tasks=tasks)
128 + results = mteb.evaluate(
129 + encoder,
130 + tasks,
131 + cache=None,
132 + show_progress_bar=False,
133 + )
134 +
135 + main_score = results[0].scores["test"][0]["main_score"]
136 + return main_score
137 +
138 +
139 +def mteb_test_embed_models(
140 + hf_runner,
141 + vllm_runner,
142 + model_info: EmbedModelInfo,
143 + vllm_extra_kwargs=None,
144 + hf_model_callback=None,
145 + atol=MTEB_EMBED_TOL,
146 +):
147 + vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)
148 +
149 + # Test embed_dims, isnan and whether to use normalize
150 + example_prompts = ["The chef prepared a delicious meal." * 1000]
151 +
152 + with vllm_runner(
153 + model_info.name,
154 + runner="pooling",
155 + max_model_len=model_info.max_model_len,
156 + **vllm_extra_kwargs,
157 + ) as vllm_model:
158 + model_config = vllm_model.llm.llm_engine.model_config
159 +
160 + # Confirm whether vllm is using the correct architecture
161 + if model_info.architecture:
162 + assert model_info.architecture in model_config.architectures
163 +
164 + # Confirm whether the important configs in model_config are correct.
165 + if model_info.pooling_type is not None:
166 + assert model_config.pooler_config.pooling_type == model_info.pooling_type
167 + if model_info.attn_type is not None:
168 + assert model_config.attn_type == model_info.attn_type
169 + if model_info.is_prefix_caching_supported is not None:
170 + assert (
171 + model_config.is_prefix_caching_supported
172 + == model_info.is_prefix_caching_supported
173 + )
174 + if model_info.is_chunked_prefill_supported is not None:
175 + assert (
176 + model_config.is_chunked_prefill_supported
177 + == model_info.is_chunked_prefill_supported
178 + )
179 +
180 + vllm_main_score = run_mteb_embed_task(
181 + VllmMtebEncoder(vllm_model), MTEB_EMBED_TASKS
182 + )
183 + vllm_dtype = vllm_model.llm.llm_engine.model_config.dtype
184 + head_dtype = model_config.head_dtype
185 +
186 + # Test embedding_size, isnan and whether to use normalize
187 + vllm_outputs = vllm_model.embed(example_prompts, truncate_prompt_tokens=-1)
188 + outputs_tensor = torch.tensor(vllm_outputs)
189 + assert not torch.any(torch.isnan(outputs_tensor))
190 + embedding_size = model_config.embedding_size
191 + assert torch.tensor(vllm_outputs).shape[-1] == embedding_size
192 +
193 + # Accelerate mteb test by setting
194 + # SentenceTransformers mteb score to a constant
195 + if model_info.mteb_score is None:
196 + with hf_runner(
197 + model_info.name,
198 + is_sentence_transformer=True,
199 + dtype=ci_envs.VLLM_CI_HF_DTYPE or model_info.hf_dtype,
200 + ) as hf_model:
201 + # e.g. setting default parameters for the encode method of hf_runner
202 + if hf_model_callback is not None:
203 + hf_model_callback(hf_model)
204 +
205 + st_main_score = run_mteb_embed_task(hf_model, MTEB_EMBED_TASKS)
206 + st_dtype = next(hf_model.model.parameters()).dtype
207 +
208 + # Check embeddings close to hf outputs
209 + hf_outputs = hf_model.encode(example_prompts)
210 + check_embeddings_close(
211 + embeddings_0_lst=hf_outputs,
212 + embeddings_1_lst=vllm_outputs,
213 + name_0="hf",
214 + name_1="vllm",
215 + tol=1e-2,
216 + )
217 + else:
218 + st_main_score = model_info.mteb_score
219 + st_dtype = "Constant"
220 +
221 + print("Model:", model_info.name)
222 + print("VLLM:", f"dtype:{vllm_dtype}", f"head_dtype:{head_dtype}", vllm_main_score)
223 + print("SentenceTransformers:", st_dtype, st_main_score)
224 + print("Difference:", st_main_score - vllm_main_score)
225 +
226 + # We are not concerned that the vllm mteb results are better
227 + # than SentenceTransformers, so we only perform one-sided testing.
228 + assert st_main_score - vllm_main_score < atol
tests/models/language/pooling_mteb_test/mteb_score_utils.py renamed
+25 -197
@@ -7,37 +7,24 @@ from pathlib import Path
7 import mteb
8 import numpy as np
9 import requests
10 -import torch
10 from mteb.models import ModelMeta
12 -from mteb.types import Array
11 from torch.utils.data import DataLoader
12
15 -import tests.ci_envs as ci_envs
13 from tests.models.utils import (
17 - EmbedModelInfo,
14 RerankModelInfo,
19 - check_embeddings_close,
15 get_vllm_extra_kwargs,
16 )
17
23 -template_home = (
24 - Path(__file__).parent.parent.parent.parent.parent
25 - / "examples/pooling/score/template"
26 -)
27 -
28 -# Most embedding models on the STS12 task (See #17175):
29 -# - Model implementation and minor changes in tensor dtype
30 -# results in differences less than 1e-4
31 -# - Different model results in differences more than 1e-3
32 -# 1e-4 is a good tolerance threshold
33 -MTEB_EMBED_TASKS = ["STS12"]
34 -MTEB_EMBED_TOL = 1e-4
35 -
18 # See #19344
19 MTEB_RERANK_TASKS = ["NFCorpus"]
20 MTEB_RERANK_LANGS = ["eng"]
21 MTEB_RERANK_TOL = 2e-3
22
23 +template_home = (
24 + Path(__file__).parent.parent.parent.parent.parent
25 + / "examples/pooling/score/template"
26 +)
27 +
28 _empty_model_meta = ModelMeta(
29 loader=None,
30 name="vllm/model",
@@ -60,84 +47,11 @@ _empty_model_meta = ModelMeta(
47 )
48
49
63 -class VllmMtebEncoder(mteb.EncoderProtocol):
50 +class MtebCrossEncoderMixin(mteb.CrossEncoderProtocol):
51 mteb_model_meta = _empty_model_meta
52
66 - def __init__(self, vllm_model):
67 - self.llm = vllm_model
68 - self.rng = np.random.default_rng(seed=42)
69 -
70 - def encode(
71 - self,
72 - inputs: DataLoader[mteb.types.BatchedInput],
73 - *args,
74 - **kwargs,
75 - ) -> np.ndarray:
76 - # Hoping to discover potential scheduling
77 - # issues by randomizing the order.
78 - sentences = [text for batch in inputs for text in batch["text"]]
79 - r = self.rng.permutation(len(sentences))
80 - sentences = [sentences[i] for i in r]
81 - outputs = self.llm.embed(sentences, use_tqdm=False)
82 - embeds = np.array(outputs)
83 - embeds = embeds[np.argsort(r)]
84 - return embeds
85 -
86 - def similarity(
87 - self,
88 - embeddings1: np.ndarray,
89 - embeddings2: np.ndarray,
90 - ) -> np.ndarray:
91 - # Cosine similarity
92 - norm1 = np.linalg.norm(embeddings1, axis=1, keepdims=True)
93 - norm2 = np.linalg.norm(embeddings2, axis=1, keepdims=True)
94 - sim = np.dot(embeddings1, embeddings2.T) / (norm1 * norm2.T)
95 - return sim
96 -
97 - def similarity_pairwise(
98 - self,
99 - embeddings1: Array,
100 - embeddings2: Array,
101 - ) -> Array:
102 - # Cosine similarity
103 - norm1 = np.linalg.norm(embeddings1, axis=1, keepdims=True)
104 - norm2 = np.linalg.norm(embeddings2, axis=1, keepdims=True)
105 - sim = np.sum(embeddings1 * embeddings2, axis=1) / (
106 - norm1.flatten() * norm2.flatten()
107 - )
108 - return sim
109 -
110 -
111 -class OpenAIClientMtebEncoder(VllmMtebEncoder):
112 - def __init__(self, model_name: str, client):
113 - self.model_name = model_name
114 - self.client = client
115 - self.rng = np.random.default_rng(seed=42)
116 -
117 - def encode(
118 - self,
119 - inputs: DataLoader[mteb.types.BatchedInput],
120 - *args,
121 - **kwargs,
122 - ) -> np.ndarray:
123 - # Hoping to discover potential scheduling
124 - # issues by randomizing the order.
125 - sentences = [text for batch in inputs for text in batch["text"]]
126 - r = self.rng.permutation(len(sentences))
127 - sentences = [sentences[i] for i in r]
128 -
129 - embeddings = self.client.embeddings.create(
130 - model=self.model_name, input=sentences
131 - )
132 - outputs = [d.embedding for d in embeddings.data]
133 - embeds = np.array(outputs)
134 - embeds = embeds[np.argsort(r)]
135 - return embeds
136 -
137 -
138 -class VllmMtebCrossEncoder(mteb.CrossEncoderProtocol):
139 - mteb_model_meta = _empty_model_meta
53
54 +class VllmMtebCrossEncoder(MtebCrossEncoderMixin):
55 def __init__(self, vllm_model):
56 self.llm = vllm_model
57 self.rng = np.random.default_rng(seed=42)
@@ -164,7 +78,7 @@ class VllmMtebCrossEncoder(mteb.CrossEncoderProtocol):
78 return scores
79
80
167 -class ScoreClientMtebEncoder(mteb.CrossEncoderProtocol):
81 +class ScoreClientMtebEncoder(MtebCrossEncoderMixin):
82 mteb_model_meta = _empty_model_meta
83
84 def __init__(self, model_name: str, url):
@@ -216,102 +130,6 @@ class RerankClientMtebEncoder(ScoreClientMtebEncoder):
130 return response["results"][0]["relevance_score"]
131
132
219 -def run_mteb_embed_task(encoder: mteb.EncoderProtocol, tasks):
220 - tasks = mteb.get_tasks(tasks=tasks)
221 - results = mteb.evaluate(
222 - encoder,
223 - tasks,
224 - cache=None,
225 - show_progress_bar=False,
226 - )
227 -
228 - main_score = results[0].scores["test"][0]["main_score"]
229 - return main_score
230 -
231 -
232 -def mteb_test_embed_models(
233 - hf_runner,
234 - vllm_runner,
235 - model_info: EmbedModelInfo,
236 - vllm_extra_kwargs=None,
237 - hf_model_callback=None,
238 - atol=MTEB_EMBED_TOL,
239 -):
240 - vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)
241 -
242 - # Test embed_dims, isnan and whether to use normalize
243 - example_prompts = ["The chef prepared a delicious meal." * 1000]
244 -
245 - with vllm_runner(
246 - model_info.name,
247 - runner="pooling",
248 - max_model_len=model_info.max_model_len,
249 - **vllm_extra_kwargs,
250 - ) as vllm_model:
251 - model_config = vllm_model.llm.llm_engine.model_config
252 -
253 - # Confirm whether vllm is using the correct architecture
254 - if model_info.architecture:
255 - assert model_info.architecture in model_config.architectures
256 -
257 - # Confirm whether vllm uses the correct default_pooling_type, which
258 - # relates to whether chunked prefill and prefix caching are enabled
259 - assert (
260 - model_config._model_info.default_pooling_type
261 - == model_info.default_pooling_type
262 - )
263 -
264 - vllm_main_score = run_mteb_embed_task(
265 - VllmMtebEncoder(vllm_model), MTEB_EMBED_TASKS
266 - )
267 - vllm_dtype = vllm_model.llm.llm_engine.model_config.dtype
268 - head_dtype = model_config.head_dtype
269 -
270 - # Test embedding_size, isnan and whether to use normalize
271 - vllm_outputs = vllm_model.embed(example_prompts, truncate_prompt_tokens=-1)
272 - outputs_tensor = torch.tensor(vllm_outputs)
273 - assert not torch.any(torch.isnan(outputs_tensor))
274 - embedding_size = model_config.embedding_size
275 - assert torch.tensor(vllm_outputs).shape[-1] == embedding_size
276 -
277 - # Accelerate mteb test by setting
278 - # SentenceTransformers mteb score to a constant
279 - if model_info.mteb_score is None:
280 - with hf_runner(
281 - model_info.name,
282 - is_sentence_transformer=True,
283 - dtype=ci_envs.VLLM_CI_HF_DTYPE or model_info.hf_dtype,
284 - ) as hf_model:
285 - # e.g. setting default parameters for the encode method of hf_runner
286 - if hf_model_callback is not None:
287 - hf_model_callback(hf_model)
288 -
289 - st_main_score = run_mteb_embed_task(hf_model, MTEB_EMBED_TASKS)
290 - st_dtype = next(hf_model.model.parameters()).dtype
291 -
292 - # Check embeddings close to hf outputs
293 - hf_outputs = hf_model.encode(example_prompts)
294 - check_embeddings_close(
295 - embeddings_0_lst=hf_outputs,
296 - embeddings_1_lst=vllm_outputs,
297 - name_0="hf",
298 - name_1="vllm",
299 - tol=1e-2,
300 - )
301 - else:
302 - st_main_score = model_info.mteb_score
303 - st_dtype = "Constant"
304 -
305 - print("Model:", model_info.name)
306 - print("VLLM:", f"dtype:{vllm_dtype}", f"head_dtype:{head_dtype}", vllm_main_score)
307 - print("SentenceTransformers:", st_dtype, st_main_score)
308 - print("Difference:", st_main_score - vllm_main_score)
309 -
310 - # We are not concerned that the vllm mteb results are better
311 - # than SentenceTransformers, so we only perform one-sided testing.
312 - assert st_main_score - vllm_main_score < atol
313 -
314 -
133 def run_mteb_rerank(cross_encoder: mteb.CrossEncoderProtocol, tasks, languages):
134 with tempfile.TemporaryDirectory() as prediction_folder:
135 bm25s = mteb.get_model("bm25s")
@@ -391,18 +209,28 @@ def mteb_test_rerank_models(
209 # Score API is only enabled for num_labels == 1
210 assert model_config.hf_config.num_labels == 1
211
394 - # Confirm whether vllm uses the correct default_pooling_type, which
395 - # relates to whether chunked prefill and prefix caching are enabled
396 - assert (
397 - model_config._model_info.default_pooling_type
398 - == model_info.default_pooling_type
399 - )
400 -
212 + # Maybe load chat_template.
213 chat_template: str | None = None
214 if model_info.chat_template_name is not None:
215 chat_template = (template_home / model_info.chat_template_name).read_text()
216 vllm_model.chat_template = chat_template
217
218 + # Confirm whether the important configs in model_config are correct.
219 + if model_info.pooling_type is not None:
220 + assert model_config.pooler_config.pooling_type == model_info.pooling_type
221 + if model_info.attn_type is not None:
222 + assert model_config.attn_type == model_info.attn_type
223 + if model_info.is_prefix_caching_supported is not None:
224 + assert (
225 + model_config.is_prefix_caching_supported
226 + == model_info.is_prefix_caching_supported
227 + )
228 + if model_info.is_chunked_prefill_supported is not None:
229 + assert (
230 + model_config.is_chunked_prefill_supported
231 + == model_info.is_chunked_prefill_supported
232 + )
233 +
234 vllm_main_score = run_mteb_rerank(
235 vllm_mteb_encoder(vllm_model),
236 tasks=MTEB_RERANK_TASKS,
tests/models/language/pooling_mteb_test/test_baai.py
+36 -32
@@ -4,90 +4,94 @@ import pytest
4
5 from tests.models.language.pooling.embed_utils import correctness_test_embed_models
6 from tests.models.utils import (
7 - CLSPoolingEmbedModelInfo,
8 - CLSPoolingRerankModelInfo,
7 EmbedModelInfo,
10 - LASTPoolingEmbedModelInfo,
8 RerankModelInfo,
9 )
10
14 -from .mteb_utils import mteb_test_embed_models, mteb_test_rerank_models
11 +from .mteb_embed_utils import mteb_test_embed_models
12 +from .mteb_score_utils import mteb_test_rerank_models
13
14 MODELS = [
15 ########## BertModel
18 - CLSPoolingEmbedModelInfo(
16 + EmbedModelInfo(
17 "BAAI/bge-base-en",
18 architecture="BertModel",
19 mteb_score=0.779336792,
20 + pooling_type="CLS",
21 + attn_type="encoder_only",
22 + is_prefix_caching_supported=False,
23 + is_chunked_prefill_supported=False,
24 enable_test=True,
25 ),
24 - CLSPoolingEmbedModelInfo(
25 - "BAAI/bge-base-zh", architecture="BertModel", enable_test=False
26 - ),
27 - CLSPoolingEmbedModelInfo(
28 - "BAAI/bge-small-en", architecture="BertModel", enable_test=False
29 - ),
30 - CLSPoolingEmbedModelInfo(
31 - "BAAI/bge-small-zh", architecture="BertModel", enable_test=False
32 - ),
33 - CLSPoolingEmbedModelInfo(
34 - "BAAI/bge-large-en", architecture="BertModel", enable_test=False
35 - ),
36 - CLSPoolingEmbedModelInfo(
37 - "BAAI/bge-large-zh", architecture="BertModel", enable_test=False
38 - ),
39 - CLSPoolingEmbedModelInfo(
26 + EmbedModelInfo("BAAI/bge-base-zh", architecture="BertModel", enable_test=False),
27 + EmbedModelInfo("BAAI/bge-small-en", architecture="BertModel", enable_test=False),
28 + EmbedModelInfo("BAAI/bge-small-zh", architecture="BertModel", enable_test=False),
29 + EmbedModelInfo("BAAI/bge-large-en", architecture="BertModel", enable_test=False),
30 + EmbedModelInfo("BAAI/bge-large-zh", architecture="BertModel", enable_test=False),
31 + EmbedModelInfo(
32 "BAAI/bge-large-zh-noinstruct", architecture="BertModel", enable_test=False
33 ),
42 - CLSPoolingEmbedModelInfo(
34 + EmbedModelInfo(
35 "BAAI/bge-base-en-v1.5", architecture="BertModel", enable_test=False
36 ),
45 - CLSPoolingEmbedModelInfo(
37 + EmbedModelInfo(
38 "BAAI/bge-base-zh-v1.5", architecture="BertModel", enable_test=False
39 ),
48 - CLSPoolingEmbedModelInfo(
40 + EmbedModelInfo(
41 "BAAI/bge-small-en-v1.5", architecture="BertModel", enable_test=False
42 ),
51 - CLSPoolingEmbedModelInfo(
43 + EmbedModelInfo(
44 "BAAI/bge-small-zh-v1.5", architecture="BertModel", enable_test=False
45 ),
54 - CLSPoolingEmbedModelInfo(
46 + EmbedModelInfo(
47 "BAAI/bge-large-en-v1.5", architecture="BertModel", enable_test=False
48 ),
57 - CLSPoolingEmbedModelInfo(
49 + EmbedModelInfo(
50 "BAAI/bge-large-zh-v1.5", architecture="BertModel", enable_test=False
51 ),
52 ########## XLMRobertaModel
61 - CLSPoolingEmbedModelInfo(
53 + EmbedModelInfo(
54 "BAAI/bge-m3",
55 architecture="XLMRobertaModel",
56 mteb_score=0.787343078,
57 + pooling_type="CLS",
58 + attn_type="encoder_only",
59 + is_prefix_caching_supported=False,
60 + is_chunked_prefill_supported=False,
61 enable_test=True,
62 ),
63 ########## Qwen2Model
68 - LASTPoolingEmbedModelInfo(
64 + EmbedModelInfo(
65 "BAAI/bge-code-v1",
66 architecture="Qwen2Model",
67 mteb_score=0.75724465,
68 dtype="float32",
69 + pooling_type="LAST",
70 + attn_type="decoder",
71 + is_prefix_caching_supported=True,
72 + is_chunked_prefill_supported=True,
73 enable_test=True,
74 ),
75 ]
76
77 RERANK_MODELS = [
78 ########## XLMRobertaForSequenceClassification
79 - CLSPoolingRerankModelInfo(
79 + RerankModelInfo(
80 "BAAI/bge-reranker-base",
81 architecture="XLMRobertaForSequenceClassification",
82 mteb_score=0.32398,
83 + pooling_type="CLS",
84 + attn_type="encoder_only",
85 + is_prefix_caching_supported=False,
86 + is_chunked_prefill_supported=False,
87 enable_test=True,
88 ),
85 - CLSPoolingRerankModelInfo(
89 + RerankModelInfo(
90 "BAAI/bge-reranker-large",
91 architecture="XLMRobertaForSequenceClassification",
92 enable_test=False,
93 ),
90 - CLSPoolingRerankModelInfo(
94 + RerankModelInfo(
95 "BAAI/bge-reranker-v2-m3",
96 architecture="XLMRobertaForSequenceClassification",
97 enable_test=False,
tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py
+8 -6
@@ -9,14 +9,12 @@ 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 (
13 - VllmMtebCrossEncoder,
14 - mteb_test_rerank_models,
15 -)
16 -from tests.models.utils import LASTPoolingRerankModelInfo, RerankModelInfo
12 +from tests.models.utils import RerankModelInfo
13 +
14 +from .mteb_score_utils import VllmMtebCrossEncoder, mteb_test_rerank_models
15
16 RERANK_MODELS = [
19 - LASTPoolingRerankModelInfo(
17 + RerankModelInfo(
18 "BAAI/bge-reranker-v2-gemma",
19 architecture="GemmaForSequenceClassification",
20 mteb_score=0.33757,
@@ -25,6 +23,10 @@ RERANK_MODELS = [
23 "classifier_from_token": ["Yes"],
24 "method": "no_post_processing",
25 },
26 + pooling_type="LAST",
27 + attn_type="decoder",
28 + is_prefix_caching_supported=True,
29 + is_chunked_prefill_supported=True,
30 ),
31 ]
32
tests/models/language/pooling_mteb_test/test_cross_encoder.py
+11 -5
@@ -3,23 +3,29 @@
3 import pytest
4
5 from tests.models.utils import (
6 - CLSPoolingRerankModelInfo,
7 - LASTPoolingRerankModelInfo,
6 RerankModelInfo,
7 )
8
11 -from .mteb_utils import mteb_test_rerank_models
9 +from .mteb_score_utils import mteb_test_rerank_models
10
11 RERANK_MODELS = [
14 - CLSPoolingRerankModelInfo(
12 + RerankModelInfo(
13 "cross-encoder/ms-marco-TinyBERT-L-2-v2",
14 mteb_score=0.32898,
15 architecture="BertForSequenceClassification",
16 + pooling_type="CLS",
17 + attn_type="encoder_only",
18 + is_prefix_caching_supported=False,
19 + is_chunked_prefill_supported=False,
20 ),
19 - LASTPoolingRerankModelInfo(
21 + RerankModelInfo(
22 "tomaarsen/Qwen3-Reranker-0.6B-seq-cls",
23 mteb_score=0.25736,
24 architecture="Qwen3ForSequenceClassification",
25 + pooling_type="LAST",
26 + attn_type="decoder",
27 + is_prefix_caching_supported=True,
28 + is_chunked_prefill_supported=True,
29 ),
30 ]
31
tests/models/language/pooling_mteb_test/test_gte.py
+45 -25
@@ -5,36 +5,32 @@ import pytest
5
6 from tests.models.language.pooling.embed_utils import correctness_test_embed_models
7 from tests.models.utils import (
8 - CLSPoolingEmbedModelInfo,
9 - CLSPoolingRerankModelInfo,
8 EmbedModelInfo,
11 - LASTPoolingEmbedModelInfo,
9 RerankModelInfo,
10 )
11
15 -from .mteb_utils import mteb_test_embed_models, mteb_test_rerank_models
12 +from .mteb_embed_utils import mteb_test_embed_models
13 +from .mteb_score_utils import mteb_test_rerank_models
14
15 MODELS = [
16 ########## BertModel
19 - CLSPoolingEmbedModelInfo(
17 + EmbedModelInfo(
18 "thenlper/gte-large",
19 mteb_score=0.76807651,
20 architecture="BertModel",
21 + pooling_type="MEAN",
22 + attn_type="encoder_only",
23 + is_prefix_caching_supported=False,
24 + is_chunked_prefill_supported=False,
25 enable_test=True,
26 ),
25 - CLSPoolingEmbedModelInfo(
26 - "thenlper/gte-base", architecture="BertModel", enable_test=False
27 - ),
28 - CLSPoolingEmbedModelInfo(
29 - "thenlper/gte-small", architecture="BertModel", enable_test=False
30 - ),
31 - CLSPoolingEmbedModelInfo(
27 + EmbedModelInfo("thenlper/gte-base", architecture="BertModel", enable_test=False),
28 + EmbedModelInfo("thenlper/gte-small", architecture="BertModel", enable_test=False),
29 + EmbedModelInfo(
30 "thenlper/gte-large-zh", architecture="BertModel", enable_test=False
31 ),
34 - CLSPoolingEmbedModelInfo(
35 - "thenlper/gte-base-zh", architecture="BertModel", enable_test=False
36 - ),
37 - CLSPoolingEmbedModelInfo(
32 + EmbedModelInfo("thenlper/gte-base-zh", architecture="BertModel", enable_test=False),
33 + EmbedModelInfo(
34 "thenlper/gte-small-zh", architecture="BertModel", enable_test=False
35 ),
36 ########### NewModel
@@ -43,48 +39,64 @@ MODELS = [
39 # - whether to use token_type_embeddings
40 # - whether to use context expansion
41 # So only test one (the most widely used) model
46 - CLSPoolingEmbedModelInfo(
42 + EmbedModelInfo(
43 "Alibaba-NLP/gte-multilingual-base",
44 architecture="GteNewModel",
45 mteb_score=0.775074696,
46 hf_overrides={"architectures": ["GteNewModel"]},
47 + pooling_type="CLS",
48 + attn_type="encoder_only",
49 + is_prefix_caching_supported=False,
50 + is_chunked_prefill_supported=False,
51 enable_test=True,
52 ),
53 - CLSPoolingEmbedModelInfo(
53 + EmbedModelInfo(
54 "Alibaba-NLP/gte-base-en-v1.5",
55 architecture="GteNewModel",
56 hf_overrides={"architectures": ["GteNewModel"]},
57 enable_test=False,
58 ),
59 - CLSPoolingEmbedModelInfo(
59 + EmbedModelInfo(
60 "Alibaba-NLP/gte-large-en-v1.5",
61 architecture="GteNewModel",
62 hf_overrides={"architectures": ["GteNewModel"]},
63 enable_test=False,
64 ),
65 ########### Qwen2ForCausalLM
66 - LASTPoolingEmbedModelInfo(
66 + EmbedModelInfo(
67 "Alibaba-NLP/gte-Qwen2-1.5B-instruct",
68 mteb_score=0.758473459018872,
69 architecture="Qwen2ForCausalLM",
70 + pooling_type="LAST",
71 + attn_type="encoder_only",
72 + is_prefix_caching_supported=False,
73 + is_chunked_prefill_supported=False,
74 enable_test=True,
75 ),
76 ########## ModernBertModel
73 - CLSPoolingEmbedModelInfo(
77 + EmbedModelInfo(
78 "Alibaba-NLP/gte-modernbert-base",
79 mteb_score=0.748193353,
80 architecture="ModernBertModel",
81 + pooling_type="CLS",
82 + attn_type="encoder_only",
83 + is_prefix_caching_supported=False,
84 + is_chunked_prefill_supported=False,
85 enable_test=True,
86 ),
87 ########## Qwen3ForCausalLM
80 - LASTPoolingEmbedModelInfo(
88 + EmbedModelInfo(
89 "Qwen/Qwen3-Embedding-0.6B",
90 mteb_score=0.771163695,
91 architecture="Qwen3ForCausalLM",
92 dtype="float32",
93 + pooling_type="LAST",
94 + attn_type="decoder",
95 + is_prefix_caching_supported=True,
96 + is_chunked_prefill_supported=True,
97 enable_test=True,
98 ),
87 - LASTPoolingEmbedModelInfo(
99 + EmbedModelInfo(
100 "Qwen/Qwen3-Embedding-4B",
101 architecture="Qwen3ForCausalLM",
102 dtype="float32",
@@ -93,18 +105,26 @@ MODELS = [
105 ]
106
107 RERANK_MODELS = [
96 - CLSPoolingRerankModelInfo(
108 + RerankModelInfo(
109 # classifier_pooling: mean
110 "Alibaba-NLP/gte-reranker-modernbert-base",
111 mteb_score=0.33386,
112 architecture="ModernBertForSequenceClassification",
113 + pooling_type="CLS",
114 + attn_type="encoder_only",
115 + is_prefix_caching_supported=False,
116 + is_chunked_prefill_supported=False,
117 enable_test=True,
118 ),
103 - CLSPoolingRerankModelInfo(
119 + RerankModelInfo(
120 "Alibaba-NLP/gte-multilingual-reranker-base",
121 mteb_score=0.33062,
122 architecture="GteNewForSequenceClassification",
123 hf_overrides={"architectures": ["GteNewForSequenceClassification"]},
124 + pooling_type="CLS",
125 + attn_type="encoder_only",
126 + is_prefix_caching_supported=False,
127 + is_chunked_prefill_supported=False,
128 enable_test=True,
129 ),
130 ]
tests/models/language/pooling_mteb_test/test_intfloat.py
+17 -13
@@ -3,40 +3,44 @@
3 import pytest
4
5 from tests.models.language.pooling.embed_utils import correctness_test_embed_models
6 -from tests.models.utils import CLSPoolingEmbedModelInfo, EmbedModelInfo
6 +from tests.models.utils import EmbedModelInfo
7
8 -from .mteb_utils import mteb_test_embed_models
8 +from .mteb_embed_utils import mteb_test_embed_models
9
10 MODELS = [
11 ########## BertModel
12 - CLSPoolingEmbedModelInfo(
12 + EmbedModelInfo(
13 "intfloat/e5-small",
14 architecture="BertModel",
15 mteb_score=0.742285423,
16 + pooling_type="MEAN",
17 + attn_type="encoder_only",
18 + is_prefix_caching_supported=False,
19 + is_chunked_prefill_supported=False,
20 enable_test=True,
21 ),
18 - CLSPoolingEmbedModelInfo(
19 - "intfloat/e5-base", architecture="BertModel", enable_test=False
20 - ),
21 - CLSPoolingEmbedModelInfo(
22 - "intfloat/e5-large", architecture="BertModel", enable_test=False
23 - ),
24 - CLSPoolingEmbedModelInfo(
22 + EmbedModelInfo("intfloat/e5-base", architecture="BertModel", enable_test=False),
23 + EmbedModelInfo("intfloat/e5-large", architecture="BertModel", enable_test=False),
24 + EmbedModelInfo(
25 "intfloat/multilingual-e5-small", architecture="BertModel", enable_test=False
26 ),
27 ########## XLMRobertaModel
28 - CLSPoolingEmbedModelInfo(
28 + EmbedModelInfo(
29 "intfloat/multilingual-e5-base",
30 architecture="XLMRobertaModel",
31 mteb_score=0.779325955,
32 + pooling_type="MEAN",
33 + attn_type="encoder_only",
34 + is_prefix_caching_supported=False,
35 + is_chunked_prefill_supported=False,
36 enable_test=True,
37 ),
34 - CLSPoolingEmbedModelInfo(
38 + EmbedModelInfo(
39 "intfloat/multilingual-e5-large",
40 architecture="XLMRobertaModel",
41 enable_test=False,
42 ),
39 - CLSPoolingEmbedModelInfo(
43 + EmbedModelInfo(
44 "intfloat/multilingual-e5-large-instruct",
45 architecture="XLMRobertaModel",
46 enable_test=False,
tests/models/language/pooling_mteb_test/test_jina.py
+12 -5
@@ -10,30 +10,37 @@ from tests.models.language.pooling.embed_utils import (
10 matryoshka_fy,
11 )
12 from tests.models.utils import (
13 - CLSPoolingEmbedModelInfo,
14 - CLSPoolingRerankModelInfo,
13 EmbedModelInfo,
14 RerankModelInfo,
15 )
16 from vllm import PoolingParams
17
20 -from .mteb_utils import mteb_test_embed_models, mteb_test_rerank_models
18 +from .mteb_embed_utils import mteb_test_embed_models
19 +from .mteb_score_utils import mteb_test_rerank_models
20
21 EMBEDDING_MODELS = [
23 - CLSPoolingEmbedModelInfo(
22 + EmbedModelInfo(
23 "jinaai/jina-embeddings-v3",
24 mteb_score=0.824413164,
25 architecture="XLMRobertaModel",
26 is_matryoshka=True,
27 + pooling_type="MEAN",
28 + attn_type="encoder_only",
29 + is_prefix_caching_supported=False,
30 + is_chunked_prefill_supported=False,
31 dtype="float32",
32 )
33 ]
34
35 RERANK_MODELS = [
33 - CLSPoolingRerankModelInfo(
36 + RerankModelInfo(
37 "jinaai/jina-reranker-v2-base-multilingual",
38 mteb_score=0.33643,
39 architecture="XLMRobertaForSequenceClassification",
40 + pooling_type="CLS",
41 + attn_type="encoder_only",
42 + is_prefix_caching_supported=False,
43 + is_chunked_prefill_supported=False,
44 )
45 ]
46
tests/models/language/pooling_mteb_test/test_mxbai_rerank.py
+8 -4
@@ -6,9 +6,9 @@ import pytest
6 import torch
7
8 from tests.conftest import HfRunner
9 -from tests.models.utils import LASTPoolingRerankModelInfo, RerankModelInfo
9 +from tests.models.utils import RerankModelInfo
10
11 -from .mteb_utils import mteb_test_rerank_models
11 +from .mteb_score_utils import mteb_test_rerank_models
12
13 mxbai_rerank_hf_overrides = {
14 "architectures": ["Qwen2ForSequenceClassification"],
@@ -17,14 +17,18 @@ mxbai_rerank_hf_overrides = {
17 }
18
19 RERANK_MODELS = [
20 - LASTPoolingRerankModelInfo(
20 + RerankModelInfo(
21 "mixedbread-ai/mxbai-rerank-base-v2",
22 architecture="Qwen2ForSequenceClassification",
23 hf_overrides=mxbai_rerank_hf_overrides,
24 mteb_score=0.273,
25 + pooling_type="LAST",
26 + attn_type="decoder",
27 + is_prefix_caching_supported=True,
28 + is_chunked_prefill_supported=True,
29 enable_test=True,
30 ),
27 - LASTPoolingRerankModelInfo(
31 + RerankModelInfo(
32 "mixedbread-ai/mxbai-rerank-large-v2",
33 architecture="Qwen2ForSequenceClassification",
34 hf_overrides=mxbai_rerank_hf_overrides,
tests/models/language/pooling_mteb_test/test_nemotron.py
+16 -6
@@ -3,29 +3,39 @@
3
4 import pytest
5
6 +from tests.models.language.pooling_mteb_test.mteb_embed_utils import (
7 + mteb_test_embed_models,
8 +)
9 +from tests.models.language.pooling_mteb_test.mteb_score_utils import (
10 + mteb_test_rerank_models,
11 +)
12 from tests.models.utils import (
13 EmbedModelInfo,
8 - LASTPoolingEmbedModelInfo,
9 - LASTPoolingRerankModelInfo,
14 RerankModelInfo,
15 )
16
13 -from .mteb_utils import mteb_test_embed_models, mteb_test_rerank_models
14 -
17 EMBEDDING_MODELS = [
16 - LASTPoolingEmbedModelInfo(
18 + EmbedModelInfo(
19 "nvidia/llama-nemotron-embed-1b-v2",
20 architecture="LlamaBidirectionalModel",
21 mteb_score=0.689164662128673,
22 + pooling_type="MEAN",
23 + attn_type="encoder_only",
24 + is_prefix_caching_supported=False,
25 + is_chunked_prefill_supported=False,
26 )
27 ]
28
29 RERANK_MODELS = [
24 - LASTPoolingRerankModelInfo(
30 + RerankModelInfo(
31 "nvidia/llama-nemotron-rerank-1b-v2",
32 architecture="LlamaBidirectionalForSequenceClassification",
33 chat_template_name="nemotron-rerank.jinja",
34 mteb_score=0.33994,
35 + pooling_type="MEAN",
36 + attn_type="encoder_only",
37 + is_prefix_caching_supported=False,
38 + is_chunked_prefill_supported=False,
39 ),
40 ]
41
tests/models/language/pooling_mteb_test/test_nomic.py
+14 -6
@@ -4,30 +4,38 @@
4 import pytest
5
6 from tests.models.language.pooling.embed_utils import correctness_test_embed_models
7 -from tests.models.utils import CLSPoolingEmbedModelInfo, EmbedModelInfo
7 +from tests.models.utils import EmbedModelInfo
8
9 -from .mteb_utils import mteb_test_embed_models
9 +from .mteb_embed_utils import mteb_test_embed_models
10
11 MODELS = [
12 - CLSPoolingEmbedModelInfo(
12 + EmbedModelInfo(
13 "nomic-ai/nomic-embed-text-v1",
14 architecture="NomicBertModel",
15 mteb_score=0.737568559,
16 enable_test=True,
17 + pooling_type="MEAN",
18 + attn_type="encoder_only",
19 + is_prefix_caching_supported=False,
20 + is_chunked_prefill_supported=False,
21 ),
18 - CLSPoolingEmbedModelInfo(
22 + EmbedModelInfo(
23 "nomic-ai/nomic-embed-text-v1.5",
24 architecture="NomicBertModel",
25 enable_test=False,
26 ),
23 - CLSPoolingEmbedModelInfo(
27 + EmbedModelInfo(
28 "nomic-ai/CodeRankEmbed", architecture="NomicBertModel", enable_test=False
29 ),
26 - CLSPoolingEmbedModelInfo(
30 + EmbedModelInfo(
31 "nomic-ai/nomic-embed-text-v2-moe",
32 architecture="NomicBertModel",
33 mteb_score=0.715488912,
34 enable_test=True,
35 + pooling_type="MEAN",
36 + attn_type="encoder_only",
37 + is_prefix_caching_supported=False,
38 + is_chunked_prefill_supported=False,
39 ),
40 ]
41
tests/models/language/pooling_mteb_test/test_qwen3_reranker.py
+8 -4
@@ -6,10 +6,10 @@ import pytest
6 import torch
7
8 from tests.conftest import HfRunner
9 -from tests.models.utils import LASTPoolingRerankModelInfo, RerankModelInfo
9 +from tests.models.utils import RerankModelInfo
10 from tests.utils import multi_gpu_test
11
12 -from .mteb_utils import mteb_test_rerank_models
12 +from .mteb_score_utils import mteb_test_rerank_models
13
14 qwen3_reranker_hf_overrides = {
15 "architectures": ["Qwen3ForSequenceClassification"],
@@ -18,14 +18,18 @@ qwen3_reranker_hf_overrides = {
18 }
19
20 RERANK_MODELS = [
21 - LASTPoolingRerankModelInfo(
21 + RerankModelInfo(
22 "Qwen/Qwen3-Reranker-0.6B",
23 architecture="Qwen3ForSequenceClassification",
24 mteb_score=0.25736,
25 hf_overrides=qwen3_reranker_hf_overrides,
26 + pooling_type="LAST",
27 + attn_type="decoder",
28 + is_prefix_caching_supported=True,
29 + is_chunked_prefill_supported=True,
30 enable_test=True,
31 ),
28 - LASTPoolingRerankModelInfo(
32 + RerankModelInfo(
33 "Qwen/Qwen3-Reranker-4B",
34 architecture="Qwen3ForSequenceClassification",
35 hf_overrides=qwen3_reranker_hf_overrides,
tests/models/language/pooling_mteb_test/test_snowflake_arctic_embed.py
+30 -10
@@ -4,62 +4,82 @@
4 import pytest
5
6 from tests.models.language.pooling.embed_utils import correctness_test_embed_models
7 -from tests.models.utils import CLSPoolingEmbedModelInfo, EmbedModelInfo
7 +from tests.models.utils import EmbedModelInfo
8
9 -from .mteb_utils import mteb_test_embed_models
9 +from .mteb_embed_utils import mteb_test_embed_models
10
11 MODELS = [
12 - CLSPoolingEmbedModelInfo(
12 + EmbedModelInfo(
13 "Snowflake/snowflake-arctic-embed-xs",
14 is_matryoshka=False,
15 architecture="BertModel",
16 mteb_score=0.714927797,
17 + pooling_type="CLS",
18 + attn_type="encoder_only",
19 + is_prefix_caching_supported=False,
20 + is_chunked_prefill_supported=False,
21 enable_test=True,
22 ),
19 - CLSPoolingEmbedModelInfo(
23 + EmbedModelInfo(
24 "Snowflake/snowflake-arctic-embed-s",
25 is_matryoshka=False,
26 architecture="BertModel",
27 enable_test=False,
28 ),
25 - CLSPoolingEmbedModelInfo(
29 + EmbedModelInfo(
30 "Snowflake/snowflake-arctic-embed-m",
31 is_matryoshka=False,
32 architecture="BertModel",
33 enable_test=False,
34 ),
31 - CLSPoolingEmbedModelInfo(
35 + EmbedModelInfo(
36 "Snowflake/snowflake-arctic-embed-m-long",
37 is_matryoshka=False,
38 architecture="NomicBertModel",
39 mteb_score=0.681146831,
40 + pooling_type="CLS",
41 + attn_type="encoder_only",
42 + is_prefix_caching_supported=False,
43 + is_chunked_prefill_supported=False,
44 enable_test=True,
45 ),
38 - CLSPoolingEmbedModelInfo(
46 + EmbedModelInfo(
47 "Snowflake/snowflake-arctic-embed-l",
48 is_matryoshka=False,
49 architecture="BertModel",
50 enable_test=False,
51 ),
44 - CLSPoolingEmbedModelInfo(
52 + EmbedModelInfo(
53 "Snowflake/snowflake-arctic-embed-m-v1.5",
54 is_matryoshka=True,
55 architecture="BertModel",
56 mteb_score=0.649088363,
57 + pooling_type="CLS",
58 + attn_type="encoder_only",
59 + is_prefix_caching_supported=False,
60 + is_chunked_prefill_supported=False,
61 enable_test=True,
62 ),
51 - CLSPoolingEmbedModelInfo(
63 + EmbedModelInfo(
64 "Snowflake/snowflake-arctic-embed-l-v2.0",
65 is_matryoshka=True,
66 architecture="XLMRobertaModel",
67 mteb_score=0.712258299,
68 + pooling_type="CLS",
69 + attn_type="encoder_only",
70 + is_prefix_caching_supported=False,
71 + is_chunked_prefill_supported=False,
72 enable_test=True,
73 ),
58 - CLSPoolingEmbedModelInfo(
74 + EmbedModelInfo(
75 "Snowflake/snowflake-arctic-embed-m-v2.0",
76 is_matryoshka=True,
77 architecture="GteModel",
78 mteb_score=0.706622444,
79 + pooling_type="CLS",
80 + attn_type="encoder_only",
81 + is_prefix_caching_supported=False,
82 + is_chunked_prefill_supported=False,
83 enable_test=True,
84 ),
85 ]
tests/models/language/pooling_mteb_test/test_st_projector.py
+11 -5
@@ -3,25 +3,31 @@
3 import pytest
4
5 from tests.models.utils import (
6 - CLSPoolingEmbedModelInfo,
6 EmbedModelInfo,
8 - LASTPoolingEmbedModelInfo,
7 )
8
11 -from .mteb_utils import mteb_test_embed_models
9 +from .mteb_embed_utils import mteb_test_embed_models
10
11 # ST models with projector (Dense) layers
12 ST_PROJECTOR_MODELS = [
15 - CLSPoolingEmbedModelInfo(
13 + EmbedModelInfo(
14 "TencentBAC/Conan-embedding-v1",
15 architecture="BertModel",
16 mteb_score=0.688611955,
17 + pooling_type="MEAN",
18 + attn_type="encoder_only",
19 + is_prefix_caching_supported=False,
20 + is_chunked_prefill_supported=False,
21 enable_test=True,
22 ),
21 - LASTPoolingEmbedModelInfo(
23 + EmbedModelInfo(
24 "google/embeddinggemma-300m",
25 architecture="Gemma3TextModel",
26 mteb_score=0.7473819294684156,
27 + pooling_type="MEAN",
28 + attn_type="encoder_only",
29 + is_prefix_caching_supported=False,
30 + is_chunked_prefill_supported=False,
31 enable_test=True,
32 dtype="float32",
33 ),
tests/models/utils.py
+5 -22
@@ -10,7 +10,7 @@ import torch
10 import torch.nn.functional as F
11 from transformers import PretrainedConfig
12
13 -from vllm.config.model import ModelConfig, ModelDType, RunnerOption
13 +from vllm.config.model import AttnTypeStr, ModelConfig, ModelDType, RunnerOption
14 from vllm.logprobs import Logprob, PromptLogprobs, SampleLogprobs
15 from vllm.multimodal.processing import InputProcessingContext
16 from vllm.tokenizers import cached_tokenizer_from_config
@@ -375,7 +375,10 @@ class ModelInfo:
375 max_model_len: int | None = None
376 hf_dtype: str = "float32"
377 hf_overrides: dict[str, Any] | None = None
378 - default_pooling_type: str = ""
378 + pooling_type: str | None = None
379 + attn_type: AttnTypeStr | None = None
380 + is_prefix_caching_supported: bool | None = None
381 + is_chunked_prefill_supported: bool | None = None
382 enable_test: bool = True
383
384
@@ -386,32 +389,12 @@ class EmbedModelInfo(ModelInfo):
389 matryoshka_dimensions: list[int] | None = None
390
391
389 -@dataclass
390 -class CLSPoolingEmbedModelInfo(EmbedModelInfo):
391 - default_pooling_type: str = "CLS"
392 -
393 -
394 -@dataclass
395 -class LASTPoolingEmbedModelInfo(EmbedModelInfo):
396 - default_pooling_type: str = "LAST"
397 -
398 -
392 @dataclass
393 class RerankModelInfo(ModelInfo):
394 mteb_score: float | None = None
395 chat_template_name: str | None = None
396
397
405 -@dataclass
406 -class CLSPoolingRerankModelInfo(RerankModelInfo):
407 - default_pooling_type: str = "CLS"
408 -
409 -
410 -@dataclass
411 -class LASTPoolingRerankModelInfo(RerankModelInfo):
412 - default_pooling_type: str = "LAST"
413 -
414 -
398 @dataclass
399 class GenerateModelInfo(ModelInfo):
400 hf_dtype: str = "auto"