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 15, 16:00 UTC
[New Model]: jinaai/jina-reranker-v2-base-multilingual (#15876)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: DarkLight1337 <tlleungac@connect.ust.hk>
wang.yuqi committed
Apr 1, 2025 at 23:32 UTC
085cbc4f9fd71d1d25495db130098483fe1d4e70
3 files changed
+86
-3
docs/source/models/supported_models.md
+1
-1
@@ -566,7 +566,7 @@ you should explicitly specify the task type to ensure that the model is used in
566
*
567
- * `XLMRobertaModel`
568
* XLM-RoBERTa-based
569
- * `intfloat/multilingual-e5-large`, etc.
569
+ * `intfloat/multilingual-e5-large`, `jinaai/jina-reranker-v2-base-multilingual`, etc.
570
*
571
*
572
:::
tests/models/embedding/language/test_jina_reranker_v2.py
new
+70
@@ -0,0 +1,70 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# ruff: noqa: E501
3
+"""Compare the scoring outputs of HF and vLLM models.
4
+
5
+Run `pytest tests/models/embedding/language/test_jina_reranker_v2.py`.
6
+"""
7
+import math
8
+
9
+import pytest
10
+
11
+MODELS = [
12
+ "jinaai/jina-reranker-v2-base-multilingual", # Roberta
13
+]
14
+
15
+TEXTS_1 = ["Organic skincare products for sensitive skin"]
16
+
17
+TEXTS_2 = [
18
+ "Organic skincare for sensitive skin with aloe vera and chamomile.",
19
+ "New makeup trends focus on bold colors and innovative techniques",
20
+ "Bio-Hautpflege für empfindliche Haut mit Aloe Vera und Kamille",
21
+ "Neue Make-up-Trends setzen auf kräftige Farben und innovative Techniken",
22
+ "Cuidado de la piel orgánico para piel sensible con aloe vera y manzanilla",
23
+ "Las nuevas tendencias de maquillaje se centran en colores vivos y técnicas innovadoras",
24
+ "针对敏感肌专门设计的天然有机护肤产品",
25
+ "新的化妆趋势注重鲜艳的颜色和创新的技巧",
26
+ "敏感肌のために特別に設計された天然有機スキンケア製品",
27
+ "新しいメイクのトレンドは鮮やかな色と革新的な技術に焦点を当てています",
28
+]
29
+
30
+
31
+@pytest.fixture(scope="module", params=MODELS)
32
+def model_name(request):
33
+ yield request.param
34
+
35
+
36
+@pytest.mark.parametrize("dtype", ["half"])
37
+def test_llm_1_to_1(vllm_runner, hf_runner, model_name, dtype: str):
38
+
39
+ text_pair = [TEXTS_1[0], TEXTS_2[0]]
40
+
41
+ with hf_runner(model_name, dtype=dtype, is_cross_encoder=True) as hf_model:
42
+ hf_outputs = hf_model.predict([text_pair]).tolist()
43
+
44
+ with vllm_runner(model_name, task="score", dtype=dtype,
45
+ max_model_len=None) as vllm_model:
46
+ vllm_outputs = vllm_model.score(text_pair[0], text_pair[1])
47
+
48
+ assert len(vllm_outputs) == 1
49
+ assert len(hf_outputs) == 1
50
+
51
+ assert math.isclose(hf_outputs[0], vllm_outputs[0], rel_tol=0.01)
52
+
53
+
54
+@pytest.mark.parametrize("dtype", ["half"])
55
+def test_llm_1_to_N(vllm_runner, hf_runner, model_name, dtype: str):
56
+
57
+ text_pairs = [[TEXTS_1[0], text] for text in TEXTS_2]
58
+
59
+ with hf_runner(model_name, dtype=dtype, is_cross_encoder=True) as hf_model:
60
+ hf_outputs = hf_model.predict(text_pairs).tolist()
61
+
62
+ with vllm_runner(model_name, task="score", dtype=dtype,
63
+ max_model_len=None) as vllm_model:
64
+ vllm_outputs = vllm_model.score(TEXTS_1[0], TEXTS_2)
65
+
66
+ assert len(vllm_outputs) == 10
67
+ assert len(hf_outputs) == 10
68
+
69
+ assert math.isclose(hf_outputs[0], vllm_outputs[0], rel_tol=0.01)
70
+ assert math.isclose(hf_outputs[1], vllm_outputs[1], rel_tol=0.01)
vllm/model_executor/models/roberta.py
+15
-2
@@ -13,7 +13,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
13
VocabParallelEmbedding)
14
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
15
from vllm.model_executor.models.bert import BertEmbeddingModel, BertModel
16
-from vllm.model_executor.models.utils import maybe_prefix
16
+from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix
17
from vllm.model_executor.pooling_metadata import PoolingMetadata
18
from vllm.sequence import IntermediateTensors, PoolerOutput
19
from vllm.transformers_utils.config import (
@@ -203,6 +203,18 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
203
_pooler: An instance of Pooler used for pooling operations.
204
"""
205
206
+ jina_to_vllm_mapper = WeightsMapper(
207
+ orig_to_new_substr={
208
+ 'emb_ln': "embeddings.LayerNorm",
209
+ 'layers': "layer",
210
+ 'mixer.Wqkv': "attention.self.qkv_proj",
211
+ 'mixer.out_proj': "attention.output.dense",
212
+ 'norm1': "attention.output.LayerNorm",
213
+ 'mlp.fc1': "intermediate.dense",
214
+ 'mlp.fc2': "output.dense",
215
+ 'norm2': "output.LayerNorm",
216
+ })
217
+
218
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
219
super().__init__()
220
config = vllm_config.model_config.hf_config
@@ -219,8 +231,9 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
231
self._pooler = CrossEncodingPooler(config, self.classifier)
232
233
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
222
-
234
bert_weights, task_weights = roberta_task_weights_filter(weights)
235
+ bert_weights = self.jina_to_vllm_mapper.apply(bert_weights)
236
+
237
self.roberta.load_weights(bert_weights)
238
239
params_dict = dict(self.named_parameters())