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 5, 02:15 UTC
Add Jina Embeddings v5 model support (fixes #38633) (#39575)
Signed-off-by: Abhijit <abroy@redhat.com> Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>
Abhijit Roy committed
Apr 16, 2026 at 12:07 UTC
2cdf86044d7e3bdcfbbb39a21a37b57ee9d4e65b
7 files changed
+218
-10
docs/models/pooling_models/embed.md
+7
@@ -45,6 +45,7 @@ You can compute pairwise similarity scores to build a similarity matrix using th
45
| `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ |
46
| `GteModel` | Arctic-Embed-2.0-M | `Snowflake/snowflake-arctic-embed-m-v2.0`. | | |
47
| `GteNewModel` | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-base`, etc. | | |
48
+| `JinaEmbeddingsV5Model`<sup>C</sup> | Qwen3-based with task-specific LoRA adapters | `jinaai/jina-embeddings-v5-text-small` (see note) | ✅︎ | ✅︎ |
49
| `LlamaBidirectionalModel`<sup>C</sup> | Llama-based with bidirectional attention | `nvidia/llama-nemotron-embed-1b-v2`, etc. | ✅︎ | ✅︎ |
50
| `LlamaModel`<sup>C</sup>, `LlamaForCausalLM`<sup>C</sup>, `MistralModel`<sup>C</sup>, etc. | Llama-based | `intfloat/e5-mistral-7b-instruct`, etc. | ✅︎ | ✅︎ |
51
| `ModernBertModel` | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | | |
@@ -73,6 +74,12 @@ You can compute pairwise similarity scores to build a similarity matrix using th
74
!!! note
75
`jinaai/jina-embeddings-v3` supports multiple tasks through LoRA, while vllm temporarily only supports text-matching tasks by merging LoRA weights.
76
77
+!!! note
78
+ `jinaai/jina-embeddings-v5-text-small` ships with four task-specific LoRA adapters
79
+ (`retrieval`, `text-matching`, `classification`, `clustering`). vLLM merges the
80
+ selected adapter into the base weights at load time. Choose the task with
81
+ `--hf-overrides '{"jina_task": "<task>"}'`; the default is `retrieval`.
82
+
83
### Multimodal Models
84
85
!!! note
tests/conftest.py
+6
@@ -364,6 +364,7 @@ class HfRunner:
364
model_name: str,
365
dtype: str = "auto",
366
*,
367
+ revision: str | None = None,
368
model_kwargs: dict[str, Any] | None = None,
369
trust_remote_code: bool = True,
370
is_sentence_transformer: bool = False,
@@ -383,6 +384,7 @@ class HfRunner:
384
self._init(
385
model_name=model_name,
386
dtype=dtype,
387
+ revision=revision,
388
model_kwargs=model_kwargs,
389
trust_remote_code=trust_remote_code,
390
is_sentence_transformer=is_sentence_transformer,
@@ -396,6 +398,7 @@ class HfRunner:
398
model_name: str,
399
dtype: str = "auto",
400
*,
401
+ revision: str | None = None,
402
model_kwargs: dict[str, Any] | None = None,
403
trust_remote_code: bool = True,
404
is_sentence_transformer: bool = False,
@@ -437,6 +440,7 @@ class HfRunner:
440
441
self.model = SentenceTransformer(
442
model_name,
443
+ revision=revision,
444
device=self.device,
445
model_kwargs=model_kwargs,
446
trust_remote_code=trust_remote_code,
@@ -447,6 +451,7 @@ class HfRunner:
451
452
self.model = CrossEncoder(
453
model_name,
454
+ revision=revision,
455
device=self.device,
456
automodel_args=model_kwargs,
457
trust_remote_code=trust_remote_code,
@@ -456,6 +461,7 @@ class HfRunner:
461
nn.Module,
462
auto_cls.from_pretrained(
463
model_name,
464
+ revision=revision,
465
trust_remote_code=trust_remote_code,
466
**model_kwargs,
467
),
tests/models/language/pooling_mteb_test/mteb_embed_utils.py
+26
-4
@@ -74,10 +74,25 @@ class MtebEmbedMixin(mteb.EncoderProtocol):
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):
78
- def __init__(self, vllm_model):
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,
@@ -87,7 +102,11 @@ class VllmMtebEncoder(MtebEmbedMixin):
102
) -> np.ndarray:
103
# Hoping to discover potential scheduling
104
# issues by randomizing the order.
90
- sentences = [text for batch in inputs for text in batch["text"]]
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)
@@ -143,6 +162,7 @@ def mteb_test_embed_models(
162
vllm_extra_kwargs=None,
163
hf_model_callback=None,
164
atol=MTEB_EMBED_TOL,
165
+ prompt_prefix: str | None = None,
166
):
167
vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)
168
@@ -182,7 +202,7 @@ def mteb_test_embed_models(
202
)
203
204
vllm_main_score = run_mteb_embed_task(
185
- VllmMtebEncoder(vllm_model), MTEB_EMBED_TASKS
205
+ VllmMtebEncoder(vllm_model, prompt_prefix=prompt_prefix), MTEB_EMBED_TASKS
206
)
207
vllm_dtype = vllm_model.llm.llm_engine.model_config.dtype
208
head_dtype = model_config.head_dtype
@@ -210,7 +230,9 @@ def mteb_test_embed_models(
230
if hf_model_callback is not None:
231
hf_model_callback(hf_model)
232
213
- st_main_score = run_mteb_embed_task(hf_model, MTEB_EMBED_TASKS)
233
+ st_main_score = run_mteb_embed_task(
234
+ HfMtebEncoder(hf_model), MTEB_EMBED_TASKS
235
+ )
236
st_dtype = next(hf_model.model.parameters()).dtype
237
238
# Check embeddings close to hf outputs
tests/models/language/pooling_mteb_test/test_jina.py
+25
-5
@@ -28,7 +28,16 @@ EMBEDDING_MODELS = [
28
attn_type="encoder_only",
29
is_prefix_caching_supported=False,
30
is_chunked_prefill_supported=False,
31
- )
31
+ ),
32
+ EmbedModelInfo(
33
+ "jinaai/jina-embeddings-v5-text-small",
34
+ mteb_score=0.794535707854956,
35
+ architecture="JinaEmbeddingsV5Model",
36
+ seq_pooling_type="LAST",
37
+ attn_type="decoder",
38
+ is_prefix_caching_supported=True,
39
+ is_chunked_prefill_supported=True,
40
+ ),
41
]
42
43
RERANK_MODELS = [
@@ -46,11 +55,18 @@ RERANK_MODELS = [
55
56
@pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
57
def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
58
+ task = "retrieval" if "v5" in model_info.name else "text-matching"
59
+ prompt_prefix: str | None = "Document: " if "v5" in model_info.name else None
60
+
61
def hf_model_callback(model):
50
- model.encode = partial(model.encode, task="text-matching")
62
+ model.encode = partial(model.encode, task=task)
63
64
mteb_test_embed_models(
53
- hf_runner, vllm_runner, model_info, hf_model_callback=hf_model_callback
65
+ hf_runner,
66
+ vllm_runner,
67
+ model_info,
68
+ hf_model_callback=hf_model_callback,
69
+ prompt_prefix=prompt_prefix,
70
)
71
72
@@ -58,8 +74,10 @@ def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -
74
def test_embed_models_correctness(
75
hf_runner, vllm_runner, model_info: EmbedModelInfo, example_prompts
76
) -> None:
77
+ task = "retrieval" if "v5" in model_info.name else "text-matching"
78
+
79
def hf_model_callback(model):
62
- model.encode = partial(model.encode, task="text-matching")
80
+ model.encode = partial(model.encode, task=task)
81
82
correctness_test_embed_models(
83
hf_runner,
@@ -97,12 +115,14 @@ def test_matryoshka(
115
# ST will strip the input texts, see test_embedding.py
116
example_prompts = [str(s).strip() for s in example_prompts]
117
118
+ task = "retrieval" if "v5" in model_info.name else "text-matching"
119
+
120
with hf_runner(
121
model_info.name,
122
dtype=dtype,
123
is_sentence_transformer=True,
124
) as hf_model:
105
- hf_outputs = hf_model.encode(example_prompts, task="text-matching")
125
+ hf_outputs = hf_model.encode(example_prompts, task=task)
126
hf_outputs = matryoshka_fy(hf_outputs, dimensions)
127
128
with vllm_runner(
tests/models/registry.py
+4
@@ -609,6 +609,10 @@ _EMBEDDING_EXAMPLE_MODELS = {
609
trust_remote_code=True,
610
hf_overrides={"architectures": ["GteNewModel"]},
611
),
612
+ "JinaEmbeddingsV5Model": _HfExamplesInfo(
613
+ "jinaai/jina-embeddings-v5-text-small",
614
+ trust_remote_code=True,
615
+ ),
616
"LlamaModel": _HfExamplesInfo("llama", is_available_online=False),
617
"LlamaBidirectionalModel": _HfExamplesInfo(
618
"nvidia/llama-nemotron-embed-1b-v2", trust_remote_code=True
vllm/model_executor/models/jina.py
+149
-1
@@ -1,14 +1,19 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
# Adapted from https://huggingface.co/jinaai/jina-reranker-v3/blob/main/modeling.py
4
+import json
5
+import logging
6
+from collections import defaultdict
7
from collections.abc import Iterable
8
9
import torch
10
+from safetensors.torch import load as safetensors_load
11
from torch import nn
12
13
from vllm.config import VllmConfig
14
from vllm.sequence import IntermediateTensors
15
from vllm.tasks import PoolingTask
16
+from vllm.transformers_utils.repo_utils import get_hf_file_bytes
17
from vllm.v1.pool.metadata import PoolingMetadata
18
19
from ..layers.pooler import DispatchPooler
@@ -18,9 +23,12 @@ from ..layers.pooler.tokwise import (
23
TokenPoolingMethodOutputItem,
24
)
25
from .interfaces import SupportsLateInteraction
21
-from .qwen3 import Qwen3Model
26
+from .interfaces_base import VllmModelForPooling
27
+from .qwen3 import Qwen3ForCausalLM, Qwen3Model
28
from .utils import AutoWeightsLoader, maybe_prefix
29
30
+logger = logging.getLogger(__name__)
31
+
32
33
class JinaForRanking(nn.Module, SupportsLateInteraction):
34
is_pooling_model = True
@@ -108,3 +116,143 @@ class JinaForRankingPool(StepPool):
116
embeds_list.append(embeds)
117
118
return embeds_list
119
+
120
+
121
+# jina-embeddings-v5-text-small wraps Qwen3-0.6B-Base with four task-specific
122
+# LoRA adapters. This implementation merges the selected adapter into the base
123
+# weights at load time to avoid any runtime dependency on peft.
124
+#
125
+# Task selection:
126
+# Pass --hf-overrides '{"jina_task": "retrieval"}' to select one of:
127
+# retrieval (default), text-matching, classification, clustering.
128
+
129
+_DEFAULT_TASK = "retrieval"
130
+_SUPPORTED_TASKS = {"retrieval", "text-matching", "classification", "clustering"}
131
+
132
+
133
+def _load_adapter(
134
+ model: str,
135
+ task: str,
136
+ revision: str | None,
137
+) -> tuple[dict, dict[str, torch.Tensor]] | None:
138
+ """Load adapter config and weights from a local path or HF repo.
139
+
140
+ Returns (adapter_config, adapter_weights) or None if not found.
141
+ """
142
+ config_bytes = get_hf_file_bytes(
143
+ f"adapters/{task}/adapter_config.json",
144
+ model,
145
+ revision,
146
+ )
147
+ if config_bytes is None:
148
+ return None
149
+
150
+ adapter_config = json.loads(config_bytes)
151
+
152
+ weights_bytes = get_hf_file_bytes(
153
+ f"adapters/{task}/adapter_model.safetensors",
154
+ model,
155
+ revision,
156
+ )
157
+ if weights_bytes is None:
158
+ return None
159
+
160
+ adapter_weights = safetensors_load(weights_bytes)
161
+ return adapter_config, adapter_weights
162
+
163
+
164
+def _build_lora_pairs(adapter_weights: dict) -> dict:
165
+ """Group raw adapter tensors into {base_key: {"A": tensor, "B": tensor}} pairs.
166
+
167
+ Transforms adapter keys like:
168
+ base_model.model.layers.0.self_attn.q_proj.lora_A.weight
169
+ Into base keys like:
170
+ layers.0.self_attn.q_proj.weight
171
+ """
172
+ lora_pairs = defaultdict(dict)
173
+ for key, tensor in adapter_weights.items():
174
+ clean_key = key
175
+ if clean_key.startswith("base_model.model."):
176
+ clean_key = clean_key[len("base_model.model.") :]
177
+
178
+ if ".lora_A." in clean_key:
179
+ base_key = clean_key.split(".lora_A.")[0] + ".weight"
180
+ lora_pairs[base_key]["A"] = tensor
181
+ elif ".lora_B." in clean_key:
182
+ base_key = clean_key.split(".lora_B.")[0] + ".weight"
183
+ lora_pairs[base_key]["B"] = tensor
184
+
185
+ return dict(lora_pairs)
186
+
187
+
188
+class JinaEmbeddingsV5Model(Qwen3ForCausalLM, VllmModelForPooling):
189
+ """Jina Embeddings V5 with task-specific LoRA adapters merged at load time.
190
+
191
+ Extends Qwen3ForCausalLM (the underlying architecture) and declares itself
192
+ as a pooling model so that as_embedding_model() does not wrap it.
193
+ """
194
+
195
+ is_pooling_model = True
196
+
197
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
198
+ super().__init__(vllm_config=vllm_config, prefix=prefix)
199
+
200
+ self._model_name = vllm_config.model_config.model
201
+ self._revision = vllm_config.model_config.revision
202
+
203
+ self._task = getattr(
204
+ vllm_config.model_config.hf_config, "jina_task", _DEFAULT_TASK
205
+ )
206
+ if self._task not in _SUPPORTED_TASKS:
207
+ logger.warning(
208
+ "Unknown jina_task=%r. Falling back to %r.",
209
+ self._task,
210
+ _DEFAULT_TASK,
211
+ )
212
+ self._task = _DEFAULT_TASK
213
+
214
+ pooler_config = vllm_config.model_config.pooler_config
215
+ assert pooler_config is not None
216
+ self.pooler = DispatchPooler.for_embedding(pooler_config)
217
+
218
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
219
+ lora_pairs: dict = {}
220
+ scaling = 1.0
221
+
222
+ result = _load_adapter(self._model_name, self._task, self._revision)
223
+ if result is None:
224
+ logger.warning(
225
+ "No adapter found for task %r in %r. Loading raw base weights.",
226
+ self._task,
227
+ self._model_name,
228
+ )
229
+ else:
230
+ adapter_config, adapter_weights = result
231
+ scaling = adapter_config["lora_alpha"] / adapter_config["r"]
232
+ lora_pairs = _build_lora_pairs(adapter_weights)
233
+ logger.info(
234
+ "Loaded %d adapter tensors for task %r (scaling=%.4f, %d LoRA pairs)",
235
+ len(adapter_weights),
236
+ self._task,
237
+ scaling,
238
+ len(lora_pairs),
239
+ )
240
+
241
+ def _merge_weights(
242
+ weights: Iterable[tuple[str, torch.Tensor]],
243
+ ) -> Iterable[tuple[str, torch.Tensor]]:
244
+ for name, tensor in weights:
245
+ clean_name = name
246
+ if clean_name.startswith("model."):
247
+ clean_name = clean_name[len("model.") :]
248
+
249
+ if clean_name in lora_pairs:
250
+ pair = lora_pairs[clean_name]
251
+ if "A" in pair and "B" in pair:
252
+ lora_A = pair["A"].to(device=tensor.device, dtype=tensor.dtype)
253
+ lora_B = pair["B"].to(device=tensor.device, dtype=tensor.dtype)
254
+ tensor = tensor + (lora_B @ lora_A) * scaling
255
+ yield name, tensor
256
+
257
+ loaded = self.model.load_weights(_merge_weights(weights))
258
+ return {f"model.{name}" for name in loaded}
vllm/model_executor/models/registry.py
+1
@@ -227,6 +227,7 @@ _EMBEDDING_MODELS = {
227
"GritLM": ("gritlm", "GritLM"),
228
"GteModel": ("bert_with_rope", "SnowflakeGteNewModel"),
229
"GteNewModel": ("bert_with_rope", "GteNewModel"),
230
+ "JinaEmbeddingsV5Model": ("jina", "JinaEmbeddingsV5Model"),
231
"LlamaBidirectionalModel": ("llama", "LlamaBidirectionalModel"),
232
"LlamaModel": ("llama", "LlamaForCausalLM"),
233
**{