# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py. # Do NOT edit this file manually as any edits will be overwritten by the generation of # the file from the modular. If any change should be done, please apply the change to the # modular_deepseek_ocr2.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # Copyright 2026 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import collections import math from collections.abc import Callable from dataclasses import dataclass from typing import Optional import torch import torch.nn.functional as F from torch import nn from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import ( use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func, ) from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check from ...utils.generic import maybe_autocast, merge_with_config_defaults from ...utils.output_capturing import capture_outputs from .configuration_deepseek_ocr2 import ( DeepseekOcr2Config, DeepseekOcr2SamVisionConfig, DeepseekOcr2TextConfig, DeepseekOcr2VisionConfig, ) logger = logging.get_logger(__name__) @dataclass class DeepseekOcr2ModelOutputWithPooling(BaseModelOutputWithPooling): """ local_last_hidden_state (`torch.FloatTensor` of shape `(total_local_patches, sequence_length, hidden_size)`, *optional*): Last hidden state from the vision encoder for local (cropped) patches. local_hidden_states (`torch.FloatTensor`, *optional*): Hidden states from all layers of the vision encoder for local patches. local_attentions (`torch.FloatTensor`, *optional*): Attention weights from all layers of the vision encoder for local patches. """ local_last_hidden_state: torch.FloatTensor | None = None local_hidden_states: torch.FloatTensor | None = None local_attentions: torch.FloatTensor | None = None @auto_docstring( custom_intro=""" Base class for Llava outputs, with hidden states and attentions. """ ) @dataclass class DeepseekOcr2ModelOutputWithPast(BaseModelOutputWithPast): r""" past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. image_hidden_states (`torch.FloatTensor`, *optional*): A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`. image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. """ image_hidden_states: torch.FloatTensor | None = None @auto_docstring( custom_intro=""" Base class for DeepseekOcr2 causal language model (or autoregressive) outputs. """ ) @dataclass class DeepseekOcr2CausalLMOutputWithPast(ModelOutput): r""" loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Language modeling loss (for next-token prediction). logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. image_hidden_states (`torch.FloatTensor`, *optional*): A `torch.FloatTensor` of size (batch_size * num_patches, num_images, sequence_length, hidden_size)`. image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state. """ loss: torch.FloatTensor | None = None logits: torch.FloatTensor | None = None past_key_values: Cache | None = None hidden_states: tuple[torch.FloatTensor] | None = None attentions: tuple[torch.FloatTensor] | None = None image_hidden_states: torch.FloatTensor | None = None @auto_docstring class DeepseekOcr2PreTrainedModel(PreTrainedModel): config: DeepseekOcr2Config base_model_prefix = "model" input_modalities = ("image", "text") supports_gradient_checkpointing = True _no_split_modules = [ "DeepseekOcr2SamVisionLayer", "DeepseekOcr2VisionEncoderLayer", "DeepseekOcr2TextDecoderLayer", ] _skip_keys_device_placement = ["past_key_values"] # SAM uses rel-pos bias, incompatible with flash attention. _supports_flash_attn = False _supports_sdpa = True _can_compile_fullgraph = True _supports_flex_attn = True _supports_attention_backend = True @torch.no_grad() def _init_weights(self, module): super()._init_weights(module) if isinstance(module, DeepseekOcr2SamVisionAttention): if module.use_rel_pos: init.zeros_(module.rel_pos_h) init.zeros_(module.rel_pos_w) elif isinstance(module, DeepseekOcr2SamVisionEncoder): if module.pos_embed is not None: init.zeros_(module.pos_embed) elif isinstance(module, DeepseekOcr2Model): embed_std = 1 / math.sqrt(self.config.text_config.hidden_size) init.normal_(module.view_separator, mean=0.0, std=embed_std) class DeepseekOcr2SamVisionAttention(nn.Module): """Multi-head Attention block with relative position embeddings.""" def __init__(self, config, window_size): super().__init__() input_size = ( (config.image_size // config.patch_size, config.image_size // config.patch_size) if window_size == 0 else (window_size, window_size) ) self.num_attention_heads = config.num_attention_heads head_dim = config.hidden_size // config.num_attention_heads self.scale = head_dim**-0.5 self.dropout = config.attention_dropout self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=config.qkv_bias) self.proj = nn.Linear(config.hidden_size, config.hidden_size) self.use_rel_pos = config.use_rel_pos if self.use_rel_pos: if input_size is None: raise ValueError("Input size must be provided if using relative positional encoding.") # initialize relative positional embeddings self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim)) self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim)) def get_rel_pos(self, q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor: """ Get relative positional embeddings according to the relative positions of query and key sizes. Args: q_size (int): size of the query. k_size (int): size of key k. rel_pos (`torch.Tensor`): relative position embeddings (L, channel). Returns: Extracted positional embeddings according to relative positions. """ max_rel_dist = int(2 * max(q_size, k_size) - 1) # Interpolate rel pos. rel_pos_resized = F.interpolate( rel_pos.reshape(1, rel_pos.shape[0], -1).transpose(1, 2), size=max_rel_dist, mode="linear", ) rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0) # Scale the coords with short length if shapes for q and k are different. q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0) k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0) relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0) return rel_pos_resized[relative_coords.long()] def get_decomposed_rel_pos( self, query: torch.Tensor, rel_pos_h: torch.Tensor, rel_pos_w: torch.Tensor, q_size: tuple[int, int], k_size: tuple[int, int], ) -> torch.Tensor: """ Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`. https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py Args: query (`torch.Tensor`): query q in the attention layer with shape (batch_size, query_height * query_width, channel). rel_pos_h (`torch.Tensor`): relative position embeddings (Lh, channel) for height axis. rel_pos_w (`torch.Tensor`): relative position embeddings (Lw, channel) for width axis. q_size (tuple): spatial sequence size of query q with (query_height, query_width). k_size (tuple): spatial sequence size of key k with (key_height, key_width). Returns: decomposed_rel_pos (`torch.Tensor`): decomposed relative position embeddings. """ query_height, query_width = q_size key_height, key_width = k_size relative_position_height = self.get_rel_pos(query_height, key_height, rel_pos_h) relative_position_width = self.get_rel_pos(query_width, key_width, rel_pos_w) batch_size, _, dim = query.shape reshaped_query = query.reshape(batch_size, query_height, query_width, dim) rel_h = torch.einsum("bhwc,hkc->bhwk", reshaped_query, relative_position_height) rel_w = torch.einsum("bhwc,wkc->bhwk", reshaped_query, relative_position_width) decomposed_rel_pos = rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :] return decomposed_rel_pos def forward(self, hidden_states: torch.Tensor, output_attentions=None) -> tuple[torch.Tensor, torch.Tensor]: batch_size, height, width, _ = hidden_states.shape # qkv with shape (3, batch_size, nHead, height * width, channel) qkv = ( self.qkv(hidden_states) .reshape(batch_size, height * width, 3, self.num_attention_heads, -1) .permute(2, 0, 3, 1, 4) ) # q, k, v with shape (batch_size * nHead, height * width, channel) query, key, value = qkv.reshape(3, batch_size * self.num_attention_heads, height * width, -1).unbind(0) attn_weights = (query * self.scale) @ key.transpose(-2, -1) if self.use_rel_pos: decomposed_rel_pos = self.get_decomposed_rel_pos( query, self.rel_pos_h, self.rel_pos_w, (height, width), (height, width) ) decomposed_rel_pos = decomposed_rel_pos.reshape_as(attn_weights) attn_weights = attn_weights + decomposed_rel_pos attn_weights = torch.nn.functional.softmax(attn_weights, dtype=torch.float32, dim=-1).to(query.dtype) attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) attn_output = (attn_probs @ value).reshape(batch_size, self.num_attention_heads, height, width, -1) attn_output = attn_output.permute(0, 2, 3, 1, 4).reshape(batch_size, height, width, -1) attn_output = self.proj(attn_output) return attn_output, attn_weights class DeepseekOcr2SamMLPBlock(nn.Module): def __init__(self, config): super().__init__() self.lin1 = nn.Linear(config.hidden_size, config.mlp_dim) self.lin2 = nn.Linear(config.mlp_dim, config.hidden_size) self.act = ACT2FN[config.hidden_act] def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.lin1(hidden_states) hidden_states = self.act(hidden_states) hidden_states = self.lin2(hidden_states) return hidden_states class DeepseekOcr2SamVisionSdpaAttention(DeepseekOcr2SamVisionAttention): """ Multi-head Attention block with relative position embeddings. Using SDPA instead of the default attention. """ def __init__(self, config, window_size): super().__init__(config, window_size) def forward(self, hidden_states: torch.Tensor, output_attentions=False) -> torch.Tensor: if output_attentions: logger.warning_once( f"{self.__class__.__name__} does not support `output_attentions=True`. The returned attention weights will " "be `None`. If you want to get attention weights, please set `attn_implementation='eager'` when loading the model." ) batch_size, height, width, _ = hidden_states.shape # qkv with shape (3, B, nHead, H * W, C) qkv = ( self.qkv(hidden_states) .reshape(batch_size, height * width, 3, self.num_attention_heads, -1) .permute(2, 0, 3, 1, 4) ) # q, k, v with shape (B * nHead, H * W, C) query, key, value = qkv.reshape(3, batch_size * self.num_attention_heads, height * width, -1).unbind(0) attn_bias = None if self.use_rel_pos: decomposed_rel_pos = self.get_decomposed_rel_pos( query, self.rel_pos_h, self.rel_pos_w, (height, width), (height, width) ) decomposed_rel_pos = decomposed_rel_pos.reshape( batch_size, self.num_attention_heads, height * width, height * width ) attn_bias = decomposed_rel_pos query = query.view(batch_size, self.num_attention_heads, height * width, -1) key = key.view(batch_size, self.num_attention_heads, height * width, -1) value = value.view(batch_size, self.num_attention_heads, height * width, -1) attn_output = torch.nn.functional.scaled_dot_product_attention(query, key, value, attn_mask=attn_bias) attn_output = ( attn_output.view(batch_size, self.num_attention_heads, height, width, -1) .permute(0, 2, 3, 1, 4) .reshape(batch_size, height, width, -1) ) attn_output = self.proj(attn_output) return attn_output, None DEEPSEEK_OCR2_SAM_VISION_ATTENTION_CLASSES = { "eager": DeepseekOcr2SamVisionAttention, "sdpa": DeepseekOcr2SamVisionSdpaAttention, } class DeepseekOcr2SamVisionLayer(GradientCheckpointingLayer): def __init__(self, config, window_size): super().__init__() self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.attn = DEEPSEEK_OCR2_SAM_VISION_ATTENTION_CLASSES[config._attn_implementation](config, window_size) self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.mlp = DeepseekOcr2SamMLPBlock(config) self.window_size = window_size def window_partition(self, hidden_states: torch.Tensor, window_size: int) -> tuple[torch.Tensor, tuple[int, int]]: """ Args: Partition into non-overlapping windows with padding if needed. hidden_states (tensor): input tokens with [batch_size, height, width, channel]. window_size (int): window size. Returns: windows: windows after partition with [batch_size * num_windows, window_size, window_size, channel]. (pad_height, pad_width): padded height and width before partition """ batch_size, height, width, channel = hidden_states.shape pad_h = (window_size - height % window_size) % window_size pad_w = (window_size - width % window_size) % window_size hidden_states = F.pad(hidden_states, (0, 0, 0, pad_w, 0, pad_h)) pad_height, pad_width = height + pad_h, width + pad_w hidden_states = hidden_states.reshape( batch_size, pad_height // window_size, window_size, pad_width // window_size, window_size, channel ) windows = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(-1, window_size, window_size, channel) return windows, (pad_height, pad_width) def window_unpartition( self, windows: torch.Tensor, window_size: int, padding_shape: tuple[int, int], original_shape: tuple[int, int] ) -> torch.Tensor: """ Args: Window unpartition into original sequences and removing padding. hidden_states (tensor): input tokens with [batch_size * num_windows, window_size, window_size, channel]. window_size (int): window size. padding_shape (Tuple): padded height and width (pad_height, pad_width). original_shape (Tuple): original height and width (height, width) before padding. Returns: hidden_states: unpartitioned sequences with [batch_size, height, width, channel]. """ pad_height, pad_width = padding_shape height, width = original_shape batch_size = windows.shape[0] // (pad_height * pad_width // window_size // window_size) hidden_states = windows.reshape( batch_size, pad_height // window_size, pad_width // window_size, window_size, window_size, -1 ) hidden_states = ( hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().reshape(batch_size, pad_height, pad_width, -1) ) hidden_states = hidden_states[:, :height, :width, :].contiguous() return hidden_states def forward(self, hidden_states: torch.Tensor) -> tuple[torch.FloatTensor]: residual = hidden_states hidden_states = self.layer_norm1(hidden_states) # Window partition if self.window_size > 0: height, width = hidden_states.shape[1], hidden_states.shape[2] hidden_states, padding_shape = self.window_partition(hidden_states, self.window_size) hidden_states, attn_weights = self.attn( hidden_states=hidden_states, ) # Reverse window partition if self.window_size > 0: hidden_states = self.window_unpartition(hidden_states, self.window_size, padding_shape, (height, width)) hidden_states = residual + hidden_states layernorm_output = self.layer_norm2(hidden_states) hidden_states = hidden_states + self.mlp(layernorm_output) return hidden_states class DeepseekOcr2SamLayerNorm(nn.LayerNorm): r"""LayerNorm that supports two data formats: channels_last (default) or channels_first. The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height, width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width). """ def __init__(self, normalized_shape, *, eps=1e-6, data_format="channels_last", **kwargs): super().__init__(normalized_shape, eps=eps, **kwargs) if data_format not in ["channels_last", "channels_first"]: raise NotImplementedError(f"Unsupported data format: {data_format}") self.data_format = data_format def forward(self, features: torch.Tensor) -> torch.Tensor: """ Args: features: Tensor of shape (batch_size, channels, height, width) OR (batch_size, height, width, channels) """ if self.data_format == "channels_first": features = features.permute(0, 2, 3, 1) features = super().forward(features) features = features.permute(0, 3, 1, 2) else: features = super().forward(features) return features class DeepseekOcr2SamVisionNeck(nn.Module): def __init__(self, config: DeepseekOcr2SamVisionConfig): super().__init__() self.config = config self.conv1 = nn.Conv2d(config.hidden_size, config.output_channels, kernel_size=1, bias=False) self.layer_norm1 = DeepseekOcr2SamLayerNorm(config.output_channels, data_format="channels_first") self.conv2 = nn.Conv2d(config.output_channels, config.output_channels, kernel_size=3, padding=1, bias=False) self.layer_norm2 = DeepseekOcr2SamLayerNorm(config.output_channels, data_format="channels_first") def forward(self, hidden_states): hidden_states = hidden_states.permute(0, 3, 1, 2) hidden_states = self.conv1(hidden_states) hidden_states = self.layer_norm1(hidden_states) hidden_states = self.conv2(hidden_states) hidden_states = self.layer_norm2(hidden_states) return hidden_states class DeepseekOcr2SamPatchEmbeddings(nn.Module): """ This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a Transformer. """ def __init__(self, config): super().__init__() image_size, patch_size = config.image_size, config.patch_size num_channels, hidden_size = config.num_channels, config.hidden_size image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size) patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size) num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0]) self.image_size = image_size self.patch_size = patch_size self.num_channels = num_channels self.num_patches = num_patches self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size) def forward(self, pixel_values): embeddings = self.projection(pixel_values).permute(0, 2, 3, 1) return embeddings class DeepseekOcr2SamVisionProj(nn.Module): """Neck and multi-scale downsampling for SAM ViT-B output.""" def __init__(self, config: DeepseekOcr2SamVisionConfig): super().__init__() self.conv1 = nn.Conv2d( config.output_channels, config.downsample_channels[0], kernel_size=3, stride=2, padding=1, bias=False, ) self.conv2 = nn.Conv2d( config.downsample_channels[0], config.downsample_channels[1], kernel_size=3, stride=2, padding=1, bias=False, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.conv1(hidden_states) hidden_states = self.conv2(hidden_states) return hidden_states class DeepseekOcr2SamVisionEncoder(DeepseekOcr2PreTrainedModel): _can_record_outputs = {"hidden_states": DeepseekOcr2SamVisionLayer, "attentions": DeepseekOcr2SamVisionAttention} def __init__(self, config: DeepseekOcr2SamVisionConfig): super().__init__(config) self.config = config self.image_size = config.image_size self.patch_embed = DeepseekOcr2SamPatchEmbeddings(config) self.pos_embed = None if config.use_abs_pos: # Initialize absolute positional embedding with pretrain image size. self.pos_embed = nn.Parameter( torch.zeros( 1, config.image_size // config.patch_size, config.image_size // config.patch_size, config.hidden_size, ) ) self.layers = nn.ModuleList() for i in range(config.num_hidden_layers): layer = DeepseekOcr2SamVisionLayer( config, window_size=config.window_size if i not in config.global_attn_indexes else 0, ) self.layers.append(layer) self.neck = DeepseekOcr2SamVisionNeck(config) self.gradient_checkpointing = False self.proj = DeepseekOcr2SamVisionProj(config) self.post_init() def get_input_embeddings(self): return self.patch_embed @merge_with_config_defaults @capture_outputs @auto_docstring def forward(self, pixel_values: torch.FloatTensor, **kwargs) -> BaseModelOutput: hidden_states = self.patch_embed(pixel_values) if self.pos_embed is not None: hidden_states = hidden_states + self.interpolate_pos_encoding( hidden_states.shape[1], hidden_states.shape[2] ) for layer_module in self.layers: hidden_states = layer_module(hidden_states) hidden_states = self.neck(hidden_states) hidden_states = self.proj(hidden_states) return BaseModelOutput(last_hidden_state=hidden_states) def interpolate_pos_encoding(self, height: int, width: int) -> torch.Tensor: """Interpolate the positional encoding to match the target spatial size using bicubic interpolation.""" if not torch.jit.is_tracing() and self.pos_embed.shape[1] == height and self.pos_embed.shape[2] == width: return self.pos_embed target_dtype = self.pos_embed.dtype pos_embed = self.pos_embed.permute(0, 3, 1, 2) pos_embed = torch.nn.functional.interpolate( pos_embed.to(torch.float32), size=(height, width), mode="bicubic", align_corners=False, antialias=True, ).to(dtype=target_dtype) pos_embed = pos_embed.permute(0, 2, 3, 1) return pos_embed class DeepseekOcr2VisionMLP(nn.Module): def __init__(self, config): super().__init__() self.config = config self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_act] def forward(self, x): down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) return down_proj @use_kernel_forward_from_hub("RMSNorm") class DeepseekOcr2VisionRMSNorm(nn.Module): def __init__(self, hidden_size, eps: float = 1e-6) -> None: """ DeepseekOcr2VisionRMSNorm is equivalent to T5LayerNorm """ super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states.to(input_dtype) def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" class DeepseekOcr2VisionRotaryEmbedding(nn.Module): inv_freq: torch.Tensor # fix linting for `register_buffer` def __init__(self, config: DeepseekOcr2VisionConfig, device=None): super().__init__() self.max_seq_len_cached = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings self.config = config self.rope_type = self.config.rope_parameters["rope_type"] rope_init_fn: Callable = self.compute_default_rope_parameters if self.rope_type != "default": rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) @staticmethod def compute_default_rope_parameters( config: DeepseekOcr2VisionConfig | None = None, device: Optional["torch.device"] = None, seq_len: int | None = None, ) -> tuple["torch.Tensor", float]: """ Computes the inverse frequencies according to the original RoPE implementation Args: config ([`~transformers.PreTrainedConfig`]): The model configuration. device (`torch.device`): The device to use for initialization of the inverse frequencies. seq_len (`int`, *optional*): The current sequence length. Unused for this type of RoPE. Returns: Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). """ base = config.rope_parameters["rope_theta"] dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads attention_factor = 1.0 # Unused in this type of RoPE # Compute the inverse frequencies inv_freq = 1.0 / ( base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) ) return inv_freq, attention_factor @torch.no_grad() @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) def forward(self, x, position_ids): inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) position_ids_expanded = position_ids[:, None, :].float() device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" with maybe_autocast(device_type=device_type, enabled=False): # Force float32 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() * self.attention_scaling sin = emb.sin() * self.attention_scaling return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @use_kernel_func_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """Applies Rotary Position Embedding to the query and key tensors. Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): The key tensor. cos (`torch.Tensor`): The cosine part of the rotary embedding. sin (`torch.Tensor`): The sine part of the rotary embedding. unsqueeze_dim (`int`, *optional*, defaults to 1): The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. Returns: `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) """ batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) def eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor | None, scaling: float, dropout: float = 0.0, **kwargs: Unpack[TransformersKwargs], ): key_states = repeat_kv(key, module.num_key_value_groups) value_states = repeat_kv(value, module.num_key_value_groups) attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling if attention_mask is not None: attn_weights = attn_weights + attention_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights @use_kernelized_func(apply_rotary_pos_emb) class DeepseekOcr2VisionAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" def __init__(self, config: DeepseekOcr2VisionConfig, layer_idx: int): super().__init__() self.layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None self.config = config self.layer_idx = layer_idx self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads self.scaling = self.head_dim**-0.5 self.attention_dropout = config.attention_dropout self.is_causal = True self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=True) self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True) self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True) self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) self.sliding_window = config.sliding_window if self.layer_type == "sliding_attention" else None def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: torch.Tensor | None, past_key_values: Cache | None = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> tuple[torch.Tensor, torch.Tensor | None]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_values is not None: key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, sliding_window=self.sliding_window, # main diff with Llama **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights class DeepseekOcr2VisionEncoderLayer(GradientCheckpointingLayer): def __init__(self, config: DeepseekOcr2VisionConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.self_attn = DeepseekOcr2VisionAttention(config=config, layer_idx=layer_idx) self.mlp = DeepseekOcr2VisionMLP(config) self.input_layernorm = DeepseekOcr2VisionRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = DeepseekOcr2VisionRMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, use_cache: bool | None = False, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, position_embeddings=position_embeddings, **kwargs, ) hidden_states = residual + hidden_states # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states @auto_docstring(custom_intro="Vision encoder for DeepSeek-OCR-2.") class DeepseekOcr2VisionEncoder(DeepseekOcr2PreTrainedModel): _can_record_outputs = { "hidden_states": DeepseekOcr2VisionEncoderLayer, "attentions": DeepseekOcr2VisionAttention, } def __init__(self, config): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.layers = nn.ModuleList( [DeepseekOcr2VisionEncoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = DeepseekOcr2VisionRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = DeepseekOcr2VisionRotaryEmbedding(config=config) self.gradient_checkpointing = False self.has_sliding_layers = "sliding_attention" in self.config.layer_types # Initialize weights and apply final processing self.post_init() @merge_with_config_defaults @capture_outputs @auto_docstring def forward( self, inputs_embeds: torch.FloatTensor, num_patches: int, position_ids: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPast: r""" num_patches (`int`): Number of image patch tokens at the beginning of the sequence. Used to build the hybrid attention mask (bidirectional over image tokens, causal over query tokens). """ if position_ids is None: position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) bsz, seq_len, _ = inputs_embeds.shape block_sequence_ids = torch.full((bsz, seq_len), -1, dtype=torch.long, device=inputs_embeds.device) block_sequence_ids[:, :num_patches] = 0 attention_mask = create_causal_mask( config=self.config, inputs_embeds=inputs_embeds, attention_mask=None, past_key_values=None, block_sequence_ids=block_sequence_ids, ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids) for encoder_layer in self.layers[: self.config.num_hidden_layers]: hidden_states = encoder_layer( hidden_states, attention_mask=attention_mask, position_embeddings=position_embeddings, position_ids=position_ids, **kwargs, ) hidden_states = self.norm(hidden_states) return BaseModelOutputWithPast(last_hidden_state=hidden_states) class DeepseekOcr2VisionModel(DeepseekOcr2PreTrainedModel): """Vision pipeline: SAM ViT-B (with neck)""" def __init__(self, config: DeepseekOcr2VisionConfig): super().__init__(config) self.sam_encoder = DeepseekOcr2SamVisionEncoder(config.sam_config) self.vision_encoder = DeepseekOcr2VisionEncoder(config.encoder_config) # Resolution-specific learnable queries self.query_768_resolution = nn.Embedding(144, config.encoder_config.hidden_size) # 12x12 for 768px self.query_1024_resolution = nn.Embedding(256, config.encoder_config.hidden_size) # 16x16 for 1024px self.post_init() @can_return_tuple @auto_docstring def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: sam_encoder_outputs = self.sam_encoder(pixel_values, **kwargs) hidden_states = sam_encoder_outputs.last_hidden_state.flatten(2).transpose(1, 2) bsz, num_patches, _ = hidden_states.shape queries = self.query_768_resolution.weight if num_patches <= 144 else self.query_1024_resolution.weight queries = queries.unsqueeze(0).expand(bsz, -1, -1) combined = torch.cat([hidden_states, queries], dim=1) encoder_outputs = self.vision_encoder( inputs_embeds=combined, num_patches=num_patches, **kwargs, ) query_features = encoder_outputs.last_hidden_state[:, num_patches:, :] return BaseModelOutput( last_hidden_state=query_features, hidden_states=encoder_outputs.hidden_states, attentions=encoder_outputs.attentions, ) class DeepseekOcr2TextRotaryEmbedding(nn.Module): inv_freq: torch.Tensor # fix linting for `register_buffer` def __init__(self, config: DeepseekOcr2TextConfig, device=None): super().__init__() self.max_seq_len_cached = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings self.config = config self.rope_type = self.config.rope_parameters["rope_type"] rope_init_fn: Callable = self.compute_default_rope_parameters if self.rope_type != "default": rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) @staticmethod def compute_default_rope_parameters( config: DeepseekOcr2TextConfig | None = None, device: Optional["torch.device"] = None, seq_len: int | None = None, ) -> tuple["torch.Tensor", float]: """ Computes the inverse frequencies according to the original RoPE implementation Args: config ([`~transformers.PreTrainedConfig`]): The model configuration. device (`torch.device`): The device to use for initialization of the inverse frequencies. seq_len (`int`, *optional*): The current sequence length. Unused for this type of RoPE. Returns: Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). """ base = config.rope_parameters["rope_theta"] dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads attention_factor = 1.0 # Unused in this type of RoPE # Compute the inverse frequencies inv_freq = 1.0 / ( base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) ) return inv_freq, attention_factor @torch.no_grad() @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) def forward(self, x, position_ids): inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) position_ids_expanded = position_ids[:, None, :].float() device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" with maybe_autocast(device_type=device_type, enabled=False): # Force float32 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() * self.attention_scaling sin = emb.sin() * self.attention_scaling return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) @use_kernelized_func(apply_rotary_pos_emb) class DeepseekOcr2TextAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" def __init__(self, config: DeepseekOcr2TextConfig, layer_idx: int): super().__init__() self.config = config self.layer_idx = layer_idx self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads self.scaling = self.head_dim**-0.5 self.attention_dropout = config.attention_dropout self.is_causal = True self.q_proj = nn.Linear( config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias ) self.k_proj = nn.Linear( config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias ) self.v_proj = nn.Linear( config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias ) self.o_proj = nn.Linear( config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias ) def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, attention_mask: torch.Tensor | None = None, past_key_values: Cache | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_values is not None: key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights class DeepseekOcr2TextMLP(nn.Module): def __init__(self, config: DeepseekOcr2TextConfig, hidden_size=None, intermediate_size=None): super().__init__() self.config = config self.hidden_size = config.hidden_size if hidden_size is None else hidden_size self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) self.act_fn = ACT2FN[config.hidden_act] def forward(self, x): down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) return down_proj @use_experts_implementation class DeepseekOcr2TextExperts(nn.Module): """Collection of expert weights stored as 3D tensors.""" def __init__(self, config): super().__init__() self.num_experts = config.n_routed_experts self.hidden_dim = config.hidden_size self.intermediate_dim = config.moe_intermediate_size self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)) self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)) self.act_fn = ACT2FN[config.hidden_act] def forward( self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor, ) -> torch.Tensor: final_hidden_states = torch.zeros_like(hidden_states) with torch.no_grad(): expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts) expert_mask = expert_mask.permute(2, 1, 0) expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() for expert_idx in expert_hit: expert_idx = expert_idx[0] if expert_idx == self.num_experts: continue top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) current_state = hidden_states[token_idx] gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1) current_hidden_states = self.act_fn(gate) * up current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx]) current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None] final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype)) return final_hidden_states class DeepseekOcr2TextMoe(nn.Module): def __init__(self, config: DeepseekOcr2TextConfig): super().__init__() self.config = config self.experts = DeepseekOcr2TextExperts(config) self.gate = nn.Linear(config.hidden_size, config.n_routed_experts, bias=False) if config.n_shared_experts is not None: intermediate_size = config.moe_intermediate_size * config.n_shared_experts self.shared_experts = DeepseekOcr2TextMLP(config=config, intermediate_size=intermediate_size) self.routed_scaling_factor = config.routed_scaling_factor self.topk_method = config.topk_method self.num_group = config.n_group self.top_k = config.num_experts_per_tok self.topk_group = config.topk_group def route_tokens_to_experts(self, router_logits): batch_size, seq_len, hidden_dim = router_logits.shape router_logits = router_logits.view(-1, hidden_dim) router_logits = router_logits.softmax(dim=-1, dtype=torch.float32) if self.topk_method == "greedy": topk_weight, topk_idx = torch.topk(router_logits, k=self.top_k, dim=-1, sorted=False) elif self.topk_method == "group_limited_greedy": group_scores = router_logits.view(batch_size * seq_len, self.num_group, -1).max(dim=-1).values group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] group_mask = torch.zeros_like(group_scores) group_mask.scatter_(1, group_idx, 1) score_mask = ( group_mask.unsqueeze(-1) .expand(batch_size * seq_len, self.num_group, self.num_experts // self.num_group) .reshape(batch_size * seq_len, -1) ) tmp_scores = router_logits.masked_fill(~score_mask.bool(), 0.0) topk_weight, topk_idx = torch.topk(tmp_scores, k=self.top_k, dim=-1, sorted=False) topk_weight = topk_weight * self.routed_scaling_factor return topk_idx, topk_weight def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: residuals = hidden_states orig_shape = hidden_states.shape router_logits = nn.functional.linear(hidden_states.type(torch.float32), self.gate.weight.type(torch.float32)) topk_indices, topk_weights = self.route_tokens_to_experts(router_logits) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states = hidden_states + self.shared_experts(residuals) return hidden_states @use_kernel_forward_from_hub("RMSNorm") class DeepseekOcr2TextRMSNorm(nn.Module): def __init__(self, hidden_size, eps: float = 1e-6) -> None: """ DeepseekOcr2TextRMSNorm is equivalent to T5LayerNorm """ super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states.to(input_dtype) def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" class DeepseekOcr2TextDecoderLayer(GradientCheckpointingLayer): def __init__(self, config, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.self_attn = DeepseekOcr2TextAttention(config=config, layer_idx=layer_idx) self.mlp = ( DeepseekOcr2TextMoe(config) if config.mlp_layer_types[layer_idx] == "sparse" else DeepseekOcr2TextMLP(config) ) self.input_layernorm = DeepseekOcr2TextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = DeepseekOcr2TextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, use_cache: bool | None = False, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, position_embeddings=position_embeddings, **kwargs, ) hidden_states = residual + hidden_states # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states @auto_docstring class DeepseekOcr2TextPreTrainedModel(PreTrainedModel): config: DeepseekOcr2TextConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["DeepseekOcr2TextDecoderLayer"] _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn = True _supports_sdpa = True _supports_flex_attn = True _can_compile_fullgraph = True _supports_attention_backend = True _can_record_outputs = { "hidden_states": DeepseekOcr2TextDecoderLayer, "attentions": DeepseekOcr2TextAttention, } @torch.no_grad() def _init_weights(self, module): super()._init_weights(module) if isinstance(module, DeepseekOcr2TextExperts): init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range) init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range) @auto_docstring class DeepseekOcr2TextModel(DeepseekOcr2TextPreTrainedModel): def __init__(self, config: DeepseekOcr2TextConfig): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList( [DeepseekOcr2TextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = DeepseekOcr2TextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) # Use (cos/sin) RoPE instead of complex RoPE to match LlamaAttention (MHA) self.rotary_emb = DeepseekOcr2TextRotaryEmbedding(config=config) self.gradient_checkpointing = False # Initialize weights and apply final processing self.post_init() @merge_with_config_defaults @capture_outputs @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, use_cache: bool | None = None, **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) if position_ids is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens position_ids = position_ids.unsqueeze(0) causal_mask = create_causal_mask( config=self.config, inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, position_ids=position_ids, ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) for decoder_layer in self.layers[: self.config.num_hidden_layers]: hidden_states = decoder_layer( hidden_states, attention_mask=causal_mask, position_embeddings=position_embeddings, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, **kwargs, ) hidden_states = self.norm(hidden_states) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values, ) @auto_docstring( custom_intro=""" The Llava-Next model which consists of a vision backbone and a language model without language modeling head. """ ) class DeepseekOcr2Model(DeepseekOcr2PreTrainedModel): base_model_prefix = "model" def __init__(self, config: DeepseekOcr2Config): super().__init__(config) self.vision_tower = DeepseekOcr2VisionModel(config.vision_config) self.multi_modal_projector = nn.Linear( config.vision_config.encoder_config.hidden_size, config.text_config.hidden_size ) self.vocab_size = config.text_config.vocab_size self.language_model = DeepseekOcr2TextModel(config.text_config) # Learnable separator between local and global views (initialized in `_init_weights`). self.view_separator = nn.Parameter(torch.empty(config.text_config.hidden_size)) self.post_init() @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPooling: r""" pixel_values_local (`torch.FloatTensor` of shape `(total_patches, 3, height, width)`, *optional*): All local patches flattened across the batch, or `None` if no local views. num_local_patches (`list[int]` or `torch.Tensor`, *optional*): Number of local patches per image, e.g. `[6, 0, 4]`. """ # torch.split requires list[int], not Tensor, for per-image variable-length splitting if isinstance(num_local_patches, torch.Tensor): num_local_patches = num_local_patches.tolist() batch_size = pixel_values.shape[0] global_vision_outputs = self.vision_tower(pixel_values, **kwargs) global_features = self.multi_modal_projector(global_vision_outputs.last_hidden_state) local_outputs = {} if pixel_values_local is not None: local_vision_outputs = self.vision_tower(pixel_values_local, **kwargs) all_local_features = self.multi_modal_projector(local_vision_outputs.last_hidden_state) per_image_local = torch.split(all_local_features, num_local_patches, dim=0) local_outputs = { "local_last_hidden_state": local_vision_outputs.last_hidden_state, "local_hidden_states": local_vision_outputs.hidden_states, "local_attentions": local_vision_outputs.attentions, } else: per_image_local = [None] * batch_size all_features = [] view_sep = self.view_separator.to(global_features.device).unsqueeze(0) for idx in range(batch_size): global_flat = global_features[idx].reshape(-1, global_features.shape[-1]) if per_image_local[idx] is not None: local_flat = per_image_local[idx].reshape(-1, per_image_local[idx].shape[-1]) all_features.append(torch.cat([local_flat, global_flat, view_sep], dim=0)) else: all_features.append(torch.cat([global_flat, view_sep], dim=0)) image_features = torch.cat(all_features, dim=0) return DeepseekOcr2ModelOutputWithPooling( last_hidden_state=global_vision_outputs.last_hidden_state, pooler_output=image_features, hidden_states=global_vision_outputs.hidden_states, attentions=global_vision_outputs.attentions, **local_outputs, ) def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor ): """ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is equal to the length of multimodal features. If the lengths are different, an error is raised. """ if input_ids is None: special_image_mask = inputs_embeds == self.get_input_embeddings()( torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) ) special_image_mask = special_image_mask.all(-1) else: special_image_mask = input_ids == self.config.image_token_id n_image_tokens = special_image_mask.sum() special_image_mask = special_image_mask.unsqueeze(-1).to(inputs_embeds.device) torch_compilable_check( n_image_tokens * inputs_embeds.shape[-1] == image_features.numel(), f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {image_features.shape[0]}", ) return special_image_mask @can_return_tuple @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, use_cache: bool | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | DeepseekOcr2ModelOutputWithPast: r""" pixel_values_local (`torch.FloatTensor`, *optional*): Local patch pixel values of shape `(total_patches, 3, H, W)`. num_local_patches (`list[int]` or `torch.Tensor`, *optional*): Number of local patches per image in the batch. """ if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) image_features = None if pixel_values is not None: image_features = self.get_image_features( pixel_values, pixel_values_local, num_local_patches, return_dict=True ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask(input_ids, inputs_embeds, image_features) inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) outputs = self.language_model( input_ids=None, inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, **kwargs, ) return DeepseekOcr2ModelOutputWithPast( last_hidden_state=outputs.last_hidden_state, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, image_hidden_states=image_features, ) @auto_docstring class DeepseekOcr2ForConditionalGeneration(DeepseekOcr2PreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} def __init__(self, config: DeepseekOcr2Config): super().__init__(config) self.model = DeepseekOcr2Model(config) self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) self.post_init() def get_output_embeddings(self) -> nn.Module: return self.lm_head def pack_image_features(self): raise NotImplementedError("DeepseekOcr2 does not use pack_image_features") @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: r""" pixel_values (`torch.FloatTensor` of shape `(batch_size, 3, height, width)`): The tensors corresponding to the global view input images. pixel_values_local (`torch.FloatTensor` of shape `(total_patches, 3, height, width)`, *optional*): All local patches flattened across the batch, or `None` if no local views. num_local_patches (`list[int]` or `torch.Tensor`, *optional*): Number of local patches per image, e.g. `[6, 0, 4]`. """ return self.model.get_image_features( pixel_values=pixel_values, pixel_values_local=pixel_values_local, num_local_patches=num_local_patches, **kwargs, ) @can_return_tuple @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], ) -> tuple | DeepseekOcr2CausalLMOutputWithPast: r""" pixel_values_local (`torch.FloatTensor`, *optional*): Local patch pixel values of shape `(total_patches, 3, H, W)`. num_local_patches (`list[int]` or `torch.Tensor`, *optional*): Number of local patches per image in the batch. """ outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, pixel_values_local=pixel_values_local, num_local_patches=num_local_patches, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, **kwargs, ) hidden_states = outputs[0] slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep hidden_states = hidden_states[:, slice_indices, :] logits = self.lm_head(hidden_states) loss = None if labels is not None: loss = self.loss_function( logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs, ) return DeepseekOcr2CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, image_hidden_states=outputs.image_hidden_states, ) def prepare_inputs_for_generation( self, input_ids, past_key_values=None, inputs_embeds=None, pixel_values=None, pixel_values_local=None, num_local_patches=None, attention_mask=None, logits_to_keep=None, is_first_iteration=False, **kwargs, ): model_inputs = super().prepare_inputs_for_generation( input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, attention_mask=attention_mask, logits_to_keep=logits_to_keep, is_first_iteration=is_first_iteration, **kwargs, ) if is_first_iteration or not kwargs.get("use_cache", True): model_inputs["pixel_values"] = pixel_values model_inputs["pixel_values_local"] = pixel_values_local model_inputs["num_local_patches"] = num_local_patches return model_inputs __all__ = [ "DeepseekOcr2ForConditionalGeneration", "DeepseekOcr2Model", "DeepseekOcr2PreTrainedModel", "DeepseekOcr2TextModel", "DeepseekOcr2TextPreTrainedModel", "DeepseekOcr2VisionModel", ]