[Model] Support jina-embeddings-v5-text-nano (EuroBERT encoder backbone) (#50688)
Signed-off-by: Omkar Shewale <omkarshewale2001@gmail.com>
Omkar Shewale committed
Aug 2, 2026 at 20:37 UTC
9a4fd57cac19b334b420ae24226c1d2151566799
6 files changed
+155
-91
docs/models/pooling_models/embed.md
+4
-2
@@ -45,7 +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) | ✅︎ | ✅︎ |
48
+| `JinaEmbeddingsV5Model`<sup>C</sup> | Qwen3-decoder or EuroBERT-encoder backbone with task-specific LoRA adapters | `jinaai/jina-embeddings-v5-text-small`, `jinaai/jina-embeddings-v5-text-nano` (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. | | |
@@ -75,7 +75,9 @@ You can compute pairwise similarity scores to build a similarity matrix using th
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
78
+ `jinaai/jina-embeddings-v5-text-small` (Qwen3 decoder) and
79
+ `jinaai/jina-embeddings-v5-text-nano` (bidirectional EuroBERT encoder,
80
+ `is_decoder=false`) ship with four task-specific LoRA adapters
81
(`retrieval`, `text-matching`, `classification`, `clustering`). vLLM merges the
82
selected adapter into the base weights at load time. Choose the task with
83
`--hf-overrides '{"jina_task": "<task>"}'`; the default is `retrieval`.
tests/models/language/pooling/test_jina_embeddings_v5.py
+15
-12
@@ -1,12 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Backbone validation for Jina Embeddings V5 (issue #50337).
3
+"""Backbone validation for Jina Embeddings V5.
4
5
The V5 family ships two backbones under one `architectures` entry: `-small` is
6
-a Qwen3 decoder (which vLLM implements), while `-nano` is a bidirectional
7
-EuroBERT encoder (which it does not). Upstream ships a separate
8
-`configuration_*.py` per repository, so the only signal distinguishing them is
9
-`is_decoder`, which the encoder variant sets to False.
6
+a Qwen3 decoder, while `-nano` is a bidirectional EuroBERT encoder. Upstream
7
+ships a separate `configuration_*.py` per repository, so the only signal
8
+distinguishing them is `is_decoder`, which the encoder variant sets to False.
9
+`JinaEmbeddingsV5ModelConfig` uses it to enable bidirectional attention for the
10
+encoder variant; `JinaEmbeddingsV5Model` then dispatches to the correct backbone.
11
"""
12
13
from types import SimpleNamespace
@@ -34,16 +35,18 @@ def test_registered_for_the_architecture():
35
36
37
@pytest.mark.cpu_test
37
-def test_encoder_backbone_is_rejected():
38
- """An encoder checkpoint must fail fast, naming the reason.
38
+def test_encoder_backbone_enables_bidirectional_attention():
39
+ """An encoder checkpoint (is_decoder=False) is supported.
40
40
- Without this, loading proceeds until every `q_norm`/`k_norm` weight is
41
- reported as uninitialized, which does not point at the actual problem.
41
+ The handler sets is_causal=False so the Llama backbone uses
42
+ EncoderOnlyAttention; JinaEmbeddingsV5Model then dispatches to the encoder
43
+ implementation.
44
"""
43
- model_config = _model_config(PretrainedConfig(is_decoder=False))
45
+ hf_config = PretrainedConfig(is_decoder=False)
46
45
- with pytest.raises(NotImplementedError, match="is_decoder=False"):
46
- JinaEmbeddingsV5ModelConfig.verify_and_update_model_config(model_config)
47
+ JinaEmbeddingsV5ModelConfig.verify_and_update_model_config(_model_config(hf_config))
48
+
49
+ assert hf_config.is_causal is False
50
51
52
@pytest.mark.cpu_test
tests/models/language/pooling_mteb_test/mteb_embed_utils.py
+11
-2
@@ -207,9 +207,18 @@ def mteb_test_embed_models(
207
vllm_dtype = vllm_model.llm.llm_engine.model_config.dtype
208
head_dtype = model_config.head_dtype
209
210
- # Test embedding_size, isnan and whether to use normalize
210
+ # Test embedding_size, isnan and whether to use normalize.
211
+ # Apply the same prompt_prefix used for scoring so this check compares
212
+ # identical effective inputs: HF's SentenceTransformer.encode() applies
213
+ # the model's default prompt (e.g. "Document: "), so vLLM must receive it
214
+ # too. This mirrors VllmMtebEncoder and changes no production behavior.
215
+ consistency_prompts = (
216
+ [prompt_prefix + p for p in example_prompts]
217
+ if prompt_prefix
218
+ else example_prompts
219
+ )
220
vllm_outputs = vllm_model.embed(
212
- example_prompts,
221
+ consistency_prompts,
222
tokenization_kwargs=dict(truncate_prompt_tokens=-1),
223
)
224
outputs_tensor = torch.tensor(vllm_outputs)
tests/models/language/pooling_mteb_test/test_jina.py
+8
@@ -38,6 +38,14 @@ EMBEDDING_MODELS = [
38
is_prefix_caching_supported=True,
39
is_chunked_prefill_supported=True,
40
),
41
+ EmbedModelInfo(
42
+ "jinaai/jina-embeddings-v5-text-nano",
43
+ architecture="JinaEmbeddingsV5Model",
44
+ seq_pooling_type="LAST",
45
+ attn_type="encoder_only",
46
+ is_prefix_caching_supported=False,
47
+ is_chunked_prefill_supported=False,
48
+ ),
49
]
50
51
RERANK_MODELS = [
vllm/model_executor/models/config.py
+5
-14
@@ -479,30 +479,21 @@ class JinaEmbeddingsV5ModelConfig(VerifyAndUpdateConfig):
479
480
@staticmethod
481
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
482
- """Reject checkpoints whose backbone is not the Qwen3 decoder.
482
+ """Enable the bidirectional encoder backbone for -nano checkpoints.
483
484
The V5 family ships more than one backbone under a single
485
`architectures` entry: the `-small` variants are Qwen3 decoders, while
486
`-nano` is a bidirectional EuroBERT encoder. Upstream ships a separate
487
`configuration_*.py` per repository, so the config carries no backbone
488
field and the encoder variants are only identifiable by
489
- `is_decoder=False`. `JinaEmbeddingsV5Model` extends `Qwen3ForCausalLM`
490
- and has no path to an encoder backbone, so without this check such a
491
- checkpoint fails later with every `q_norm` and `k_norm` weight reported
492
- as uninitialized, which does not point at the actual problem.
493
-
494
- Raises:
495
- NotImplementedError: If the checkpoint uses an encoder backbone.
489
+ `is_decoder=False`. For encoder checkpoints, set `is_causal=False` so the
490
+ Llama backbone uses EncoderOnlyAttention; `JinaEmbeddingsV5Model` then
491
+ dispatches to its encoder implementation.
492
"""
493
if getattr(model_config.hf_config, "is_decoder", True):
494
return
495
500
- raise NotImplementedError(
501
- "This jina-embeddings-v5 checkpoint uses a bidirectional encoder "
502
- "backbone (is_decoder=False). JinaEmbeddingsV5Model is built on "
503
- "Qwen3 and cannot serve it, so only the Qwen3-based decoder "
504
- "variants are supported, such as jina-embeddings-v5-text-small."
505
- )
496
+ model_config.hf_config.is_causal = False
497
498
499
class JinaForRankingConfig(VerifyAndUpdateConfig):
vllm/model_executor/models/jina.py
+112
-61
@@ -24,6 +24,7 @@ from ..layers.pooler.tokwise import (
24
)
25
from .interfaces import SupportsLateInteraction
26
from .interfaces_base import VllmModelForPooling
27
+from .llama import LlamaForCausalLM
28
from .qwen3 import Qwen3ForCausalLM, Qwen3Model
29
from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
30
@@ -185,11 +186,79 @@ def _build_lora_pairs(adapter_weights: dict) -> dict:
186
return dict(lora_pairs)
187
188
188
-class JinaEmbeddingsV5Model(Qwen3ForCausalLM, VllmModelForPooling):
189
- """Jina Embeddings V5 with task-specific LoRA adapters merged at load time.
189
+def _setup_jina_v5_task_and_pooler(model: nn.Module, vllm_config: VllmConfig) -> None:
190
+ """Shared init for jina-embeddings-v5 wrappers: select task + build pooler."""
191
+ model._model_name = vllm_config.model_config.model
192
+ model._revision = vllm_config.model_config.revision
193
191
- Extends Qwen3ForCausalLM (the underlying architecture) and declares itself
192
- as a pooling model so that as_embedding_model() does not wrap it.
194
+ model._task = getattr(
195
+ vllm_config.model_config.hf_config, "jina_task", _DEFAULT_TASK
196
+ )
197
+ if model._task not in _SUPPORTED_TASKS:
198
+ logger.warning(
199
+ "Unknown jina_task=%r. Falling back to %r.",
200
+ model._task,
201
+ _DEFAULT_TASK,
202
+ )
203
+ model._task = _DEFAULT_TASK
204
+
205
+ pooler_config = vllm_config.model_config.pooler_config
206
+ assert pooler_config is not None
207
+ model.pooler = DispatchPooler.for_embedding(pooler_config)
208
+
209
+
210
+def _load_jina_v5_weights(
211
+ model: nn.Module, weights: Iterable[tuple[str, torch.Tensor]]
212
+) -> set[str]:
213
+ """Shared loader: merge the selected task LoRA adapter into the base weights."""
214
+ lora_pairs: dict = {}
215
+ scaling = 1.0
216
+
217
+ result = _load_adapter(model._model_name, model._task, model._revision)
218
+ if result is None:
219
+ logger.warning(
220
+ "No adapter found for task %r in %r. Loading raw base weights.",
221
+ model._task,
222
+ model._model_name,
223
+ )
224
+ else:
225
+ adapter_config, adapter_weights = result
226
+ scaling = adapter_config["lora_alpha"] / adapter_config["r"]
227
+ lora_pairs = _build_lora_pairs(adapter_weights)
228
+ logger.info(
229
+ "Loaded %d adapter tensors for task %r (scaling=%.4f, %d LoRA pairs)",
230
+ len(adapter_weights),
231
+ model._task,
232
+ scaling,
233
+ len(lora_pairs),
234
+ )
235
+
236
+ def _merge_weights(
237
+ weights: Iterable[tuple[str, torch.Tensor]],
238
+ ) -> Iterable[tuple[str, torch.Tensor]]:
239
+ for name, tensor in weights:
240
+ clean_name = name
241
+ if clean_name.startswith("model."):
242
+ clean_name = clean_name[len("model.") :]
243
+
244
+ if clean_name in lora_pairs:
245
+ pair = lora_pairs[clean_name]
246
+ if "A" in pair and "B" in pair:
247
+ lora_A = pair["A"].to(device=tensor.device, dtype=tensor.dtype)
248
+ lora_B = pair["B"].to(device=tensor.device, dtype=tensor.dtype)
249
+ tensor = tensor + (lora_B @ lora_A) * scaling
250
+ yield name, tensor
251
+
252
+ loader = AutoWeightsLoader(model, ignore_unexpected_prefixes=["lm_head."])
253
+ weights = _merge_weights(weights)
254
+ return loader.load_weights(weights, mapper=model.hf_to_vllm_mapper)
255
+
256
+
257
+class JinaEmbeddingsV5DecoderModel(Qwen3ForCausalLM, VllmModelForPooling):
258
+ """jina-embeddings-v5 with a Qwen3 decoder backbone (e.g. -small).
259
+
260
+ Task-specific LoRA adapters are merged into the base weights at load time.
261
+ Declares itself a pooling model so that as_embedding_model() does not wrap it.
262
"""
263
264
is_pooling_model = True
@@ -199,64 +268,46 @@ class JinaEmbeddingsV5Model(Qwen3ForCausalLM, VllmModelForPooling):
268
269
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
270
super().__init__(vllm_config=vllm_config, prefix=prefix)
271
+ _setup_jina_v5_task_and_pooler(self, vllm_config)
272
203
- self._model_name = vllm_config.model_config.model
204
- self._revision = vllm_config.model_config.revision
273
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
274
+ return _load_jina_v5_weights(self, weights)
275
206
- self._task = getattr(
207
- vllm_config.model_config.hf_config, "jina_task", _DEFAULT_TASK
208
- )
209
- if self._task not in _SUPPORTED_TASKS:
210
- logger.warning(
211
- "Unknown jina_task=%r. Falling back to %r.",
212
- self._task,
213
- _DEFAULT_TASK,
214
- )
215
- self._task = _DEFAULT_TASK
216
-
217
- pooler_config = vllm_config.model_config.pooler_config
218
- assert pooler_config is not None
219
- self.pooler = DispatchPooler.for_embedding(pooler_config)
276
+
277
+class JinaEmbeddingsV5EncoderModel(LlamaForCausalLM, VllmModelForPooling):
278
+ """jina-embeddings-v5 with a bidirectional EuroBERT (Llama) encoder backbone.
279
+
280
+ Used by encoder checkpoints such as jina-embeddings-v5-text-nano
281
+ (``is_decoder=False``). EuroBERT is architecturally a bidirectional Llama, so
282
+ the LlamaModel backbone switches to EncoderOnlyAttention when the config
283
+ carries ``is_causal=False`` (set by ``JinaEmbeddingsV5ModelConfig``).
284
+ """
285
+
286
+ is_pooling_model = True
287
+ hf_to_vllm_mapper = LlamaForCausalLM.hf_to_vllm_mapper | WeightsMapper(
288
+ orig_to_new_prefix={"": "model."}
289
+ )
290
+
291
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
292
+ super().__init__(vllm_config=vllm_config, prefix=prefix)
293
+ _setup_jina_v5_task_and_pooler(self, vllm_config)
294
295
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
222
- lora_pairs: dict = {}
223
- scaling = 1.0
224
-
225
- result = _load_adapter(self._model_name, self._task, self._revision)
226
- if result is None:
227
- logger.warning(
228
- "No adapter found for task %r in %r. Loading raw base weights.",
229
- self._task,
230
- self._model_name,
231
- )
232
- else:
233
- adapter_config, adapter_weights = result
234
- scaling = adapter_config["lora_alpha"] / adapter_config["r"]
235
- lora_pairs = _build_lora_pairs(adapter_weights)
236
- logger.info(
237
- "Loaded %d adapter tensors for task %r (scaling=%.4f, %d LoRA pairs)",
238
- len(adapter_weights),
239
- self._task,
240
- scaling,
241
- len(lora_pairs),
242
- )
243
-
244
- def _merge_weights(
245
- weights: Iterable[tuple[str, torch.Tensor]],
246
- ) -> Iterable[tuple[str, torch.Tensor]]:
247
- for name, tensor in weights:
248
- clean_name = name
249
- if clean_name.startswith("model."):
250
- clean_name = clean_name[len("model.") :]
251
-
252
- if clean_name in lora_pairs:
253
- pair = lora_pairs[clean_name]
254
- if "A" in pair and "B" in pair:
255
- lora_A = pair["A"].to(device=tensor.device, dtype=tensor.dtype)
256
- lora_B = pair["B"].to(device=tensor.device, dtype=tensor.dtype)
257
- tensor = tensor + (lora_B @ lora_A) * scaling
258
- yield name, tensor
259
-
260
- loader = AutoWeightsLoader(self, ignore_unexpected_prefixes=["lm_head."])
261
- weights = _merge_weights(weights)
262
- return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
296
+ return _load_jina_v5_weights(self, weights)
297
+
298
+
299
+class JinaEmbeddingsV5Model(JinaEmbeddingsV5DecoderModel):
300
+ """Dispatcher for the jina-embeddings-v5 family.
301
+
302
+ The family ships two backbones under one ``architectures`` entry: Qwen3
303
+ decoders (-small) and bidirectional EuroBERT encoders (-nano), told apart by
304
+ ``is_decoder``. Inherits the decoder implementation so registry introspection
305
+ still sees a valid pooling model, and ``__new__`` swaps in the encoder
306
+ variant for encoder checkpoints.
307
+ """
308
+
309
+ def __new__(cls, *, vllm_config: VllmConfig, prefix: str = ""):
310
+ is_decoder = getattr(vllm_config.model_config.hf_config, "is_decoder", True)
311
+ if not is_decoder:
312
+ return JinaEmbeddingsV5EncoderModel(vllm_config=vllm_config, prefix=prefix)
313
+ return super().__new__(cls)