main
py 103 lines 3.36 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Mapping
5
6 import torch
7 import torch.nn as nn
8 from transformers.activations import GELUActivation
9
10 from vllm.config import VllmConfig
11 from vllm.config.multimodal import MultiModalDummyOptions
12 from vllm.inputs import MultiModalDataDict
13 from vllm.multimodal import MULTIMODAL_REGISTRY
14
15 from .llava_next import (
16 LlavaDummyInputsBuilder,
17 LlavaNextMultiModalProcessor,
18 LlavaNextProcessingInfo,
19 )
20 from .llava_onevision import LlavaOnevisionForConditionalGeneration
21 from .utils import WeightsMapper
22
23
24 class RVLProcessingInfo(LlavaNextProcessingInfo):
25 def get_hf_config(self):
26 return self.ctx.get_hf_config()
27
28 def get_hf_processor(self, **kwargs: object):
29 return self.ctx.get_hf_processor(**kwargs)
30
31
32 class RVLDummyInputsBuilder(LlavaDummyInputsBuilder[RVLProcessingInfo]):
33 def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
34 num_images = mm_counts.get("image", 0)
35 image_token = "<image>"
36
37 return image_token * num_images
38
39 def get_dummy_mm_data(
40 self,
41 seq_len: int,
42 mm_counts: Mapping[str, int],
43 mm_options: MultiModalDummyOptions,
44 ) -> MultiModalDataDict:
45 target_width, target_height = self.info.get_image_size_with_most_features()
46
47 return {
48 "image": self._get_dummy_images(
49 width=target_width,
50 height=target_height,
51 num_images=mm_counts.get("image", 0),
52 overrides=mm_options.get("image"),
53 ),
54 }
55
56
57 class RVLMultiModalProjector(nn.Module):
58 def __init__(self, config):
59 super().__init__()
60 self.pre_norm = nn.LayerNorm(config.vision_config.hidden_size, eps=1e-06)
61 self.linear_1 = nn.Linear(
62 config.vision_config.hidden_size,
63 config.text_config.hidden_size,
64 bias=True,
65 )
66 self.act = GELUActivation()
67 self.linear_2 = nn.Linear(
68 config.text_config.hidden_size,
69 config.text_config.hidden_size,
70 bias=True,
71 )
72
73 def forward(self, image_feature: torch.Tensor) -> torch.Tensor:
74 image_feature = self.pre_norm(image_feature)
75 hidden_states = self.linear_1(image_feature)
76 hidden_states = self.act(hidden_states)
77 hidden_states = self.linear_2(hidden_states)
78
79 return hidden_states
80
81
82 @MULTIMODAL_REGISTRY.register_processor(
83 LlavaNextMultiModalProcessor,
84 info=RVLProcessingInfo,
85 dummy_inputs=RVLDummyInputsBuilder,
86 )
87 class RForConditionalGeneration(LlavaOnevisionForConditionalGeneration):
88 hf_to_vllm_mapper = WeightsMapper(
89 orig_to_new_prefix={
90 # mapping for new names in checkpoint saved after transformers
91 # v4.52
92 "model.language_model.": "language_model.model.",
93 "model.vision_tower.": "vision_tower.",
94 "model.multi_modal_projector.": "multi_modal_projector.",
95 "model.image_newline": "image_newline",
96 "lm_head.": "language_model.lm_head.",
97 }
98 )
99
100 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
101 super().__init__(vllm_config=vllm_config, prefix=prefix)
102 config = vllm_config.model_config.hf_config
103 self.multi_modal_projector = RVLMultiModalProjector(config)