Sync from v0.13

vllm/model_executor/models/jina_vl.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Iterable, Mapping
+
+import torch
+import torch.nn as nn
+from transformers import BatchFeature
+
+from vllm.config import ModelConfig, VllmConfig
+from vllm.inputs import TokensPrompt
+from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
+from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.sequence import IntermediateTensors
+
+from .interfaces import SupportsCrossEncoding, SupportsMultiModal, SupportsScoreTemplate
+from .qwen2_vl import (
+    Qwen2VLDummyInputsBuilder,
+    Qwen2VLForConditionalGeneration,
+    Qwen2VLMultiModalProcessor,
+    Qwen2VLProcessingInfo,
+)
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
+
+logger = init_logger(__name__)
+
+
+class JinaVLScorer(nn.Module):
+    def __init__(self, model_config: "ModelConfig"):
+        super().__init__()
+        config = model_config.hf_config.get_text_config()
+        head_dtype = model_config.head_dtype
+        self.dense = ColumnParallelLinear(
+            config.hidden_size, config.hidden_size, params_dtype=head_dtype, bias=True
+        )
+        self.out_proj = RowParallelLinear(
+            config.hidden_size, config.num_labels, params_dtype=head_dtype, bias=True
+        )
+
+    def forward(self, x, **kwargs):
+        x, _ = self.dense(x)
+        x = torch.relu(x)
+        x, _ = self.out_proj(x)
+        return x
+
+
+class JinaVLMultiModalProcessor(Qwen2VLMultiModalProcessor):
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, object],
+        tok_kwargs: Mapping[str, object],
+    ) -> BatchFeature:
+        # NOTE: We should reverse the order of the mm_data because the
+        # query prompt is placed after the document prompt in the score
+        # template for JinaVLForRanking model, but in mm_data they are
+        # stored in the opposite order (query first, then document).
+        for _, value in mm_data.items():
+            value.reverse()
+        return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs)
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    JinaVLMultiModalProcessor,
+    info=Qwen2VLProcessingInfo,
+    dummy_inputs=Qwen2VLDummyInputsBuilder,
+)
+class JinaVLForSequenceClassification(
+    Qwen2VLForConditionalGeneration,
+    SupportsCrossEncoding,
+    SupportsMultiModal,
+    SupportsScoreTemplate,
+):
+    is_pooling_model = True
+    weight_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "score.0.": "score.dense.",
+            "score.2.": "score.out_proj.",
+            # mapping for new names in checkpoint saved after transformers v4.52
+            "model.language_model.": "language_model.model.",
+            "visual.": "visual.",
+            # mapping for original checkpoint
+            "lm_head.": "language_model.lm_head.",
+            "model.": "language_model.model.",
+        }
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "qwen2_vl")
+        )
+        pooler_config = vllm_config.model_config.pooler_config
+        assert pooler_config is not None
+
+        self.score = JinaVLScorer(vllm_config.model_config)
+        self.pooler = DispatchPooler(
+            {
+                "token_classify": Pooler.for_token_classify(
+                    pooler_config, classifier=self.score
+                ),
+                "classify": Pooler.for_classify(
+                    pooler_config, classifier=self.score, act_fn="classify"
+                ),
+                "score": Pooler.for_classify(
+                    pooler_config, classifier=self.score, act_fn="score"
+                ),
+            }
+        )
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        if modality.startswith("image"):
+            return "<|vision_start|><|image_pad|><|vision_end|>"
+
+        raise ValueError("Only image modality is supported")
+
+    @classmethod
+    def get_score_template(cls, query: str, document: str) -> str | None:
+        return f"**Document**:\n{document}\n**Query**:\n{query}"
+
+    @classmethod
+    def post_process_tokens(cls, prompt: TokensPrompt) -> None:
+        # add score target token at the end of prompt tokens
+        prompt["prompt_token_ids"].append(100)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor:
+        hidden_states = super().forward(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.weight_mapper)