@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 160 lines 5.22 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5
6 from tests.models.language.pooling.embed_utils import correctness_test_embed_models
7 from tests.models.utils import (
8 EmbedModelInfo,
9 RerankModelInfo,
10 )
11 from vllm.platforms import current_platform
12
13 from .mteb_embed_utils import mteb_test_embed_models
14 from .mteb_score_utils import mteb_test_rerank_models
15
16 MODELS = [
17 ########## BertModel
18 EmbedModelInfo(
19 "thenlper/gte-large",
20 mteb_score=0.76807651,
21 mteb_tol=2e-3,
22 architecture="BertModel",
23 seq_pooling_type="MEAN",
24 attn_type="encoder_only",
25 is_prefix_caching_supported=False,
26 is_chunked_prefill_supported=False,
27 enable_test=True,
28 ),
29 EmbedModelInfo("thenlper/gte-base", architecture="BertModel", enable_test=False),
30 EmbedModelInfo("thenlper/gte-small", architecture="BertModel", enable_test=False),
31 EmbedModelInfo(
32 "thenlper/gte-large-zh", architecture="BertModel", enable_test=False
33 ),
34 EmbedModelInfo("thenlper/gte-base-zh", architecture="BertModel", enable_test=False),
35 EmbedModelInfo(
36 "thenlper/gte-small-zh", architecture="BertModel", enable_test=False
37 ),
38 ########### NewModel
39 # These three architectures are almost the same, but not exactly the same.
40 # For example,
41 # - whether to use token_type_embeddings
42 # - whether to use context expansion
43 # So only test one (the most widely used) model
44 EmbedModelInfo(
45 "Alibaba-NLP/gte-multilingual-base",
46 architecture="GteNewModel",
47 mteb_score=0.775074696,
48 mteb_tol=2e-3,
49 hf_overrides={"architectures": ["GteNewModel"]},
50 seq_pooling_type="CLS",
51 attn_type="encoder_only",
52 is_prefix_caching_supported=False,
53 is_chunked_prefill_supported=False,
54 enable_test=True,
55 ),
56 EmbedModelInfo(
57 "Alibaba-NLP/gte-base-en-v1.5",
58 architecture="GteNewModel",
59 hf_overrides={"architectures": ["GteNewModel"]},
60 enable_test=False,
61 ),
62 EmbedModelInfo(
63 "Alibaba-NLP/gte-large-en-v1.5",
64 architecture="GteNewModel",
65 hf_overrides={"architectures": ["GteNewModel"]},
66 enable_test=False,
67 ),
68 ########### Qwen2ForCausalLM
69 EmbedModelInfo(
70 "Alibaba-NLP/gte-Qwen2-1.5B-instruct",
71 mteb_score=0.758473459018872,
72 architecture="Qwen2ForCausalLM",
73 seq_pooling_type="LAST",
74 attn_type="encoder_only",
75 is_prefix_caching_supported=False,
76 is_chunked_prefill_supported=False,
77 # Skip: numerical regression with transformers v5.
78 enable_test=False,
79 ),
80 ########## ModernBertModel
81 EmbedModelInfo(
82 "Alibaba-NLP/gte-modernbert-base",
83 mteb_score=0.748193353,
84 mteb_tol=2e-3,
85 architecture="ModernBertModel",
86 seq_pooling_type="CLS",
87 attn_type="encoder_only",
88 is_prefix_caching_supported=False,
89 is_chunked_prefill_supported=False,
90 enable_test=True,
91 ),
92 ########## Qwen3ForCausalLM
93 EmbedModelInfo(
94 "Qwen/Qwen3-Embedding-0.6B",
95 mteb_score=0.771163695,
96 mteb_tol=2e-3,
97 architecture="Qwen3ForCausalLM",
98 seq_pooling_type="LAST",
99 attn_type="decoder",
100 is_prefix_caching_supported=True,
101 is_chunked_prefill_supported=True,
102 enable_test=True,
103 ),
104 EmbedModelInfo(
105 "Qwen/Qwen3-Embedding-4B",
106 architecture="Qwen3ForCausalLM",
107 enable_test=False,
108 ),
109 ]
110
111 RERANK_MODELS = [
112 RerankModelInfo(
113 # classifier_pooling: mean
114 "Alibaba-NLP/gte-reranker-modernbert-base",
115 mteb_score=0.33386,
116 mteb_tol=1e-2,
117 architecture="ModernBertForSequenceClassification",
118 seq_pooling_type="CLS",
119 attn_type="encoder_only",
120 is_prefix_caching_supported=False,
121 is_chunked_prefill_supported=False,
122 enable_test=True,
123 ),
124 RerankModelInfo(
125 "Alibaba-NLP/gte-multilingual-reranker-base",
126 mteb_score=0.33062,
127 mteb_tol=1e-2,
128 architecture="GteNewForSequenceClassification",
129 hf_overrides={"architectures": ["GteNewForSequenceClassification"]},
130 seq_pooling_type="CLS",
131 attn_type="encoder_only",
132 is_prefix_caching_supported=False,
133 is_chunked_prefill_supported=False,
134 enable_test=True,
135 ),
136 ]
137
138
139 @pytest.mark.flaky(reruns=2)
140 @pytest.mark.parametrize("model_info", MODELS)
141 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
142 mteb_test_embed_models(hf_runner, vllm_runner, model_info)
143
144
145 @pytest.mark.parametrize("model_info", MODELS)
146 def test_embed_models_correctness(
147 hf_runner, vllm_runner, model_info: EmbedModelInfo, example_prompts
148 ) -> None:
149 correctness_test_embed_models(hf_runner, vllm_runner, model_info, example_prompts)
150
151
152 @pytest.mark.flaky(reruns=2)
153 @pytest.mark.parametrize("model_info", RERANK_MODELS)
154 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
155 vllm_extra_kwargs = {}
156 if current_platform.is_rocm():
157 vllm_extra_kwargs["attention_backend"] = "TRITON_ATTN"
158 mteb_test_rerank_models(
159 vllm_runner, model_info, vllm_extra_kwargs=vllm_extra_kwargs
160 )