| 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 | # 5e-4 is a good tolerance threshold |
| 23 | MTEB_EMBED_TASKS = ["STS12"] |
| 24 | MTEB_EMBED_TOL = 5e-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 HfMtebEncoder(MtebEmbedMixin): |
| 78 | def __init__(self, model): |
| 79 | self.model = model |
| 80 | |
| 81 | def encode( |
| 82 | self, |
| 83 | inputs: DataLoader[mteb.types.BatchedInput], |
| 84 | *args, |
| 85 | **kwargs, |
| 86 | ) -> np.ndarray: |
| 87 | sentences = [text for batch in inputs for text in batch["text"]] |
| 88 | return self.model.encode(sentences) |
| 89 | |
| 90 | |
| 91 | class VllmMtebEncoder(MtebEmbedMixin): |
| 92 | def __init__(self, vllm_model, prompt_prefix: str | None = None): |
| 93 | self.llm = vllm_model |
| 94 | self.rng = np.random.default_rng(seed=42) |
| 95 | self.prompt_prefix = prompt_prefix |
| 96 | |
| 97 | def encode( |
| 98 | self, |
| 99 | inputs: DataLoader[mteb.types.BatchedInput], |
| 100 | *args, |
| 101 | **kwargs, |
| 102 | ) -> np.ndarray: |
| 103 | # Hoping to discover potential scheduling |
| 104 | # issues by randomizing the order. |
| 105 | sentences = [ |
| 106 | self.prompt_prefix + text if self.prompt_prefix else text |
| 107 | for batch in inputs |
| 108 | for text in batch["text"] |
| 109 | ] |
| 110 | r = self.rng.permutation(len(sentences)) |
| 111 | sentences = [sentences[i] for i in r] |
| 112 | outputs = self.llm.embed(sentences, use_tqdm=False) |
| 113 | embeds = np.array(outputs) |
| 114 | embeds = embeds[np.argsort(r)] |
| 115 | return embeds |
| 116 | |
| 117 | |
| 118 | class OpenAIClientMtebEncoder(MtebEmbedMixin): |
| 119 | def __init__(self, model_name: str, client): |
| 120 | self.model_name = model_name |
| 121 | self.client = client |
| 122 | self.rng = np.random.default_rng(seed=42) |
| 123 | |
| 124 | def encode( |
| 125 | self, |
| 126 | inputs: DataLoader[mteb.types.BatchedInput], |
| 127 | *args, |
| 128 | **kwargs, |
| 129 | ) -> np.ndarray: |
| 130 | # Hoping to discover potential scheduling |
| 131 | # issues by randomizing the order. |
| 132 | sentences = [text for batch in inputs for text in batch["text"]] |
| 133 | r = self.rng.permutation(len(sentences)) |
| 134 | sentences = [sentences[i] for i in r] |
| 135 | |
| 136 | embeddings = self.client.embeddings.create( |
| 137 | model=self.model_name, input=sentences |
| 138 | ) |
| 139 | outputs = [d.embedding for d in embeddings.data] |
| 140 | embeds = np.array(outputs) |
| 141 | embeds = embeds[np.argsort(r)] |
| 142 | return embeds |
| 143 | |
| 144 | |
| 145 | def run_mteb_embed_task(encoder: mteb.EncoderProtocol, tasks): |
| 146 | tasks = mteb.get_tasks(tasks=tasks) |
| 147 | results = mteb.evaluate( |
| 148 | encoder, |
| 149 | tasks, |
| 150 | cache=None, |
| 151 | show_progress_bar=False, |
| 152 | ) |
| 153 | |
| 154 | main_score = results[0].scores["test"][0]["main_score"] |
| 155 | return main_score |
| 156 | |
| 157 | |
| 158 | def mteb_test_embed_models( |
| 159 | hf_runner, |
| 160 | vllm_runner, |
| 161 | model_info: EmbedModelInfo, |
| 162 | vllm_extra_kwargs=None, |
| 163 | hf_model_callback=None, |
| 164 | atol: float | None = None, |
| 165 | prompt_prefix: str | None = None, |
| 166 | vllm_model_callback=None, |
| 167 | ): |
| 168 | if atol is None: |
| 169 | atol = ( |
| 170 | model_info.mteb_tol if model_info.mteb_tol is not None else MTEB_EMBED_TOL |
| 171 | ) |
| 172 | vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs) |
| 173 | |
| 174 | # Test embed_dims, isnan and whether to use normalize |
| 175 | example_prompts = ["The chef prepared a delicious meal." * 1000] |
| 176 | |
| 177 | with vllm_runner( |
| 178 | model_info.name, |
| 179 | revision=model_info.revision, |
| 180 | runner="pooling", |
| 181 | max_model_len=model_info.max_model_len, |
| 182 | **vllm_extra_kwargs, |
| 183 | ) as vllm_model: |
| 184 | model_config = vllm_model.llm.llm_engine.model_config |
| 185 | |
| 186 | if vllm_model_callback is not None: |
| 187 | vllm_model_callback(vllm_model) |
| 188 | |
| 189 | # Confirm whether vllm is using the correct architecture |
| 190 | if model_info.architecture: |
| 191 | assert model_info.architecture in model_config.architectures |
| 192 | |
| 193 | # Confirm whether the important configs in model_config are correct. |
| 194 | pooler_config = model_config.pooler_config |
| 195 | if model_info.seq_pooling_type is not None: |
| 196 | assert pooler_config.seq_pooling_type == model_info.seq_pooling_type |
| 197 | if model_info.tok_pooling_type is not None: |
| 198 | assert pooler_config.tok_pooling_type == model_info.tok_pooling_type |
| 199 | if model_info.attn_type is not None: |
| 200 | assert model_config.attn_type == model_info.attn_type |
| 201 | if model_info.is_prefix_caching_supported is not None: |
| 202 | assert ( |
| 203 | model_config.is_prefix_caching_supported |
| 204 | == model_info.is_prefix_caching_supported |
| 205 | ) |
| 206 | if model_info.is_chunked_prefill_supported is not None: |
| 207 | assert ( |
| 208 | model_config.is_chunked_prefill_supported |
| 209 | == model_info.is_chunked_prefill_supported |
| 210 | ) |
| 211 | |
| 212 | vllm_main_score = run_mteb_embed_task( |
| 213 | VllmMtebEncoder(vllm_model, prompt_prefix=prompt_prefix), MTEB_EMBED_TASKS |
| 214 | ) |
| 215 | vllm_dtype = vllm_model.llm.llm_engine.model_config.dtype |
| 216 | head_dtype = model_config.head_dtype |
| 217 | |
| 218 | # Test embedding_size, isnan and whether to use normalize. |
| 219 | # Apply the same prompt_prefix used for scoring so this check compares |
| 220 | # identical effective inputs: HF's SentenceTransformer.encode() applies |
| 221 | # the model's default prompt (e.g. "Document: "), so vLLM must receive it |
| 222 | # too. This mirrors VllmMtebEncoder and changes no production behavior. |
| 223 | consistency_prompts = ( |
| 224 | [prompt_prefix + p for p in example_prompts] |
| 225 | if prompt_prefix |
| 226 | else example_prompts |
| 227 | ) |
| 228 | vllm_outputs = vllm_model.embed( |
| 229 | consistency_prompts, |
| 230 | tokenization_kwargs=dict(truncate_prompt_tokens=-1), |
| 231 | ) |
| 232 | outputs_tensor = torch.tensor(vllm_outputs) |
| 233 | assert not torch.any(torch.isnan(outputs_tensor)) |
| 234 | embedding_size = model_config.embedding_size |
| 235 | assert torch.tensor(vllm_outputs).shape[-1] == embedding_size |
| 236 | |
| 237 | # Accelerate mteb test by setting |
| 238 | # SentenceTransformers mteb score to a constant |
| 239 | if model_info.mteb_score is None: |
| 240 | with hf_runner( |
| 241 | model_info.name, |
| 242 | revision=model_info.revision, |
| 243 | is_sentence_transformer=True, |
| 244 | dtype=ci_envs.VLLM_CI_HF_DTYPE or model_info.hf_dtype, |
| 245 | ) as hf_model: |
| 246 | # e.g. setting default parameters for the encode method of hf_runner |
| 247 | if hf_model_callback is not None: |
| 248 | hf_model_callback(hf_model) |
| 249 | |
| 250 | st_main_score = run_mteb_embed_task( |
| 251 | HfMtebEncoder(hf_model), MTEB_EMBED_TASKS |
| 252 | ) |
| 253 | st_dtype = next(hf_model.model.parameters()).dtype |
| 254 | |
| 255 | # Check embeddings close to hf outputs |
| 256 | hf_outputs = hf_model.encode(example_prompts) |
| 257 | check_embeddings_close( |
| 258 | embeddings_0_lst=hf_outputs, |
| 259 | embeddings_1_lst=vllm_outputs, |
| 260 | name_0="hf", |
| 261 | name_1="vllm", |
| 262 | tol=1e-2, |
| 263 | ) |
| 264 | else: |
| 265 | st_main_score = model_info.mteb_score |
| 266 | st_dtype = "Constant" |
| 267 | |
| 268 | print("Model:", model_info.name) |
| 269 | print("VLLM:", f"dtype:{vllm_dtype}", f"head_dtype:{head_dtype}", vllm_main_score) |
| 270 | print("SentenceTransformers:", st_dtype, st_main_score) |
| 271 | print("Difference:", st_main_score - vllm_main_score) |
| 272 | |
| 273 | # We are not concerned that the vllm mteb results are better |
| 274 | # than SentenceTransformers, so we only perform one-sided testing. |
| 275 | diff = st_main_score - vllm_main_score |
| 276 | assert diff < atol, ( |
| 277 | f"diff={diff:.6g} tol={atol} model={model_info.name} " |
| 278 | f"(st={st_main_score}, vllm={vllm_main_score})" |
| 279 | ) |