| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from functools import partial |
| 4 | |
| 5 | import pytest |
| 6 | |
| 7 | from tests.models.language.pooling.embed_utils import ( |
| 8 | check_embeddings_close, |
| 9 | correctness_test_embed_models, |
| 10 | matryoshka_fy, |
| 11 | ) |
| 12 | from tests.models.utils import ( |
| 13 | EmbedModelInfo, |
| 14 | RerankModelInfo, |
| 15 | ) |
| 16 | from vllm import PoolingParams |
| 17 | from vllm.platforms import current_platform |
| 18 | |
| 19 | from .mteb_embed_utils import mteb_test_embed_models |
| 20 | from .mteb_score_utils import mteb_test_rerank_models |
| 21 | |
| 22 | EMBEDDING_MODELS = [ |
| 23 | EmbedModelInfo( |
| 24 | "jinaai/jina-embeddings-v3", |
| 25 | mteb_score=0.824413164, |
| 26 | mteb_tol=2e-3, |
| 27 | architecture="XLMRobertaModel", |
| 28 | is_matryoshka=True, |
| 29 | seq_pooling_type="MEAN", |
| 30 | attn_type="encoder_only", |
| 31 | is_prefix_caching_supported=False, |
| 32 | is_chunked_prefill_supported=False, |
| 33 | ), |
| 34 | EmbedModelInfo( |
| 35 | "jinaai/jina-embeddings-v5-text-small", |
| 36 | mteb_score=0.794535707854956, |
| 37 | mteb_tol=2e-3, |
| 38 | architecture="JinaEmbeddingsV5Model", |
| 39 | seq_pooling_type="LAST", |
| 40 | attn_type="decoder", |
| 41 | is_prefix_caching_supported=True, |
| 42 | is_chunked_prefill_supported=True, |
| 43 | ), |
| 44 | EmbedModelInfo( |
| 45 | "jinaai/jina-embeddings-v5-text-nano", |
| 46 | architecture="JinaEmbeddingsV5Model", |
| 47 | mteb_tol=2e-3, |
| 48 | dtype="bfloat16" if current_platform.is_rocm() else "auto", |
| 49 | seq_pooling_type="LAST", |
| 50 | attn_type="encoder_only", |
| 51 | is_prefix_caching_supported=False, |
| 52 | is_chunked_prefill_supported=False, |
| 53 | ), |
| 54 | ] |
| 55 | |
| 56 | RERANK_MODELS = [ |
| 57 | RerankModelInfo( |
| 58 | "jinaai/jina-reranker-v2-base-multilingual", |
| 59 | mteb_score=0.33643, |
| 60 | mteb_tol=1e-2, |
| 61 | architecture="XLMRobertaForSequenceClassification", |
| 62 | seq_pooling_type="CLS", |
| 63 | attn_type="encoder_only", |
| 64 | is_prefix_caching_supported=False, |
| 65 | is_chunked_prefill_supported=False, |
| 66 | ) |
| 67 | ] |
| 68 | |
| 69 | |
| 70 | @pytest.mark.flaky(reruns=2) |
| 71 | @pytest.mark.parametrize("model_info", EMBEDDING_MODELS) |
| 72 | def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None: |
| 73 | task = "retrieval" if "v5" in model_info.name else "text-matching" |
| 74 | prompt_prefix: str | None = "Document: " if "v5" in model_info.name else None |
| 75 | |
| 76 | def hf_model_callback(model): |
| 77 | model.encode = partial(model.encode, task=task) |
| 78 | |
| 79 | mteb_test_embed_models( |
| 80 | hf_runner, |
| 81 | vllm_runner, |
| 82 | model_info, |
| 83 | hf_model_callback=hf_model_callback, |
| 84 | prompt_prefix=prompt_prefix, |
| 85 | ) |
| 86 | |
| 87 | |
| 88 | @pytest.mark.parametrize("model_info", EMBEDDING_MODELS) |
| 89 | def test_embed_models_correctness( |
| 90 | hf_runner, vllm_runner, model_info: EmbedModelInfo, example_prompts |
| 91 | ) -> None: |
| 92 | task = "retrieval" if "v5" in model_info.name else "text-matching" |
| 93 | |
| 94 | def hf_model_callback(model): |
| 95 | model.encode = partial(model.encode, task=task) |
| 96 | |
| 97 | correctness_test_embed_models( |
| 98 | hf_runner, |
| 99 | vllm_runner, |
| 100 | model_info, |
| 101 | example_prompts, |
| 102 | hf_model_callback=hf_model_callback, |
| 103 | ) |
| 104 | |
| 105 | |
| 106 | @pytest.mark.flaky(reruns=2) |
| 107 | @pytest.mark.parametrize("model_info", RERANK_MODELS) |
| 108 | def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None: |
| 109 | mteb_test_rerank_models(vllm_runner, model_info) |
| 110 | |
| 111 | |
| 112 | @pytest.mark.skip( |
| 113 | reason="jinaai/jina-embeddings-v3 custom XLMRobertaLoRA model on HF hub " |
| 114 | "is incompatible with transformers v5 (missing all_tied_weights_keys)" |
| 115 | ) |
| 116 | @pytest.mark.parametrize("model_info", EMBEDDING_MODELS) |
| 117 | @pytest.mark.parametrize("dtype", ["half"]) |
| 118 | @pytest.mark.parametrize("dimensions", [16, 32]) |
| 119 | def test_matryoshka( |
| 120 | hf_runner, |
| 121 | vllm_runner, |
| 122 | model_info, |
| 123 | dtype: str, |
| 124 | dimensions: int, |
| 125 | example_prompts, |
| 126 | monkeypatch, |
| 127 | ) -> None: |
| 128 | if not model_info.is_matryoshka: |
| 129 | pytest.skip("Model is not matryoshka") |
| 130 | |
| 131 | # ST will strip the input texts, see test_embedding.py |
| 132 | example_prompts = [str(s).strip() for s in example_prompts] |
| 133 | |
| 134 | task = "retrieval" if "v5" in model_info.name else "text-matching" |
| 135 | |
| 136 | with hf_runner( |
| 137 | model_info.name, |
| 138 | dtype=dtype, |
| 139 | is_sentence_transformer=True, |
| 140 | ) as hf_model: |
| 141 | hf_outputs = hf_model.encode(example_prompts, task=task) |
| 142 | hf_outputs = matryoshka_fy(hf_outputs, dimensions) |
| 143 | |
| 144 | with vllm_runner( |
| 145 | model_info.name, runner="pooling", dtype=dtype, max_model_len=None |
| 146 | ) as vllm_model: |
| 147 | assert vllm_model.llm.llm_engine.model_config.is_matryoshka |
| 148 | |
| 149 | matryoshka_dimensions = ( |
| 150 | vllm_model.llm.llm_engine.model_config.matryoshka_dimensions |
| 151 | ) |
| 152 | assert matryoshka_dimensions is not None |
| 153 | |
| 154 | if dimensions not in matryoshka_dimensions: |
| 155 | with pytest.raises(ValueError): |
| 156 | vllm_model.embed( |
| 157 | example_prompts, pooling_params=PoolingParams(dimensions=dimensions) |
| 158 | ) |
| 159 | else: |
| 160 | vllm_outputs = vllm_model.embed( |
| 161 | example_prompts, pooling_params=PoolingParams(dimensions=dimensions) |
| 162 | ) |
| 163 | |
| 164 | check_embeddings_close( |
| 165 | embeddings_0_lst=hf_outputs, |
| 166 | embeddings_1_lst=vllm_outputs, |
| 167 | name_0="hf", |
| 168 | name_1="vllm", |
| 169 | tol=1e-2, |
| 170 | ) |