blip2.py 24.2 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
from collections.abc import Iterable, Mapping, Sequence
5
from typing import Annotated, Literal, TypeAlias
6
7
8

import torch
import torch.nn as nn
9
10
11
12
13
14
from transformers import (
    BatchFeature,
    Blip2Config,
    Blip2QFormerConfig,
    apply_chunking_to_forward,
)
15

16
from vllm.config import CacheConfig, VllmConfig
17
from vllm.config.multimodal import BaseDummyOptions
18
19
20
from vllm.model_executor.layers.activation import get_act_fn
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.multimodal import MULTIMODAL_REGISTRY
21
22
23
24
25
from vllm.multimodal.inputs import (
    MultiModalDataDict,
    MultiModalFieldConfig,
    MultiModalKwargsItems,
)
26
from vllm.multimodal.parse import MultiModalDataItems
27
from vllm.multimodal.processing import (
28
    BaseDummyInputsBuilder,
29
30
31
32
33
34
    BaseMultiModalProcessor,
    BaseProcessingInfo,
    PromptIndexTargets,
    PromptInsertion,
    PromptUpdate,
)
35
from vllm.sequence import IntermediateTensors
36
from vllm.utils.tensor_schema import TensorSchema, TensorShape
37

38
from .blip import BlipVisionModel, get_blip_num_patches
39
40
from .interfaces import (
    MultiModalEmbeddings,
41
    SupportsLoRA,
42
43
44
45
    SupportsMultiModal,
    SupportsPP,
    SupportsQuant,
)
46
from .module_mapping import MultiModelKeys
47
from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
48
49


50
51
52
53
54
55
56
57
class Blip2ImagePixelInputs(TensorSchema):
    """
    Dimensions:
        - bn: Batch size * number of images
        - c: Number of channels (3)
        - h: Height of each image
        - w: Width of each image
    """
58

59
    type: Literal["pixel_values"]
60
    data: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")]
61
62


63
class Blip2ImageEmbeddingInputs(TensorSchema):
64
    """
65
66
67
68
69
    Dimensions:
        - bn: Batch size * number of images
        - f: Image feature size
        - h: Hidden size (must match the hidden size of language model backbone)
    """
70

71
72
    type: Literal["image_embeds"]
    data: Annotated[torch.Tensor, TensorShape("bn", "f", "h")]
73
74


75
Blip2ImageInputs: TypeAlias = Blip2ImagePixelInputs | Blip2ImageEmbeddingInputs
76

77
78
79
80
81
82

class Blip2QFormerMultiHeadAttention(nn.Module):
    def __init__(
        self,
        config: Blip2QFormerConfig,
        *,
83
84
        quant_config: QuantizationConfig | None,
        cache_config: CacheConfig | None,
85
        is_cross_attention: bool = False,
86
        prefix: str = "",
87
88
89
90
91
92
93
94
95
96
97
98
    ) -> None:
        super().__init__()

        self.config = config

        if config.hidden_size % config.num_attention_heads != 0:
            raise ValueError(
                f"The hidden size ({config.hidden_size}) is not a multiple of "
                f"the number of attention heads ({config.num_attention_heads})"
            )

        self.num_attention_heads = config.num_attention_heads
99
        self.attention_head_size = config.hidden_size // config.num_attention_heads
100
101
102
103
104
105
106
107
108
109
110
        self.all_head_size = self.num_attention_heads * self.attention_head_size
        self.scaling = self.attention_head_size**-0.5

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        if is_cross_attention:
            kv_hidden_size = config.encoder_hidden_size
        else:
            kv_hidden_size = config.hidden_size
        self.key = nn.Linear(kv_hidden_size, self.all_head_size)
        self.value = nn.Linear(kv_hidden_size, self.all_head_size)

111
112
113
        self.position_embedding_type = getattr(
            config, "position_embedding_type", "absolute"
        )
114
        if self.position_embedding_type != "absolute":
115
116
117
            raise NotImplementedError(
                f"Unsupported position_embedding_type: {self.position_embedding_type}"
            )
118
119
120
121

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

    def transpose_for_scores(self, x):
122
        x = x.view(*x.size()[:-1], self.num_attention_heads, self.attention_head_size)
123
124
125
126
127
        return x.permute(0, 2, 1, 3)

    def forward(
        self,
        hidden_states: torch.Tensor,
128
        encoder_hidden_states: torch.FloatTensor | None = None,
129
130
131
132
    ):
        is_cross_attention = encoder_hidden_states is not None

        if is_cross_attention:
133
134
            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
135
136
137
138
139
140
141
142
        else:
            key_layer = self.transpose_for_scores(self.key(hidden_states))
            value_layer = self.transpose_for_scores(self.value(hidden_states))

        mixed_query_layer = self.query(hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer)

143
144
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
        attention_probs = torch.softmax(attention_scores * self.scaling, dim=-1)
145
146
147
148
149
150
151
152

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs_dropped = self.dropout(attention_probs)

        context_layer = torch.matmul(attention_probs_dropped, value_layer)

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
153
154
155
        context_layer = context_layer.view(
            *context_layer.size()[:-2], self.all_head_size
        )
156
157
158
159
160

        return context_layer


class Blip2QFormerSelfOutput(nn.Module):
161
    def __init__(self, config: Blip2QFormerConfig, prefix: str = "") -> None:
162
163
164
        super().__init__()

        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
165
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(
        self,
        hidden_states: torch.Tensor,
        input_tensor: torch.Tensor,
    ) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states


class Blip2QFormerAttention(nn.Module):
    def __init__(
        self,
        config: Blip2QFormerConfig,
        *,
184
185
        quant_config: QuantizationConfig | None,
        cache_config: CacheConfig | None,
186
        is_cross_attention: bool = False,
187
        prefix: str = "",
188
189
190
191
192
193
194
195
    ) -> None:
        super().__init__()

        self.attention = Blip2QFormerMultiHeadAttention(
            config,
            quant_config=quant_config,
            cache_config=cache_config,
            is_cross_attention=is_cross_attention,
196
            prefix=f"{prefix}.attention",
197
198
        )

199
        self.output = Blip2QFormerSelfOutput(config, prefix=f"{prefix}.output")
200
201
202
203

    def forward(
        self,
        hidden_states: torch.Tensor,
204
        encoder_hidden_states: torch.FloatTensor | None = None,
205
    ) -> tuple[torch.Tensor]:
206
207
208
209
210
211
212
213
214
215
        self_output = self.attention(
            hidden_states,
            encoder_hidden_states=encoder_hidden_states,
        )
        attention_output = self.output(self_output, hidden_states)

        return attention_output


class Blip2QFormerIntermediate(nn.Module):
216
    def __init__(self, config: Blip2QFormerConfig, prefix: str = "") -> None:
217
218
219
220
221
222
223
224
225
226
227
228
        super().__init__()

        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
        self.intermediate_act_fn = get_act_fn(config.hidden_act)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        return hidden_states


class Blip2QFormerOutput(nn.Module):
229
    def __init__(self, config: Blip2QFormerConfig, prefix: str = "") -> None:
230
231
232
        super().__init__()

        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
233
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(
        self,
        hidden_states: torch.Tensor,
        input_tensor: torch.Tensor,
    ) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states


class Blip2QFormerLayer(nn.Module):
    def __init__(
        self,
        config: Blip2QFormerConfig,
        *,
252
253
        quant_config: QuantizationConfig | None,
        cache_config: CacheConfig | None,
254
        layer_idx: int,
255
        prefix: str = "",
256
257
258
259
260
    ) -> None:
        super().__init__()

        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1
261
262
263
264
265
266
        self.attention = Blip2QFormerAttention(
            config,
            quant_config=quant_config,
            cache_config=cache_config,
            prefix=f"{prefix}.attention",
        )
267
268
269
270
271
272
273
274

        self.layer_idx = layer_idx

        if layer_idx % config.cross_attention_frequency == 0:
            self.crossattention = Blip2QFormerAttention(
                config,
                quant_config=quant_config,
                cache_config=cache_config,
275
                is_cross_attention=True,
276
277
                prefix=f"{prefix}.crossattention",
            )
278
279
280
281
            self.has_cross_attention = True
        else:
            self.has_cross_attention = False

282
        self.intermediate_query = Blip2QFormerIntermediate(
283
284
285
            config, prefix=f"{prefix}.intermediate_query"
        )
        self.output_query = Blip2QFormerOutput(config, prefix=f"{prefix}.output_query")
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317

    def forward(
        self,
        hidden_states: torch.FloatTensor,
        encoder_hidden_states: torch.FloatTensor,
        query_length: int,
    ):
        attention_output = self.attention(hidden_states)

        if query_length > 0:
            query_attention_output = attention_output[:, :query_length, :]

            if self.has_cross_attention:
                query_attention_output = self.crossattention(
                    query_attention_output,
                    encoder_hidden_states=encoder_hidden_states,
                )

            layer_output = apply_chunking_to_forward(
                self.feed_forward_chunk_query,
                self.chunk_size_feed_forward,
                self.seq_len_dim,
                query_attention_output,
            )

            if attention_output.shape[1] > query_length:
                layer_output_text = apply_chunking_to_forward(
                    self.feed_forward_chunk,
                    self.chunk_size_feed_forward,
                    self.seq_len_dim,
                    attention_output[:, query_length:, :],
                )
318
                layer_output = torch.cat([layer_output, layer_output_text], dim=1)
319
320
321
322
323
324
325
326
327
328
        else:
            layer_output = apply_chunking_to_forward(
                self.feed_forward_chunk,
                self.chunk_size_feed_forward,
                self.seq_len_dim,
                attention_output,
            )

        return layer_output

329
    def feed_forward_chunk(self, attention_output: torch.Tensor) -> torch.Tensor:
330
331
332
333
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output)
        return layer_output

334
    def feed_forward_chunk_query(self, attention_output: torch.Tensor) -> torch.Tensor:
335
336
337
338
339
340
341
342
343
344
        intermediate_output = self.intermediate_query(attention_output)
        layer_output = self.output_query(intermediate_output, attention_output)
        return layer_output


class Blip2QFormerEncoder(nn.Module):
    def __init__(
        self,
        config: Blip2QFormerConfig,
        *,
345
346
        quant_config: QuantizationConfig | None,
        cache_config: CacheConfig | None,
347
        prefix: str = "",
348
349
350
351
352
    ) -> None:
        super().__init__()

        self.config = config

353
354
355
356
357
358
359
360
361
362
363
364
        self.layer = nn.ModuleList(
            [
                Blip2QFormerLayer(
                    config,
                    quant_config=quant_config,
                    cache_config=cache_config,
                    layer_idx=layer_idx,
                    prefix=f"{prefix}.layer.{layer_idx}",
                )
                for layer_idx in range(config.num_hidden_layers)
            ]
        )
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389

    def forward(
        self,
        hidden_states: torch.FloatTensor,
        encoder_hidden_states: torch.FloatTensor,
        query_length: int,
    ) -> torch.Tensor:
        for i in range(self.config.num_hidden_layers):
            layer_module = self.layer[i]

            hidden_states = layer_module(
                hidden_states,
                encoder_hidden_states=encoder_hidden_states,
                query_length=query_length,
            )

        return hidden_states


# Adapted from https://github.com/huggingface/transformers/blob/v4.41.2/src/transformers/models/blip_2/modeling_blip_2.py#L1025
class Blip2QFormerModel(nn.Module):
    def __init__(
        self,
        config: Blip2QFormerConfig,
        *,
390
391
        quant_config: QuantizationConfig | None,
        cache_config: CacheConfig | None,
392
        prefix: str = "",
393
394
395
396
397
    ) -> None:
        super().__init__()

        self.config = config

398
        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
399
400
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

401
402
403
404
405
406
        self.encoder = Blip2QFormerEncoder(
            config,
            quant_config=quant_config,
            cache_config=cache_config,
            prefix=f"{prefix}.encoder",
        )
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426

    def forward(
        self,
        query_embeds: torch.FloatTensor,
        encoder_hidden_states: torch.FloatTensor,
    ) -> torch.Tensor:
        query_length = query_embeds.shape[1]

        embedding_output = self.layernorm(query_embeds)
        embedding_output = self.dropout(embedding_output)

        sequence_output = self.encoder(
            embedding_output,
            encoder_hidden_states=encoder_hidden_states,
            query_length=query_length,
        )

        return sequence_output


427
428
class Blip2ProcessingInfo(BaseProcessingInfo):
    def get_hf_config(self):
429
        return self.ctx.get_hf_config(Blip2Config)
430

431
    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
432
433
        return {"image": 1}

434
435
436
437
438
439
    def get_num_image_tokens(self) -> int:
        hf_config = self.get_hf_config()
        return hf_config.num_query_tokens


class Blip2DummyInputsBuilder(BaseDummyInputsBuilder[Blip2ProcessingInfo]):
440
441
442
443
    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        return ""

    def get_dummy_mm_data(
444
445
446
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
447
        mm_options: Mapping[str, BaseDummyOptions] | None = None,
448
    ) -> MultiModalDataDict:
449
        hf_config = self.info.get_hf_config()
450
451
452
453
454
        vision_config = hf_config.vision_config

        max_image_size = vision_config.image_size
        num_images = mm_counts.get("image", 0)

455
456
        image_overrides = mm_options.get("image") if mm_options else None

457
        return {
458
459
460
461
462
463
            "image": self._get_dummy_images(
                width=max_image_size,
                height=max_image_size,
                num_images=num_images,
                overrides=image_overrides,
            )
464
465
466
        }


467
class Blip2MultiModalProcessor(BaseMultiModalProcessor[Blip2ProcessingInfo]):
468
469
470
471
472
    def _call_hf_processor(
        self,
        prompt: str,
        mm_data: Mapping[str, object],
        mm_kwargs: Mapping[str, object],
473
        tok_kwargs: Mapping[str, object],
474
475
476
477
478
479
480
481
482
483
484
    ) -> BatchFeature:
        if not mm_data:
            # HF processor always adds placeholders even when there's no image
            tokenizer = self.info.get_tokenizer()
            prompt_ids = tokenizer.encode(prompt)
            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")

        return super()._call_hf_processor(
            prompt=prompt,
            mm_data=mm_data,
            mm_kwargs=mm_kwargs,
485
            tok_kwargs=tok_kwargs,
486
487
        )

488
489
490
491
492
493
494
495
496
    def _get_mm_fields_config(
        self,
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
        return dict(
            pixel_values=MultiModalFieldConfig.batched("image"),
            image_embeds=MultiModalFieldConfig.batched("image"),
        )
497

498
    def _get_prompt_updates(
499
500
501
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
502
        out_mm_kwargs: MultiModalKwargsItems,
503
    ) -> Sequence[PromptUpdate]:
504
505
506
        tokenizer = self.info.get_tokenizer()
        vocab = tokenizer.get_vocab()

507
        image_token_id = vocab["<image>"]
508
        num_image_tokens = self.info.get_num_image_tokens()
509
        image_tokens = [image_token_id] * num_image_tokens
510
511

        return [
512
            PromptInsertion(
513
                modality="image",
514
                target=PromptIndexTargets.start(),
515
                insertion=image_tokens,
516
517
            )
        ]
518
519


520
521
522
523
524
525
@MULTIMODAL_REGISTRY.register_processor(
    Blip2MultiModalProcessor,
    info=Blip2ProcessingInfo,
    dummy_inputs=Blip2DummyInputsBuilder,
)
class Blip2ForConditionalGeneration(
526
    nn.Module, SupportsLoRA, SupportsMultiModal, SupportsPP, SupportsQuant
527
):
528
    @classmethod
529
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
530
531
532
533
534
        if modality.startswith("image"):
            return None

        raise ValueError("Only image modality is supported")

535
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
536
        super().__init__()
537
538
539
540
        config = vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
        multimodal_config = vllm_config.model_config.multimodal_config
541
542
        self.config = config
        self.multimodal_config = multimodal_config
543
544
545
546
547
548
549
550
        vision_config = config.vision_config
        self._vision_tokens_per_image = (
            get_blip_num_patches(
                image_size=vision_config.image_size,
                patch_size=vision_config.patch_size,
            )
            + 1  # include class token
        )
551

552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
        with self._mark_tower_model(vllm_config, "image"):
            self.vision_model = BlipVisionModel(vision_config, quant_config)
            self.query_tokens = nn.Parameter(
                torch.zeros(
                    1, config.num_query_tokens, config.qformer_config.hidden_size
                )
            )
            self.qformer = Blip2QFormerModel(
                config.qformer_config,
                cache_config=cache_config,
                quant_config=quant_config,
                prefix=f"{prefix}.qformer",
            )
            self.language_projection = nn.Linear(
                config.qformer_config.hidden_size,
                config.text_config.hidden_size,
                bias=True,
            )
570

571
572
573
574
575
576
        with self._mark_language_model(vllm_config):
            self.language_model = init_vllm_registered_model(
                vllm_config=vllm_config,
                hf_config=config.text_config,
                prefix=maybe_prefix(prefix, "language_model"),
            )
577

578
        self.make_empty_intermediate_tensors = (
579
580
            self.language_model.make_empty_intermediate_tensors
        )
581

582
    def _parse_and_validate_image_input(
583
        self, **kwargs: object
584
    ) -> Blip2ImageInputs | None:
585
        pixel_values = kwargs.pop("pixel_values", None)
586
        image_embeds = kwargs.pop("image_embeds", None)
587

588
        if pixel_values is None and image_embeds is None:
589
590
            return None

591
        if pixel_values is not None:
592
            expected_h = expected_w = self.config.vision_config.image_size
593
594
595
596
597
            return Blip2ImagePixelInputs(
                type="pixel_values",
                data=pixel_values,
                resolve_bindings={"h": expected_h, "w": expected_w},
            )
598
599
600
601

        if image_embeds is not None:
            return Blip2ImageEmbeddingInputs(
                type="image_embeds",
602
                data=image_embeds,
603
604
605
            )

        raise AssertionError("This line should be unreachable.")
606

607
608
609
    def _image_pixels_to_features(
        self, vision_model: BlipVisionModel, pixel_values: torch.Tensor
    ) -> torch.Tensor:
610
611
612
613
614
615
        # NOTE: we skip the step to select the vision feature layer since
        # this is already done inside the vision tower
        image_features = vision_model(pixel_values)

        return image_features

616
    def _process_image_pixels(self, inputs: Blip2ImagePixelInputs) -> torch.Tensor:
617
618
619
620
        pixel_values = inputs["data"]

        return self._image_pixels_to_features(self.vision_model, pixel_values)

621
    def _process_image_input(self, image_input: Blip2ImageInputs) -> torch.Tensor:
622
623
624
        if image_input["type"] == "image_embeds":
            return image_input["data"]

625
626
        image_features = self._process_image_pixels(image_input)

627
        query_tokens = self.query_tokens.expand(image_features.shape[0], -1, -1)
628
629
630
631
632
633
634
        query_output = self.qformer(
            query_embeds=query_tokens,
            encoder_hidden_states=image_features,
        )

        return self.language_projection(query_output)

635
    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
636
637
        image_input = self._parse_and_validate_image_input(**kwargs)
        if image_input is None:
638
            return []
639
640
641
        vision_embeddings = self._process_image_input(image_input)
        return vision_embeddings

642
643
644
645
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
646
647
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
648
        **kwargs: object,
649
    ) -> IntermediateTensors:
650
651
652
653
654
655
656
657
658
659
660
661
        """Run forward pass for BLIP-2.

        One key thing to understand is the `input_ids` already accounts for the
        positions of the to-be-inserted image embeddings.

        Concretely, consider a text prompt:
        `"Question: What's the content of the image? Answer:"`.

        Tokenizer outputs:
        `[2, 45641, 35, 653, 18, 5, 1383, 9, 5, 2274, 116, 31652, 35]`.

        To reserve space in KV cache, we have to insert placeholder tokens
662
        before they are inputted to the model, so the input processor prepends
663
664
665
666
667
668
669
670
671
672
673
674
        dummy tokens (denoted as `50265`), resulting in:
        `[50265, ..., 50265, 2, 45641, 35, ..., 31652, 35]`.

        We insert 32 tokens since it corresponds to the number of query
        embeddings outputted by the Q-Former and inputted to the language model.

        This way, the `positions` and `attn_metadata` are consistent
        with the `input_ids`.

        Args:
            input_ids: Flattened (concatenated) input_ids corresponding to a
                batch.
675

676
        Info:
samzong's avatar
samzong committed
677
            [`Blip2ImageInputs`][vllm.model_executor.models.blip2.Blip2ImageInputs]
678
        """
679

680
        if intermediate_tensors is not None:
681
            inputs_embeds = None
682

683
684
685
        hidden_states = self.language_model.model(
            input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
        )
686
687
688

        return hidden_states

689
690
691
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
692
    ) -> torch.Tensor | None:
693
        return self.language_model.compute_logits(hidden_states)
694

695
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
696
        loader = AutoWeightsLoader(self)
697
        return loader.load_weights(weights)
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730

    def get_mm_mapping(self) -> MultiModelKeys:
        return MultiModelKeys.from_string_field(
            language_model="language_model",
            connector=["qformer", "language_projection"],
            tower_model="vision_model",
        )

    def get_num_mm_encoder_tokens(
        self,
        num_image_tokens: int,
    ) -> int:
        if num_image_tokens <= 0:
            return 0
        assert num_image_tokens % self.config.num_query_tokens == 0, (
            "The number of image tokens must be a multiple of "
            "the number of query tokens."
        )
        num_images = num_image_tokens / self.config.num_query_tokens
        return num_images * self._vision_tokens_per_image

    def get_num_mm_connector_tokens(
        self,
        num_vision_tokens: int,
    ) -> int:
        if num_vision_tokens <= 0:
            return 0
        assert num_vision_tokens % self._vision_tokens_per_image == 0, (
            "The number of vision tokens must be a multiple of "
            "the number of tokens per image."
        )
        num_images = num_vision_tokens / self._vision_tokens_per_image
        return num_images * self.config.num_query_tokens