@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 170 lines 6.19 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections.abc import Iterable, Mapping, Sequence
4
5 import torch
6 import torch.nn as nn
7 from transformers import (
8 BatchFeature,
9 PixtralVisionConfig,
10 )
11
12 from vllm.config import VllmConfig
13 from vllm.model_executor.models.mistral3 import (
14 Mistral3DummyInputsBuilder,
15 Mistral3ForConditionalGeneration,
16 Mistral3MultiModalProjector,
17 Mistral3ProcessingInfo,
18 init_vision_tower_for_mistral3,
19 )
20 from vllm.model_executor.models.pixtral import PixtralHFEncoderInfo
21 from vllm.model_executor.models.utils import (
22 AutoWeightsLoader,
23 WeightsMapper,
24 init_vllm_registered_model,
25 maybe_prefix,
26 )
27 from vllm.multimodal import MULTIMODAL_REGISTRY
28 from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
29 from vllm.multimodal.parse import ImageProcessorItems, MultiModalDataItems
30 from vllm.multimodal.processing import (
31 BaseMultiModalProcessor,
32 PromptReplacement,
33 PromptUpdate,
34 PromptUpdateDetails,
35 )
36
37
38 class LightOnOCRProcessingInfo(Mistral3ProcessingInfo):
39 def get_vision_encoder_info(
40 self,
41 mm_processor_kwargs: Mapping[str, object] | None = None,
42 ) -> PixtralHFEncoderInfo:
43 # LightOnOCR prompt updates use vision_config.image_size directly;
44 # preserve that behavior rather than inheriting Mistral3's processor
45 # longest_edge override.
46 return PixtralHFEncoderInfo(self.get_hf_config())
47
48
49 class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[LightOnOCRProcessingInfo]):
50 def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
51 return self.dummy_inputs.get_dummy_text(mm_counts)
52
53 def _postprocess_hf_mm_data(
54 self,
55 mm_data: Mapping[str, object],
56 hf_processor_mm_kwargs: Mapping[str, object],
57 processed_data: BatchFeature,
58 ) -> BatchFeature:
59 if not mm_data:
60 return processed_data
61
62 # un-pad pixel_values per-image so caches remain independent.
63 pixel_values = processed_data.get("pixel_values")
64 if pixel_values is not None:
65 image_sizes = processed_data["image_sizes"]
66 assert len(pixel_values) == len(image_sizes)
67 processed_data["pixel_values"] = [
68 p[:, :h, :w] for p, (h, w) in zip(pixel_values, image_sizes)
69 ]
70
71 return processed_data
72
73 def _get_mm_fields_config(
74 self,
75 hf_inputs: BatchFeature,
76 hf_processor_mm_kwargs: Mapping[str, object],
77 ) -> Mapping[str, MultiModalFieldConfig]:
78 return dict(
79 pixel_values=MultiModalFieldConfig.batched("image"),
80 image_embeds=MultiModalFieldConfig.batched("image"),
81 )
82
83 def _get_prompt_updates(
84 self,
85 mm_items: MultiModalDataItems,
86 hf_processor_mm_kwargs: Mapping[str, object],
87 out_mm_kwargs: MultiModalKwargsItems,
88 ) -> Sequence[PromptUpdate]:
89 hf_config = self.info.get_hf_config()
90 image_token_id = hf_config.image_token_index
91
92 assert isinstance(hf_config.vision_config, PixtralVisionConfig)
93 encoder_info = PixtralHFEncoderInfo(hf_config)
94
95 def replace(item_idx: int):
96 images = mm_items.get_items("image", ImageProcessorItems)
97 size = images.get_image_size(item_idx)
98 ncols, nrows = encoder_info.get_patch_grid_size(
99 image_width=size.width, image_height=size.height
100 )
101 # break/end tokens are not used in LightOnOCR
102 tokens = [image_token_id] * (ncols * nrows)
103 return PromptUpdateDetails.select_token_id(tokens, image_token_id)
104
105 return [
106 PromptReplacement(
107 modality="image", target=[image_token_id], replacement=replace
108 )
109 ]
110
111
112 @MULTIMODAL_REGISTRY.register_processor(
113 LightOnOCRMultiModalProcessor,
114 info=LightOnOCRProcessingInfo,
115 dummy_inputs=Mistral3DummyInputsBuilder,
116 )
117 class LightOnOCRForConditionalGeneration(Mistral3ForConditionalGeneration):
118 supports_mm_device_do_normalize = False
119
120 hf_to_vllm_mapper = WeightsMapper(
121 orig_to_new_prefix={
122 "model.vision_encoder.": "vision_tower.",
123 "model.vision_projection.": "multi_modal_projector.",
124 "lm_head.": "language_model.lm_head.",
125 "model.language_model.": "language_model.model.",
126 }
127 )
128
129 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
130 nn.Module.__init__(self)
131
132 config = vllm_config.model_config.hf_config
133 quant_config = vllm_config.quant_config
134 multimodal_config = vllm_config.model_config.multimodal_config
135
136 self.config = config
137 self.multimodal_config = multimodal_config
138
139 with self._mark_tower_model(vllm_config, "image"):
140 self.vision_tower = init_vision_tower_for_mistral3(
141 config,
142 quant_config=quant_config,
143 require_post_norm=False,
144 prefix=maybe_prefix(prefix, "vision_tower"),
145 )
146 self.multi_modal_projector = Mistral3MultiModalProjector(
147 vision_hidden_size=config.vision_config.hidden_size,
148 text_hidden_size=config.text_config.hidden_size,
149 projector_hidden_act=config.projector_hidden_act,
150 spatial_merge_size=config.spatial_merge_size,
151 patch_size=config.vision_config.patch_size,
152 multimodal_projector_bias=config.multimodal_projector_bias,
153 quant_config=quant_config,
154 prefix=maybe_prefix(prefix, "multi_modal_projector"),
155 )
156
157 with self._mark_language_model(vllm_config):
158 self.language_model = init_vllm_registered_model(
159 vllm_config=vllm_config,
160 hf_config=config.text_config,
161 prefix=maybe_prefix(prefix, "language_model"),
162 )
163
164 self.make_empty_intermediate_tensors = (
165 self.language_model.make_empty_intermediate_tensors
166 )
167
168 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
169 loader = AutoWeightsLoader(self)
170 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)