| 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.model_executor.layers.logits_processor import LogitsProcessor |
| 15 | from vllm.model_executor.layers.pooler import PoolingParamsUpdate |
| 16 | from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead |
| 17 | from vllm.sequence import IntermediateTensors |
| 18 | from vllm.tasks import PoolingTask |
| 19 | from vllm.transformers_utils.repo_utils import get_hf_file_bytes |
| 20 | from vllm.utils.torch_utils import async_tensor_h2d |
| 21 | from vllm.v1.pool.metadata import PoolingMetadata |
| 22 | |
| 23 | from ..layers.pooler import DispatchPooler |
| 24 | from ..layers.pooler.tokwise import ( |
| 25 | StepPool, |
| 26 | TokenPooler, |
| 27 | TokenPoolingMethodOutputItem, |
| 28 | ) |
| 29 | from .interfaces import SupportsLateInteraction |
| 30 | from .interfaces_base import VllmModelForPooling |
| 31 | from .llama import LlamaForCausalLM |
| 32 | from .qwen3 import Qwen3ForCausalLM, Qwen3Model |
| 33 | from .utils import ( |
| 34 | AutoWeightsLoader, |
| 35 | StageMissingLayer, |
| 36 | WeightsMapper, |
| 37 | maybe_prefix, |
| 38 | no_init_weights, |
| 39 | ) |
| 40 | |
| 41 | logger = logging.getLogger(__name__) |
| 42 | |
| 43 | |
| 44 | class JinaForRanking(nn.Module, SupportsLateInteraction): |
| 45 | is_pooling_model = True |
| 46 | |
| 47 | hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"lm_head.": None}) |
| 48 | |
| 49 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 50 | super().__init__() |
| 51 | config = vllm_config.model_config.hf_config |
| 52 | quant_config = vllm_config.quant_config |
| 53 | |
| 54 | self.config = config |
| 55 | self.projector_dim: int = config.embedding_size |
| 56 | |
| 57 | self.vllm_config = vllm_config |
| 58 | self.quant_config = quant_config |
| 59 | self.model = Qwen3Model( |
| 60 | vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") |
| 61 | ) |
| 62 | |
| 63 | self.projector = nn.Sequential( |
| 64 | nn.Linear(config.hidden_size, config.hidden_size // 2, bias=False), |
| 65 | nn.ReLU(), |
| 66 | nn.Linear(config.hidden_size // 2, self.projector_dim, bias=False), |
| 67 | ) |
| 68 | |
| 69 | self.pooler = DispatchPooler( |
| 70 | { |
| 71 | "token_embed": TokenPooler( |
| 72 | pooling=JinaForRankingPool(self.projector), |
| 73 | ) |
| 74 | } |
| 75 | ) |
| 76 | |
| 77 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 78 | return self.model.embed_input_ids(input_ids) |
| 79 | |
| 80 | def forward( |
| 81 | self, |
| 82 | input_ids: torch.Tensor | None, |
| 83 | positions: torch.Tensor, |
| 84 | intermediate_tensors: IntermediateTensors | None = None, |
| 85 | inputs_embeds: torch.Tensor | None = None, |
| 86 | ) -> torch.Tensor | IntermediateTensors: |
| 87 | hidden_states = self.model( |
| 88 | input_ids, positions, intermediate_tensors, inputs_embeds |
| 89 | ) |
| 90 | return hidden_states |
| 91 | |
| 92 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 93 | loader = AutoWeightsLoader(self) |
| 94 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 95 | |
| 96 | |
| 97 | class JinaForRankingPool(StepPool): |
| 98 | def __init__(self, projector: nn.Sequential): |
| 99 | super().__init__() |
| 100 | |
| 101 | self.doc_token_id = 151670 |
| 102 | self.query_token_id = 151671 |
| 103 | self.projector = projector |
| 104 | |
| 105 | def get_supported_tasks(self) -> set[PoolingTask]: |
| 106 | return {"token_embed"} |
| 107 | |
| 108 | def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: |
| 109 | return PoolingParamsUpdate(requires_token_ids=True) |
| 110 | |
| 111 | def forward( |
| 112 | self, |
| 113 | hidden_states: torch.Tensor, |
| 114 | pooling_metadata: PoolingMetadata, |
| 115 | ) -> list[TokenPoolingMethodOutputItem]: |
| 116 | pooled_data_lst = super().forward(hidden_states, pooling_metadata) |
| 117 | prompt_token_ids_cpu = pooling_metadata.get_prompt_token_ids_cpu() |
| 118 | |
| 119 | embeds_list = list[torch.Tensor | None]() |
| 120 | for data, token_ids_cpu in zip(pooled_data_lst, prompt_token_ids_cpu): |
| 121 | # for unfinished chunked prefill |
| 122 | if data is None: |
| 123 | embeds_list.append(None) |
| 124 | else: |
| 125 | doc_match = torch.eq(token_ids_cpu, self.doc_token_id) |
| 126 | docs_indexes_cpu = torch.where(doc_match)[0] |
| 127 | query_indexes_cpu = torch.where( |
| 128 | torch.eq(token_ids_cpu, self.query_token_id) |
| 129 | )[0] |
| 130 | |
| 131 | # The JinaForRanking model concatenates docs first, then query. |
| 132 | # Let's stay consistent with this novel design. |
| 133 | indexes_cpu = torch.cat([docs_indexes_cpu, query_indexes_cpu]) |
| 134 | indexes = async_tensor_h2d( |
| 135 | indexes_cpu, device=data.device, dtype=torch.long |
| 136 | ) |
| 137 | embeds = self.projector(data[indexes]) |
| 138 | embeds_list.append(embeds) |
| 139 | |
| 140 | return embeds_list |
| 141 | |
| 142 | |
| 143 | # jina-embeddings-v5-text-small wraps Qwen3-0.6B-Base with four task-specific |
| 144 | # LoRA adapters. This implementation merges the selected adapter into the base |
| 145 | # weights at load time to avoid any runtime dependency on peft. |
| 146 | # |
| 147 | # Task selection: |
| 148 | # Pass --hf-overrides '{"jina_task": "retrieval"}' to select one of: |
| 149 | # retrieval (default), text-matching, classification, clustering. |
| 150 | |
| 151 | _DEFAULT_TASK = "retrieval" |
| 152 | _SUPPORTED_TASKS = {"retrieval", "text-matching", "classification", "clustering"} |
| 153 | |
| 154 | |
| 155 | def _load_adapter( |
| 156 | model: str, |
| 157 | task: str, |
| 158 | revision: str | None, |
| 159 | ) -> tuple[dict, dict[str, torch.Tensor]] | None: |
| 160 | """Load adapter config and weights from a local path or HF repo. |
| 161 | |
| 162 | Returns (adapter_config, adapter_weights) or None if not found. |
| 163 | """ |
| 164 | config_bytes = get_hf_file_bytes( |
| 165 | f"adapters/{task}/adapter_config.json", |
| 166 | model, |
| 167 | revision, |
| 168 | ) |
| 169 | if config_bytes is None: |
| 170 | return None |
| 171 | |
| 172 | adapter_config = json.loads(config_bytes) |
| 173 | |
| 174 | weights_bytes = get_hf_file_bytes( |
| 175 | f"adapters/{task}/adapter_model.safetensors", |
| 176 | model, |
| 177 | revision, |
| 178 | ) |
| 179 | if weights_bytes is None: |
| 180 | return None |
| 181 | |
| 182 | adapter_weights = safetensors_load(weights_bytes) |
| 183 | return adapter_config, adapter_weights |
| 184 | |
| 185 | |
| 186 | def _build_lora_pairs( |
| 187 | adapter_weights: dict[str, torch.Tensor], |
| 188 | ) -> dict[str, dict[str, torch.Tensor]]: |
| 189 | """Group raw adapter tensors into {base_key: {"A": tensor, "B": tensor}} pairs. |
| 190 | |
| 191 | Transforms adapter keys like: |
| 192 | base_model.model.layers.0.self_attn.q_proj.lora_A.weight |
| 193 | Into base keys like: |
| 194 | layers.0.self_attn.q_proj.weight |
| 195 | """ |
| 196 | lora_pairs: defaultdict[str, dict[str, torch.Tensor]] = defaultdict(dict) |
| 197 | for key, tensor in adapter_weights.items(): |
| 198 | clean_key = key |
| 199 | if clean_key.startswith("base_model.model."): |
| 200 | clean_key = clean_key[len("base_model.model.") :] |
| 201 | |
| 202 | if ".lora_A." in clean_key: |
| 203 | base_key = clean_key.split(".lora_A.")[0] + ".weight" |
| 204 | lora_pairs[base_key]["A"] = tensor |
| 205 | elif ".lora_B." in clean_key: |
| 206 | base_key = clean_key.split(".lora_B.")[0] + ".weight" |
| 207 | lora_pairs[base_key]["B"] = tensor |
| 208 | |
| 209 | return dict(lora_pairs) |
| 210 | |
| 211 | |
| 212 | def _setup_jina_v5_task_and_pooler(model: nn.Module, vllm_config: VllmConfig) -> None: |
| 213 | """Shared init for jina-embeddings-v5 wrappers: select task + build pooler.""" |
| 214 | model._model_name = vllm_config.model_config.model |
| 215 | model._revision = vllm_config.model_config.revision |
| 216 | |
| 217 | model._task = getattr( |
| 218 | vllm_config.model_config.hf_config, "jina_task", _DEFAULT_TASK |
| 219 | ) |
| 220 | if model._task not in _SUPPORTED_TASKS: |
| 221 | logger.warning( |
| 222 | "Unknown jina_task=%r. Falling back to %r.", |
| 223 | model._task, |
| 224 | _DEFAULT_TASK, |
| 225 | ) |
| 226 | model._task = _DEFAULT_TASK |
| 227 | |
| 228 | pooler_config = vllm_config.model_config.pooler_config |
| 229 | assert pooler_config is not None |
| 230 | model.pooler = DispatchPooler.for_embedding(pooler_config) |
| 231 | |
| 232 | |
| 233 | def _load_jina_v5_weights( |
| 234 | model: nn.Module, weights: Iterable[tuple[str, torch.Tensor]] |
| 235 | ) -> set[str]: |
| 236 | """Shared loader: merge the selected task LoRA adapter into the base weights.""" |
| 237 | lora_pairs: dict[str, dict[str, torch.Tensor]] = {} |
| 238 | scaling = 1.0 |
| 239 | |
| 240 | result = _load_adapter(model._model_name, model._task, model._revision) |
| 241 | if result is None: |
| 242 | logger.warning( |
| 243 | "No adapter found for task %r in %r. Loading raw base weights.", |
| 244 | model._task, |
| 245 | model._model_name, |
| 246 | ) |
| 247 | else: |
| 248 | adapter_config, adapter_weights = result |
| 249 | scaling = adapter_config["lora_alpha"] / adapter_config["r"] |
| 250 | lora_pairs = _build_lora_pairs(adapter_weights) |
| 251 | logger.info( |
| 252 | "Loaded %d adapter tensors for task %r (scaling=%.4f, %d LoRA pairs)", |
| 253 | len(adapter_weights), |
| 254 | model._task, |
| 255 | scaling, |
| 256 | len(lora_pairs), |
| 257 | ) |
| 258 | |
| 259 | def _merge_weights( |
| 260 | weights: Iterable[tuple[str, torch.Tensor]], |
| 261 | ) -> Iterable[tuple[str, torch.Tensor]]: |
| 262 | for name, tensor in weights: |
| 263 | clean_name = name |
| 264 | if clean_name.startswith("model."): |
| 265 | clean_name = clean_name[len("model.") :] |
| 266 | |
| 267 | if clean_name in lora_pairs: |
| 268 | pair = lora_pairs[clean_name] |
| 269 | if "A" in pair and "B" in pair: |
| 270 | lora_A = pair["A"].to(device=tensor.device, dtype=tensor.dtype) |
| 271 | lora_B = pair["B"].to(device=tensor.device, dtype=tensor.dtype) |
| 272 | tensor = tensor + (lora_B @ lora_A) * scaling |
| 273 | yield name, tensor |
| 274 | |
| 275 | loader = AutoWeightsLoader(model, ignore_unexpected_prefixes=["lm_head."]) |
| 276 | weights = _merge_weights(weights) |
| 277 | return loader.load_weights(weights, mapper=model.hf_to_vllm_mapper) |
| 278 | |
| 279 | |
| 280 | class JinaEmbeddingsV5DecoderModel(Qwen3ForCausalLM, VllmModelForPooling): |
| 281 | """jina-embeddings-v5 with a Qwen3 decoder backbone (e.g. -small). |
| 282 | |
| 283 | Task-specific LoRA adapters are merged into the base weights at load time. |
| 284 | Declares itself a pooling model so that as_embedding_model() does not wrap it. |
| 285 | """ |
| 286 | |
| 287 | is_pooling_model = True |
| 288 | hf_to_vllm_mapper = Qwen3ForCausalLM.hf_to_vllm_mapper | WeightsMapper( |
| 289 | orig_to_new_prefix={"": "model."} |
| 290 | ) |
| 291 | |
| 292 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 293 | with no_init_weights( |
| 294 | self, |
| 295 | lambda mod: StageMissingLayer("output", mod), |
| 296 | targets=(LogitsProcessor, ParallelLMHead), |
| 297 | ): |
| 298 | super().__init__(vllm_config=vllm_config, prefix=prefix) |
| 299 | _setup_jina_v5_task_and_pooler(self, vllm_config) |
| 300 | |
| 301 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 302 | return _load_jina_v5_weights(self, weights) |
| 303 | |
| 304 | |
| 305 | class JinaEmbeddingsV5EncoderModel(LlamaForCausalLM, VllmModelForPooling): |
| 306 | """jina-embeddings-v5 with a bidirectional EuroBERT (Llama) encoder backbone. |
| 307 | |
| 308 | Used by encoder checkpoints such as jina-embeddings-v5-text-nano |
| 309 | (``is_decoder=False``). EuroBERT is architecturally a bidirectional Llama, so |
| 310 | the LlamaModel backbone switches to EncoderOnlyAttention when the config |
| 311 | carries ``is_causal=False`` (set by ``JinaEmbeddingsV5ModelConfig``). |
| 312 | """ |
| 313 | |
| 314 | is_pooling_model = True |
| 315 | hf_to_vllm_mapper = LlamaForCausalLM.hf_to_vllm_mapper | WeightsMapper( |
| 316 | orig_to_new_prefix={"": "model."} |
| 317 | ) |
| 318 | |
| 319 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 320 | with no_init_weights( |
| 321 | self, |
| 322 | lambda mod: StageMissingLayer("output", mod), |
| 323 | targets=(LogitsProcessor, ParallelLMHead), |
| 324 | ): |
| 325 | super().__init__(vllm_config=vllm_config, prefix=prefix) |
| 326 | _setup_jina_v5_task_and_pooler(self, vllm_config) |
| 327 | |
| 328 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 329 | return _load_jina_v5_weights(self, weights) |
| 330 | |
| 331 | |
| 332 | class JinaEmbeddingsV5Model(JinaEmbeddingsV5DecoderModel): |
| 333 | """Dispatcher for the jina-embeddings-v5 family. |
| 334 | |
| 335 | The family ships two backbones under one ``architectures`` entry: Qwen3 |
| 336 | decoders (-small) and bidirectional EuroBERT encoders (-nano), told apart by |
| 337 | ``is_decoder``. Inherits the decoder implementation so registry introspection |
| 338 | still sees a valid pooling model, and ``__new__`` swaps in the encoder |
| 339 | variant for encoder checkpoints. |
| 340 | """ |
| 341 | |
| 342 | def __new__(cls, *, vllm_config: VllmConfig, prefix: str = ""): |
| 343 | is_decoder = getattr(vllm_config.model_config.hf_config, "is_decoder", True) |
| 344 | if not is_decoder: |
| 345 | return JinaEmbeddingsV5EncoderModel(vllm_config=vllm_config, prefix=prefix) |
| 346 | return super().__new__(cls) |