@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 171 lines 5.66 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections.abc import Iterable
4
5 import torch
6 import torch.nn as nn
7
8 from vllm.config import ModelConfig, VllmConfig
9 from vllm.inputs import TokensPrompt
10 from vllm.logger import init_logger
11 from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
12 from vllm.model_executor.layers.pooler import DispatchPooler
13 from vllm.multimodal import MULTIMODAL_REGISTRY
14 from vllm.multimodal.inputs import MultiModalKwargsItems
15 from vllm.multimodal.processing.processor import (
16 MultiModalProcessingResult,
17 ProcessorInputs,
18 TimingContext,
19 )
20 from vllm.sequence import IntermediateTensors
21
22 from .interfaces import SupportsCrossEncoding, SupportsMultiModal, SupportsScoreTemplate
23 from .qwen2_vl import (
24 Qwen2VLDummyInputsBuilder,
25 Qwen2VLForConditionalGeneration,
26 Qwen2VLMultiModalProcessor,
27 Qwen2VLProcessingInfo,
28 )
29 from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
30
31 logger = init_logger(__name__)
32
33
34 class JinaVLScorer(nn.Module):
35 def __init__(self, model_config: "ModelConfig", prefix: str = ""):
36 super().__init__()
37 config = model_config.hf_config.get_text_config()
38 head_dtype = model_config.head_dtype
39 self.dense = ColumnParallelLinear(
40 config.hidden_size,
41 config.hidden_size,
42 params_dtype=head_dtype,
43 bias=True,
44 prefix=f"{prefix}.dense",
45 )
46 self.out_proj = RowParallelLinear(
47 config.hidden_size,
48 config.num_labels,
49 params_dtype=head_dtype,
50 bias=True,
51 prefix=f"{prefix}.out_proj",
52 )
53
54 def forward(self, x, **kwargs):
55 x, _ = self.dense(x)
56 x = torch.relu(x)
57 x, _ = self.out_proj(x)
58 return x
59
60
61 class JinaVLMultiModalProcessor(Qwen2VLMultiModalProcessor):
62 def _cached_apply_hf_processor(
63 self,
64 inputs: ProcessorInputs,
65 timing_ctx: TimingContext,
66 ) -> MultiModalProcessingResult:
67 mm_res = super()._cached_apply_hf_processor(inputs, timing_ctx)
68
69 # Score inputs are query-first, while the prompt template is document-first.
70 mm_kwargs = MultiModalKwargsItems(
71 {
72 modality: list(reversed(items))
73 for modality, items in mm_res.kwargs.items()
74 }
75 )
76 mm_hashes = {
77 modality: list(reversed(hashes))
78 for modality, hashes in mm_res.hashes.items()
79 }
80 mm_prompt_updates = {
81 modality: [
82 [
83 self._recompute_cached_prompt_update(update, item_idx)
84 for update in updates
85 ]
86 for item_idx, updates in enumerate(reversed(item_updates))
87 ]
88 for modality, item_updates in mm_res.prompt_updates.items()
89 }
90
91 return MultiModalProcessingResult(
92 prompt_ids=mm_res.prompt_ids,
93 kwargs=mm_kwargs,
94 hashes=mm_hashes,
95 prompt_updates=mm_prompt_updates,
96 )
97
98
99 @MULTIMODAL_REGISTRY.register_processor(
100 JinaVLMultiModalProcessor,
101 info=Qwen2VLProcessingInfo,
102 dummy_inputs=Qwen2VLDummyInputsBuilder,
103 )
104 class JinaVLForSequenceClassification(
105 Qwen2VLForConditionalGeneration,
106 SupportsCrossEncoding,
107 SupportsMultiModal,
108 SupportsScoreTemplate,
109 ):
110 is_pooling_model = True
111 weight_mapper = WeightsMapper(
112 orig_to_new_prefix={
113 "score.0.": "score.dense.",
114 "score.2.": "score.out_proj.",
115 # mapping for new names in checkpoint saved after transformers v4.52
116 "model.language_model.": "language_model.model.",
117 "visual.": "visual.",
118 # mapping for original checkpoint
119 "lm_head.": "language_model.lm_head.",
120 "model.": "language_model.model.",
121 }
122 )
123
124 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
125 super().__init__(
126 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "qwen2_vl")
127 )
128 pooler_config = vllm_config.model_config.pooler_config
129 assert pooler_config is not None
130
131 self.score = JinaVLScorer(
132 vllm_config.model_config, prefix=maybe_prefix(prefix, "score")
133 )
134 self.pooler = DispatchPooler.for_seq_cls(pooler_config, classifier=self.score)
135
136 @classmethod
137 def get_placeholder_str(cls, modality: str, i: int) -> str | None:
138 if modality.startswith("image"):
139 return "<|vision_start|><|image_pad|><|vision_end|>"
140
141 raise ValueError("Only image modality is supported")
142
143 @classmethod
144 def get_score_template(cls, query: str, document: str) -> str | None:
145 return f"**Document**:\n{document}\n**Query**:\n{query}"
146
147 @classmethod
148 def post_process_tokens(cls, prompt: TokensPrompt) -> None:
149 # add score target token at the end of prompt tokens
150 prompt["prompt_token_ids"].append(100)
151
152 def forward(
153 self,
154 input_ids: torch.Tensor | None,
155 positions: torch.Tensor,
156 intermediate_tensors: IntermediateTensors | None = None,
157 inputs_embeds: torch.Tensor | None = None,
158 **kwargs: object,
159 ) -> torch.Tensor:
160 hidden_states = super().forward(
161 input_ids=input_ids,
162 positions=positions,
163 intermediate_tensors=intermediate_tensors,
164 inputs_embeds=inputs_embeds,
165 **kwargs,
166 )
167 return hidden_states
168
169 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
170 loader = AutoWeightsLoader(self)
171 return loader.load_weights(weights, mapper=self.weight_mapper)