Aluode/PerceptionLabPortable

1#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨2#           This file was automatically generated from src/transformers/models/got_ocr2/modular_got_ocr2.py.3#               Do NOT edit this file manually as any edits will be overwritten by the generation of4#             the file from the modular. If any change should be done, please apply the change to the5#                          modular_got_ocr2.py file directly. One of our CI enforces this.6#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨7# coding=utf-88# Copyright 2024 HuggingFace Inc. team. All rights reserved.9#10# Licensed under the Apache License, Version 2.0 (the "License");11# you may not use this file except in compliance with the License.12# You may obtain a copy of the License at13#14#     http://www.apache.org/licenses/LICENSE-2.015#16# Unless required by applicable law or agreed to in writing, software17# distributed under the License is distributed on an "AS IS" BASIS,18# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.19# See the License for the specific language governing permissions and20# limitations under the License.21 22 23import collections24from dataclasses import dataclass25from typing import Optional, Union26 27import torch28import torch.nn as nn29import torch.nn.functional as F30 31from transformers.utils.generic import check_model_inputs32 33from ...activations import ACT2FN34from ...cache_utils import Cache35from ...generation import GenerationMixin36from ...modeling_layers import GradientCheckpointingLayer37from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput38from ...modeling_utils import PreTrainedModel39from ...processing_utils import Unpack40from ...utils import TransformersKwargs, auto_docstring, can_return_tuple41from ..auto import AutoModel42from .configuration_got_ocr2 import GotOcr2Config, GotOcr2VisionConfig43 44 45class GotOcr2MLPBlock(nn.Module):46    def __init__(self, config):47        super().__init__()48        self.lin1 = nn.Linear(config.hidden_size, config.mlp_dim)49        self.lin2 = nn.Linear(config.mlp_dim, config.hidden_size)50        self.act = ACT2FN[config.hidden_act]51 52    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:53        hidden_states = self.lin1(hidden_states)54        hidden_states = self.act(hidden_states)55        hidden_states = self.lin2(hidden_states)56        return hidden_states57 58 59class GotOcr2VisionAttention(nn.Module):60    """Multi-head Attention block with relative position embeddings."""61 62    def __init__(self, config, window_size):63        super().__init__()64        input_size = (65            (config.image_size // config.patch_size, config.image_size // config.patch_size)66            if window_size == 067            else (window_size, window_size)68        )69 70        self.num_attention_heads = config.num_attention_heads71        head_dim = config.hidden_size // config.num_attention_heads72        self.scale = head_dim**-0.573        self.dropout = config.attention_dropout74 75        self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=config.qkv_bias)76        self.proj = nn.Linear(config.hidden_size, config.hidden_size)77 78        self.use_rel_pos = config.use_rel_pos79        if self.use_rel_pos:80            if input_size is None:81                raise ValueError("Input size must be provided if using relative positional encoding.")82 83            # initialize relative positional embeddings84            self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim))85            self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))86 87    def get_rel_pos(self, q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:88        """89        Get relative positional embeddings according to the relative positions of90            query and key sizes.91 92        Args:93            q_size (int):94                size of the query.95            k_size (int):96                size of key k.97            rel_pos (`torch.Tensor`):98                relative position embeddings (L, channel).99 100        Returns:101            Extracted positional embeddings according to relative positions.102        """103        max_rel_dist = int(2 * max(q_size, k_size) - 1)104        # Interpolate rel pos.105        rel_pos_resized = F.interpolate(106            rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),107            size=max_rel_dist,108            mode="linear",109        )110        rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)111 112        # Scale the coords with short length if shapes for q and k are different.113        q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)114        k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)115        relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)116 117        return rel_pos_resized[relative_coords.long()]118 119    def get_decomposed_rel_pos(120        self,121        query: torch.Tensor,122        rel_pos_h: torch.Tensor,123        rel_pos_w: torch.Tensor,124        q_size: tuple[int, int],125        k_size: tuple[int, int],126    ) -> torch.Tensor:127        """128        Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.129        https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py130 131        Args:132            query (`torch.Tensor`):133                query q in the attention layer with shape (batch_size, query_height * query_width, channel).134            rel_pos_h (`torch.Tensor`):135                relative position embeddings (Lh, channel) for height axis.136            rel_pos_w (`torch.Tensor`):137                relative position embeddings (Lw, channel) for width axis.138            q_size (tuple):139                spatial sequence size of query q with (query_height, query_width).140            k_size (tuple):141                spatial sequence size of key k with (key_height, key_width).142 143        Returns:144            decomposed_rel_pos (`torch.Tensor`):145                decomposed relative position embeddings.146        """147        query_height, query_width = q_size148        key_height, key_width = k_size149        relative_position_height = self.get_rel_pos(query_height, key_height, rel_pos_h)150        relative_position_width = self.get_rel_pos(query_width, key_width, rel_pos_w)151 152        batch_size, _, dim = query.shape153        reshaped_query = query.reshape(batch_size, query_height, query_width, dim)154        rel_h = torch.einsum("bhwc,hkc->bhwk", reshaped_query, relative_position_height)155        rel_w = torch.einsum("bhwc,wkc->bhwk", reshaped_query, relative_position_width)156 157        decomposed_rel_pos = rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :]158 159        return decomposed_rel_pos160 161    def forward(self, hidden_states: torch.Tensor, output_attentions=None) -> tuple[torch.Tensor, torch.Tensor]:162        batch_size, height, width, _ = hidden_states.shape163        # qkv with shape (3, batch_size, nHead, height * width, channel)164        qkv = (165            self.qkv(hidden_states)166            .reshape(batch_size, height * width, 3, self.num_attention_heads, -1)167            .permute(2, 0, 3, 1, 4)168        )169        # q, k, v with shape (batch_size * nHead, height * width, channel)170        query, key, value = qkv.reshape(3, batch_size * self.num_attention_heads, height * width, -1).unbind(0)171 172        attn_weights = (query * self.scale) @ key.transpose(-2, -1)173 174        if self.use_rel_pos:175            decomposed_rel_pos = self.get_decomposed_rel_pos(176                query, self.rel_pos_h, self.rel_pos_w, (height, width), (height, width)177            )178            decomposed_rel_pos = decomposed_rel_pos.reshape_as(attn_weights)179            attn_weights = attn_weights + decomposed_rel_pos180 181        attn_weights = torch.nn.functional.softmax(attn_weights, dtype=torch.float32, dim=-1).to(query.dtype)182 183        attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)184 185        attn_output = (attn_probs @ value).reshape(batch_size, self.num_attention_heads, height, width, -1)186        attn_output = attn_output.permute(0, 2, 3, 1, 4).reshape(batch_size, height, width, -1)187 188        attn_output = self.proj(attn_output)189        return attn_output, attn_weights190 191 192class GotOcr2VisionLayer(GradientCheckpointingLayer):193    def __init__(self, config, window_size):194        super().__init__()195        self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)196        self.attn = GotOcr2VisionAttention(config, window_size)197        self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)198        self.mlp = GotOcr2MLPBlock(config)199        self.window_size = window_size200 201    def window_partition(self, hidden_states: torch.Tensor, window_size: int) -> tuple[torch.Tensor, tuple[int, int]]:202        """203        Args:204        Partition into non-overlapping windows with padding if needed.205            hidden_states (tensor): input tokens with [batch_size, height, width, channel]. window_size (int): window206            size.207 208        Returns:209            windows: windows after partition with [batch_size * num_windows, window_size, window_size, channel].210            (pad_height, pad_width): padded height and width before partition211        """212        batch_size, height, width, channel = hidden_states.shape213 214        pad_h = (window_size - height % window_size) % window_size215        pad_w = (window_size - width % window_size) % window_size216        hidden_states = F.pad(hidden_states, (0, 0, 0, pad_w, 0, pad_h))217        pad_height, pad_width = height + pad_h, width + pad_w218 219        hidden_states = hidden_states.reshape(220            batch_size, pad_height // window_size, window_size, pad_width // window_size, window_size, channel221        )222        windows = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(-1, window_size, window_size, channel)223        return windows, (pad_height, pad_width)224 225    def window_unpartition(226        self, windows: torch.Tensor, window_size: int, padding_shape: tuple[int, int], original_shape: tuple[int, int]227    ) -> torch.Tensor:228        """229        Args:230        Window unpartition into original sequences and removing padding.231            hidden_states (tensor):232                input tokens with [batch_size * num_windows, window_size, window_size, channel].233            window_size (int):234                window size.235            padding_shape (Tuple):236                padded height and width (pad_height, pad_width).237            original_shape (Tuple): original height and width (height, width) before padding.238 239        Returns:240            hidden_states: unpartitioned sequences with [batch_size, height, width, channel].241        """242        pad_height, pad_width = padding_shape243        height, width = original_shape244        batch_size = windows.shape[0] // (pad_height * pad_width // window_size // window_size)245        hidden_states = windows.reshape(246            batch_size, pad_height // window_size, pad_width // window_size, window_size, window_size, -1247        )248        hidden_states = (249            hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(batch_size, pad_height, pad_width, -1)250        )251 252        hidden_states = hidden_states[:, :height, :width, :].contiguous()253        return hidden_states254 255    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.FloatTensor]:256        residual = hidden_states257        hidden_states = self.layer_norm1(hidden_states)258        # Window partition259        if self.window_size > 0:260            height, width = hidden_states.shape[1], hidden_states.shape[2]261            hidden_states, padding_shape = self.window_partition(hidden_states, self.window_size)262 263        hidden_states, attn_weights = self.attn(264            hidden_states=hidden_states,265        )266        # Reverse window partition267        if self.window_size > 0:268            hidden_states = self.window_unpartition(hidden_states, self.window_size, padding_shape, (height, width))269 270        hidden_states = residual + hidden_states271        layernorm_output = self.layer_norm2(hidden_states)272        hidden_states = hidden_states + self.mlp(layernorm_output)273        return hidden_states274 275 276@auto_docstring277class GotOcr2PreTrainedModel(PreTrainedModel):278    config: GotOcr2Config279    base_model_prefix = ""280    supports_gradient_checkpointing = True281    _skip_keys_device_placement = "past_key_values"282    _supports_flash_attn = False283    _supports_sdpa = False284 285    _can_compile_fullgraph = True286    _supports_flex_attn = False287    _supports_attention_backend = True288 289    def _init_weights(self, module):290        super()._init_weights(module)291        if isinstance(module, GotOcr2VisionAttention):292            if module.use_rel_pos:293                module.rel_pos_h.data.zero_()294                module.rel_pos_w.data.zero_()295        elif isinstance(module, GotOcr2VisionEncoder):296            if module.pos_embed is not None:297                module.pos_embed.data.zero_()298 299 300@dataclass301@auto_docstring(302    custom_intro="""303    Base class for got_ocr2 vision model's outputs that also contains image embeddings obtained by applying the projection304    layer to the pooler_output.305    """306)307class GotOcr2VisionEncoderOutput(ModelOutput):308    r"""309    image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):310        The image embeddings obtained by applying the projection layer to the pooler_output.311    """312 313    image_embeds: Optional[torch.FloatTensor] = None314    last_hidden_state: Optional[torch.FloatTensor] = None315    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None316    attentions: Optional[tuple[torch.FloatTensor, ...]] = None317 318 319class GotOcr2PatchEmbeddings(nn.Module):320    """321    This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial322    `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a323    Transformer.324    """325 326    def __init__(self, config):327        super().__init__()328        image_size, patch_size = config.image_size, config.patch_size329        num_channels, hidden_size = config.num_channels, config.hidden_size330        image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)331        patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)332        num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])333        self.image_size = image_size334        self.patch_size = patch_size335        self.num_channels = num_channels336        self.num_patches = num_patches337 338        self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)339 340    def forward(self, pixel_values):341        batch_size, num_channels, height, width = pixel_values.shape342        if num_channels != self.num_channels:343            raise ValueError(344                "Make sure that the channel dimension of the pixel values match with the one set in the configuration."345            )346        if height != self.image_size[0] or width != self.image_size[1]:347            raise ValueError(348                f"Input image size ({height}*{width}) doesn't match model ({self.image_size[0]}*{self.image_size[1]})."349            )350        embeddings = self.projection(pixel_values).permute(0, 2, 3, 1)351        return embeddings352 353 354class GotOcr2LayerNorm(nn.LayerNorm):355    r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.356    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height,357    width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width).358    """359 360    def __init__(self, normalized_shape, *, eps=1e-6, data_format="channels_last", **kwargs):361        super().__init__(normalized_shape, eps=eps, **kwargs)362        if data_format not in ["channels_last", "channels_first"]:363            raise NotImplementedError(f"Unsupported data format: {data_format}")364        self.data_format = data_format365 366    def forward(self, features: torch.Tensor) -> torch.Tensor:367        """368        Args:369            features: Tensor of shape (batch_size, channels, height, width) OR (batch_size, height, width, channels)370        """371        if self.data_format == "channels_first":372            features = features.permute(0, 2, 3, 1)373            features = super().forward(features)374            features = features.permute(0, 3, 1, 2)375        else:376            features = super().forward(features)377        return features378 379 380class GotOcr2VisionNeck(nn.Module):381    def __init__(self, config: GotOcr2VisionConfig):382        super().__init__()383        self.config = config384 385        self.conv1 = nn.Conv2d(config.hidden_size, config.output_channels, kernel_size=1, bias=False)386        self.layer_norm1 = GotOcr2LayerNorm(config.output_channels, data_format="channels_first")387        self.conv2 = nn.Conv2d(config.output_channels, config.output_channels, kernel_size=3, padding=1, bias=False)388        self.layer_norm2 = GotOcr2LayerNorm(config.output_channels, data_format="channels_first")389 390    def forward(self, hidden_states):391        hidden_states = hidden_states.permute(0, 3, 1, 2)392        hidden_states = self.conv1(hidden_states)393        hidden_states = self.layer_norm1(hidden_states)394 395        hidden_states = self.conv2(hidden_states)396        hidden_states = self.layer_norm2(hidden_states)397        return hidden_states398 399 400class GotOcr2VisionEncoder(GotOcr2PreTrainedModel):401    _can_record_outputs = {"hidden_states": GotOcr2VisionLayer, "attentions": GotOcr2VisionAttention}402 403    def __init__(self, config: GotOcr2VisionConfig):404        super().__init__(config)405        self.config = config406        self.image_size = config.image_size407        self.patch_embed = GotOcr2PatchEmbeddings(config)408 409        self.pos_embed = None410        if config.use_abs_pos:411            # Initialize absolute positional embedding with pretrain image size.412            self.pos_embed = nn.Parameter(413                torch.zeros(414                    1,415                    config.image_size // config.patch_size,416                    config.image_size // config.patch_size,417                    config.hidden_size,418                )419            )420 421        self.layers = nn.ModuleList()422        for i in range(config.num_hidden_layers):423            layer = GotOcr2VisionLayer(424                config,425                window_size=config.window_size if i not in config.global_attn_indexes else 0,426            )427            self.layers.append(layer)428 429        self.neck = GotOcr2VisionNeck(config)430 431        self.gradient_checkpointing = False432 433    def get_input_embeddings(self):434        return self.patch_embed435 436    @check_model_inputs(tie_last_hidden_states=False)437    def forward(438        self, pixel_values: Optional[torch.FloatTensor] = None, **kwargs: Unpack[TransformersKwargs]439    ) -> GotOcr2VisionEncoderOutput:440        if pixel_values is None:441            raise ValueError("You have to specify pixel_values")442 443        hidden_states = self.patch_embed(pixel_values)444        if self.pos_embed is not None:445            hidden_states = hidden_states + self.pos_embed446        for layer_module in self.layers:447            hidden_states = layer_module(hidden_states)448        hidden_states = self.neck(hidden_states)449        return GotOcr2VisionEncoderOutput(450            last_hidden_state=hidden_states,451        )452 453 454class GotOcr2MultiModalProjector(nn.Module):455    def __init__(self, config: GotOcr2Config):456        super().__init__()457        vision_output_channels = config.vision_config.output_channels458        language_hidden_size = config.text_config.hidden_size459        self.conv_upsampler1 = nn.Conv2d(460            vision_output_channels, vision_output_channels * 2, kernel_size=3, stride=2, padding=1, bias=False461        )462        self.conv_upsampler2 = nn.Conv2d(463            vision_output_channels * 2, language_hidden_size, kernel_size=3, stride=2, padding=1, bias=False464        )465        self.multimodal_projector = nn.Linear(language_hidden_size, language_hidden_size)466 467    def forward(self, vision_embeddings: torch.Tensor) -> torch.Tensor:468        hidden_state = self.conv_upsampler1(vision_embeddings)469        hidden_state = self.conv_upsampler2(hidden_state)470        hidden_state = hidden_state.flatten(2).permute(0, 2, 1)471        hidden_state = self.multimodal_projector(hidden_state)472        return hidden_state473 474 475@dataclass476@auto_docstring(477    custom_intro="""478    Base class for GotOcr2 causal language model (or autoregressive) outputs.479    """480)481class GotOcr2CausalLMOutputWithPast(ModelOutput):482    r"""483    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):484        Language modeling loss (for next-token prediction).485    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):486        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).487    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):488        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).489 490        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see491        `past_key_values` input) to speed up sequential decoding.492    image_hidden_states (`torch.FloatTensor`, *optional*):493        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.494        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.495    """496 497    loss: Optional[torch.FloatTensor] = None498    logits: Optional[torch.FloatTensor] = None499    past_key_values: Optional[Cache] = None500    hidden_states: Optional[tuple[torch.FloatTensor]] = None501    attentions: Optional[tuple[torch.FloatTensor]] = None502    image_hidden_states: Optional[torch.FloatTensor] = None503 504 505@dataclass506@auto_docstring(507    custom_intro="""508    Base class for GotOcr2 outputs, with hidden states and attentions.509    """510)511class GotOcr2ModelOutputWithPast(BaseModelOutputWithPast):512    r"""513    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):514        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).515 516        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see517        `past_key_values` input) to speed up sequential decoding.518    image_hidden_states (`torch.FloatTensor`, *optional*):519        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.520        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.521    """522 523    image_hidden_states: Optional[torch.FloatTensor] = None524 525 526@auto_docstring(527    custom_intro="""528    The GotOcr2 model which consists of a vision backbone and a language model, without a language modeling head.529    """530)531class GotOcr2Model(GotOcr2PreTrainedModel):532    _checkpoint_conversion_mapping = {"language_model.model": "language_model"}533 534    def __init__(self, config: GotOcr2Config):535        super().__init__(config)536        self.vision_tower = GotOcr2VisionEncoder(config.vision_config)537 538        self.multi_modal_projector = GotOcr2MultiModalProjector(config)539        self.language_model = AutoModel.from_config(config.text_config)540        self.post_init()541 542    def get_input_embeddings(self):543        return self.language_model.get_input_embeddings()544 545    def set_input_embeddings(self, value):546        self.language_model.set_input_embeddings(value)547 548    def set_decoder(self, decoder):549        self.language_model = decoder550 551    def get_decoder(self):552        return self.language_model553 554    def get_image_features(555        self,556        pixel_values: torch.FloatTensor,557    ):558        """559        Obtains image last hidden states from the vision tower and apply multimodal projection.560 561        Args:562            pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`)563        Returns:564            image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).565        """566        image_outputs = self.vision_tower(pixel_values).last_hidden_state567        return self.multi_modal_projector(image_outputs)568 569    def get_placeholder_mask(570        self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor571    ):572        """573        Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is574        equal to the length of multimodal features. If the lengths are different, an error is raised.575        """576        if input_ids is None:577            special_image_mask = inputs_embeds == self.get_input_embeddings()(578                torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)579            )580            special_image_mask = special_image_mask.all(-1)581        else:582            special_image_mask = input_ids == self.config.image_token_id583 584        n_image_tokens = special_image_mask.sum()585        special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)586        n_image_features = image_features.shape[0] * image_features.shape[1]587        if inputs_embeds[special_image_mask].numel() != image_features.numel():588            raise ValueError(589                f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"590            )591        return special_image_mask592 593    @can_return_tuple594    @auto_docstring595    def forward(596        self,597        input_ids: Optional[torch.LongTensor] = None,598        pixel_values: Optional[torch.FloatTensor] = None,599        attention_mask: Optional[torch.Tensor] = None,600        position_ids: Optional[torch.LongTensor] = None,601        past_key_values: Optional[Cache] = None,602        inputs_embeds: Optional[torch.FloatTensor] = None,603        use_cache: Optional[bool] = None,604        output_attentions: Optional[bool] = None,605        output_hidden_states: Optional[bool] = None,606        return_dict: Optional[bool] = None,607        cache_position: Optional[torch.LongTensor] = None,608        **kwargs: Unpack[TransformersKwargs],609    ) -> Union[tuple, GotOcr2ModelOutputWithPast]:610        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions611        output_hidden_states = (612            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states613        )614        return_dict = return_dict if return_dict is not None else self.config.use_return_dict615 616        if (input_ids is None) ^ (inputs_embeds is not None):617            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")618 619        if inputs_embeds is None:620            inputs_embeds = self.get_input_embeddings()(input_ids)621 622        if pixel_values is not None:623            image_features = self.get_image_features(pixel_values=pixel_values.to(inputs_embeds.dtype))624            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)625            special_image_mask = self.get_placeholder_mask(626                input_ids, inputs_embeds=inputs_embeds, image_features=image_features627            )628            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)629 630        outputs = self.language_model(631            attention_mask=attention_mask,632            position_ids=position_ids,633            past_key_values=past_key_values,634            inputs_embeds=inputs_embeds,635            use_cache=use_cache,636            output_attentions=output_attentions,637            output_hidden_states=output_hidden_states,638            return_dict=True,639            cache_position=cache_position,640            **kwargs,641        )642 643        return GotOcr2ModelOutputWithPast(644            last_hidden_state=outputs.last_hidden_state,645            past_key_values=outputs.past_key_values,646            hidden_states=outputs.hidden_states,647            attentions=outputs.attentions,648            image_hidden_states=image_features if pixel_values is not None else None,649        )650 651 652@auto_docstring(653    custom_intro="""654    The GOT_OCR2 model which consists of a vision backbone and a language model.655    """656)657class GotOcr2ForConditionalGeneration(GotOcr2PreTrainedModel, GenerationMixin):658    _checkpoint_conversion_mapping = {659        "^language_model.model": "model.language_model",660        "^vision_tower": "model.vision_tower",661        "^multi_modal_projector": "model.multi_modal_projector",662        "^language_model.lm_head": "lm_head",663    }664    _tied_weights_keys = ["lm_head.weight"]665 666    def __init__(self, config: GotOcr2Config):667        super().__init__(config)668        self.model = GotOcr2Model(config)669        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)670        self.post_init()671 672    def get_input_embeddings(self):673        return self.model.get_input_embeddings()674 675    def set_input_embeddings(self, value):676        self.model.set_input_embeddings(value)677 678    def get_output_embeddings(self) -> nn.Module:679        return self.lm_head680 681    def set_decoder(self, decoder):682        self.model.set_decoder(decoder)683 684    def get_decoder(self):685        return self.model.get_decoder()686 687    def get_image_features(688        self,689        pixel_values: torch.FloatTensor,690        vision_feature_layer: Optional[Union[int, list[int]]] = None,691        vision_feature_select_strategy: Optional[str] = None,692        **kwargs,693    ):694        return self.model.get_image_features(695            pixel_values=pixel_values,696            vision_feature_layer=vision_feature_layer,697            vision_feature_select_strategy=vision_feature_select_strategy,698            **kwargs,699        )700 701    # Make modules available through conditional class for BC702    @property703    def language_model(self):704        return self.model.language_model705 706    @property707    def vision_tower(self):708        return self.model.vision_tower709 710    @property711    def multi_modal_projector(self):712        return self.model.multi_modal_projector713 714    @can_return_tuple715    @auto_docstring716    def forward(717        self,718        input_ids: Optional[torch.LongTensor] = None,719        pixel_values: Optional[torch.FloatTensor] = None,720        attention_mask: Optional[torch.Tensor] = None,721        position_ids: Optional[torch.LongTensor] = None,722        past_key_values: Optional[Cache] = None,723        inputs_embeds: Optional[torch.FloatTensor] = None,724        labels: Optional[torch.LongTensor] = None,725        use_cache: Optional[bool] = None,726        output_attentions: Optional[bool] = None,727        output_hidden_states: Optional[bool] = None,728        return_dict: Optional[bool] = None,729        cache_position: Optional[torch.LongTensor] = None,730        logits_to_keep: Union[int, torch.Tensor] = 0,731        **kwargs: Unpack[TransformersKwargs],732    ) -> Union[tuple, GotOcr2CausalLMOutputWithPast]:733        r"""734        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):735            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,736            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored737            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.738 739        Example:740 741        ```python742        >>> from PIL import Image743        >>> import requests744        >>> from transformers import AutoProcessor, GotOcr2ForConditionalGeneration, TextStreamer745 746        >>> model = GotOcr2ForConditionalGeneration.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf").to("cuda")747        >>> processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR-2.0-hf")748 749        >>> url = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png"750        >>> image = Image.open(requests.get(url, stream=True).raw)751 752        >>> inputs = processor(image, return_tensors="pt", color="green").to("cuda")753 754        >>> # Generate755        >>> streamer = TextStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)756        >>> generate_ids = model.generate(757        ...     **inputs,758        ...     do_sample=False,759        ...     tokenizer = processor.tokenizer,760        ...     stop_strings='<|im_end|>',761        ...     streamer=streamer,762        ...     max_new_tokens=4096,763        ... )764        "You should keep in mind what features from the module should be used, especially765        when you're planning to sell a template."766        ```"""767        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions768        output_hidden_states = (769            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states770        )771        return_dict = return_dict if return_dict is not None else self.config.use_return_dict772 773        outputs = self.model(774            input_ids=input_ids,775            pixel_values=pixel_values,776            attention_mask=attention_mask,777            position_ids=position_ids,778            past_key_values=past_key_values,779            inputs_embeds=inputs_embeds,780            use_cache=use_cache,781            output_attentions=output_attentions,782            output_hidden_states=output_hidden_states,783            return_dict=True,784            cache_position=cache_position,785            logits_to_keep=logits_to_keep,786            **kwargs,787        )788 789        hidden_states = outputs[0]790        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss791        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep792        logits = self.lm_head(hidden_states[:, slice_indices, :])793 794        loss = None795        if labels is not None:796            loss = self.loss_function(797                logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs798            )799 800        return GotOcr2CausalLMOutputWithPast(801            loss=loss,802            logits=logits,803            past_key_values=outputs.past_key_values,804            hidden_states=outputs.hidden_states,805            attentions=outputs.attentions,806            image_hidden_states=outputs.image_hidden_states,807        )808 809    def prepare_inputs_for_generation(810        self,811        input_ids,812        past_key_values=None,813        inputs_embeds=None,814        pixel_values=None,815        attention_mask=None,816        cache_position=None,817        logits_to_keep=None,818        **kwargs,819    ):820        # Overwritten -- in specific circumstances we don't want to forward image inputs to the model821 822        model_inputs = super().prepare_inputs_for_generation(823            input_ids,824            past_key_values=past_key_values,825            inputs_embeds=inputs_embeds,826            attention_mask=attention_mask,827            cache_position=cache_position,828            logits_to_keep=logits_to_keep,829            **kwargs,830        )831 832        if cache_position[0] == 0:833            # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore834            # Otherwise we need pixel values to be passed to model835            model_inputs["pixel_values"] = pixel_values836 837        return model_inputs838 839 840__all__ = ["GotOcr2PreTrainedModel", "GotOcr2Model", "GotOcr2ForConditionalGeneration"]841