From 13b058583047352feaec6f4693004b13167b0aff Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 09:55:12 +0200 Subject: [PATCH 001/195] Add Unlimited OCR --- docs/source/en/_toctree.yml | 2 + docs/source/en/model_doc/unlimited_ocr.md | 91 + src/transformers/models/__init__.py | 1 + src/transformers/models/auto/auto_mappings.py | 9 + src/transformers/models/auto/modeling_auto.py | 2 + .../models/unlimited_ocr/__init__.py | 28 + .../configuration_unlimited_ocr.py | 306 +++ .../unlimited_ocr/modeling_unlimited_ocr.py | 1707 +++++++++++++++++ .../unlimited_ocr/modular_unlimited_ocr.py | 223 +++ tests/models/unlimited_ocr/__init__.py | 0 .../test_modeling_unlimited_ocr.py | 247 +++ 11 files changed, 2616 insertions(+) create mode 100644 docs/source/en/model_doc/unlimited_ocr.md create mode 100644 src/transformers/models/unlimited_ocr/__init__.py create mode 100644 src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py create mode 100644 src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py create mode 100644 src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py create mode 100644 tests/models/unlimited_ocr/__init__.py create mode 100644 tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 7130a9cfff9f..550ea5d35bb8 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -879,6 +879,8 @@ title: UL2 - local: model_doc/umt5 title: UMT5 + - local: model_doc/unlimited_ocr + title: UnlimitedOcr - local: model_doc/vaultgemma title: VaultGemma - local: model_doc/xmod diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md new file mode 100644 index 000000000000..6d762b4a9b8e --- /dev/null +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -0,0 +1,91 @@ + +*This model was contributed to Hugging Face Transformers on 2026-06-23.* + + +# UnlimitedOcr + +## Overview + +The UnlimitedOcr model was proposed in []() by . + + +The abstract from the paper is the following: + + + +Tips: + + + +This model was contributed by [INSERT YOUR HF USERNAME HERE](https://huggingface.co/). +The original code can be found [here](). + +## Usage examples + + + +## UnlimitedOcrConfig + +[[autodoc]] UnlimitedOcrConfig + +## UnlimitedOcrTextConfig + +[[autodoc]] UnlimitedOcrTextConfig + +## UnlimitedOcrVisionConfig + +[[autodoc]] UnlimitedOcrVisionConfig + +## UnlimitedOcrVisionEncoderConfig + +[[autodoc]] UnlimitedOcrVisionEncoderConfig + +## UnlimitedOcrSamVisionConfig + +[[autodoc]] UnlimitedOcrSamVisionConfig + +## UnlimitedOcrForConditionalGeneration + +[[autodoc]] UnlimitedOcrForConditionalGeneration + +## UnlimitedOcrModel + +[[autodoc]] UnlimitedOcrModel + - forward + +## UnlimitedOcrPreTrainedModel + +[[autodoc]] UnlimitedOcrPreTrainedModel + - forward + +## UnlimitedOcrTextModel + +[[autodoc]] UnlimitedOcrTextModel + - forward + +## UnlimitedOcrTextPreTrainedModel + +[[autodoc]] UnlimitedOcrTextPreTrainedModel + - forward + +## UnlimitedOcrVisionModel + +[[autodoc]] UnlimitedOcrVisionModel + - forward \ No newline at end of file diff --git a/src/transformers/models/__init__.py b/src/transformers/models/__init__.py index cb25ec5c29d5..8d6f7a1e64f1 100644 --- a/src/transformers/models/__init__.py +++ b/src/transformers/models/__init__.py @@ -448,6 +448,7 @@ from .unispeech import * from .unispeech_sat import * from .univnet import * + from .unlimited_ocr import * from .upernet import * from .uvdoc import * from .vaultgemma import * diff --git a/src/transformers/models/auto/auto_mappings.py b/src/transformers/models/auto/auto_mappings.py index 13382ceefc1e..3154e30c53e0 100644 --- a/src/transformers/models/auto/auto_mappings.py +++ b/src/transformers/models/auto/auto_mappings.py @@ -622,6 +622,11 @@ ("unispeech", "UniSpeechConfig"), ("unispeech-sat", "UniSpeechSatConfig"), ("univnet", "UnivNetConfig"), + ("unlimited_ocr", "UnlimitedOcrConfig"), + ("unlimited_ocr_encoder", "UnlimitedOcrVisionEncoderConfig"), + ("unlimited_ocr_sam_vision_model", "UnlimitedOcrSamVisionConfig"), + ("unlimited_ocr_text", "UnlimitedOcrTextConfig"), + ("unlimited_ocr_vision", "UnlimitedOcrVisionConfig"), ("upernet", "UperNetConfig"), ("uvdoc", "UVDocConfig"), ("uvdoc_backbone", "UVDocBackboneConfig"), @@ -891,6 +896,10 @@ ("t5gemma2_text", "t5gemma2"), ("table-transformer", "table_transformer"), ("unispeech-sat", "unispeech_sat"), + ("unlimited_ocr_encoder", "unlimited_ocr"), + ("unlimited_ocr_sam_vision_model", "unlimited_ocr"), + ("unlimited_ocr_text", "unlimited_ocr"), + ("unlimited_ocr_vision", "unlimited_ocr"), ("uvdoc_backbone", "uvdoc"), ("video_llama_3_vision", "video_llama_3"), ("videoprism_text_model", "videoprism"), diff --git a/src/transformers/models/auto/modeling_auto.py b/src/transformers/models/auto/modeling_auto.py index a9187f0b8984..428de798488a 100644 --- a/src/transformers/models/auto/modeling_auto.py +++ b/src/transformers/models/auto/modeling_auto.py @@ -494,6 +494,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("unispeech", "UniSpeechModel"), ("unispeech-sat", "UniSpeechSatModel"), ("univnet", "UnivNetModel"), + ("unlimited_ocr", "UnlimitedOcrModel"), ("uvdoc", "UVDocModel"), ("vaultgemma", "VaultGemmaModel"), ("vibevoice_acoustic_tokenizer", "VibeVoiceAcousticTokenizerModel"), @@ -1093,6 +1094,7 @@ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin): ("smolvlm", "SmolVLMForConditionalGeneration"), ("t5gemma2", "T5Gemma2ForConditionalGeneration"), ("udop", "UdopForConditionalGeneration"), + ("unlimited_ocr", "UnlimitedOcrForConditionalGeneration"), ("video_llama_3", "VideoLlama3ForConditionalGeneration"), ("video_llava", "VideoLlavaForConditionalGeneration"), ("vipllava", "VipLlavaForConditionalGeneration"), diff --git a/src/transformers/models/unlimited_ocr/__init__.py b/src/transformers/models/unlimited_ocr/__init__.py new file mode 100644 index 000000000000..6a34d47b1edc --- /dev/null +++ b/src/transformers/models/unlimited_ocr/__init__.py @@ -0,0 +1,28 @@ +# 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. + +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_unlimited_ocr import * + from .modeling_unlimited_ocr import * +else: + import sys + + _file = globals()["__file__"] + sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py new file mode 100644 index 000000000000..54d2171bc010 --- /dev/null +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -0,0 +1,306 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/unlimited_ocr/modular_unlimited_ocr.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_unlimited_ocr.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. + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="deepseek-community/DeepSeek-OCR-2") +@strict +class UnlimitedOcrSamVisionConfig(PreTrainedConfig): + r""" + output_channels (`int`, *optional*, defaults to 256): + The number of output channels in the SAM neck. + window_size (`int`, *optional*, defaults to 14): + Window size for windowed attention layers. + global_attn_indexes (`list[int]`, *optional*, defaults to `[2, 5, 8, 11]`): + Indices of encoder layers that use global (non-windowed) attention. + mlp_dim (`int`, *optional*): + Dimensionality of the MLP layer in each vision encoder block. Defaults to `hidden_size * mlp_ratio`. + downsample_channels (`list[int]`, *optional*): + The channel dimensions for the multi-scale downsampling neck layers. Defaults to `[512, 896]`. + """ + + base_config_key = "sam_config" + model_type = "unlimited_ocr_sam_vision_model" + + hidden_size: int = 768 + output_channels: int = 256 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + num_channels: int = 3 + image_size: int | list[int] | tuple[int, int] = 1024 + patch_size: int | list[int] | tuple[int, int] = 16 + hidden_act: str = "gelu" + layer_norm_eps: float = 1e-06 + attention_dropout: float | int = 0.0 + initializer_range: float = 1e-10 + qkv_bias: bool = True + mlp_ratio: float = 4.0 + use_abs_pos: bool = True + use_rel_pos: bool = True + window_size: int = 14 + global_attn_indexes: list[int] | tuple[int, ...] = (2, 5, 8, 11) + mlp_dim: int | None = None + + downsample_channels: list[int] | None = None + + def __post_init__(self, **kwargs): + if self.downsample_channels is None: + self.downsample_channels = [512, 896] + self.mlp_dim = int(self.hidden_size * self.mlp_ratio) if self.mlp_dim is None else self.mlp_dim + self.scale = self.hidden_size // 2 + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="deepseek-community/DeepSeek-OCR-2") +@strict +class UnlimitedOcrVisionEncoderConfig(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import UnlimitedOcrConfig + + >>> config = UnlimitedOcrConfig() + >>> encoder_config = config.vision_config.encoder_config + ```""" + + model_type = "unlimited_ocr_encoder" + keys_to_ignore_at_inference = ["past_key_values"] + + # Default tensor parallel plan for base model `UnlimitedOcrVisionEncoder` + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + vocab_size: int = 151936 + hidden_size: int = 4096 + intermediate_size: int = 22016 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + num_key_value_heads: int | None = 32 + hidden_act: str = "silu" + max_position_embeddings: int = 32768 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-6 + use_cache: bool = True + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + use_sliding_window: bool = False + sliding_window: int | None = 4096 + max_window_layers: int = 28 + layer_types: list[str] | None = None + attention_dropout: float | int = 0.0 + pad_token_id: int | None = None + bos_token_id: int | None = None + eos_token_id: int | list[int] | None = None + + base_config_key = "encoder_config" + + def __post_init__(self, **kwargs): + self.sliding_window = self.sliding_window if self.use_sliding_window else None + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + if self.layer_types is None: + self.layer_types = [ + "sliding_attention" + if self.sliding_window is not None and i >= self.max_window_layers + else "full_attention" + for i in range(self.num_hidden_layers) + ] + + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="deepseek-community/DeepSeek-OCR-2") +@strict +class UnlimitedOcrVisionConfig(PreTrainedConfig): + r""" + sam_config (`dict` or `UnlimitedOcrSamVisionConfig`, *optional*): + Configuration for the SAM vision encoder. Defaults to `UnlimitedOcrSamVisionConfig()`. + encoder_config (`dict` or `UnlimitedOcrVisionEncoderConfig`, *optional*): + Configuration for the DeepSeek-OCR-2 vision encoder. Defaults to `UnlimitedOcrVisionEncoderConfig()`. + """ + + model_type = "unlimited_ocr_vision" + base_config_key = "vision_config" + sub_configs = { + "sam_config": UnlimitedOcrSamVisionConfig, + "encoder_config": UnlimitedOcrVisionEncoderConfig, + } + + sam_config: dict | PreTrainedConfig | None = None + encoder_config: dict | PreTrainedConfig | None = None + + def __post_init__(self, **kwargs): + if self.sam_config is None: + self.sam_config = UnlimitedOcrSamVisionConfig() + elif isinstance(self.sam_config, dict): + self.sam_config = UnlimitedOcrSamVisionConfig(**self.sam_config) + + if self.encoder_config is None: + self.encoder_config = UnlimitedOcrVisionEncoderConfig() + elif isinstance(self.encoder_config, dict): + self.encoder_config = UnlimitedOcrVisionEncoderConfig(**self.encoder_config) + + super().__post_init__(**kwargs) + + +@auto_docstring(checkpoint="deepseek-community/DeepSeek-OCR-2") +@strict +class UnlimitedOcrTextConfig(PreTrainedConfig): + r""" + n_group (`int`, *optional*): + Number of groups for grouped top-k expert routing. + topk_method (`str`, *optional*, defaults to `"greedy"`): + Method for selecting top-k experts in MoE layers. + mlp_layer_types (`list[str]`, *optional*): + MLP type (`"dense"` or `"sparse"`) for each decoder layer, e.g. `["dense", "sparse", "sparse", ...]`. + """ + + model_type = "unlimited_ocr_text" + keys_to_ignore_at_inference = ["past_key_values"] + + # Override DeepseekV2's MLA TP plan with standard MHA projections + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.shared_experts.gate_proj": "colwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + vocab_size: int = 32000 + hidden_size: int = 4096 + intermediate_size: int = 11008 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + num_key_value_heads: int | None = None + hidden_act: str = "silu" + max_position_embeddings: int = 2048 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-6 + use_cache: bool = True + pad_token_id: int | None = None + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + pretraining_tp: int | None = 1 + tie_word_embeddings: bool = False + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | None = 0.0 + mlp_bias: bool = False + head_dim: int | None = None + n_group: int | None = None + n_routed_experts: int = 64 + n_shared_experts: int = 2 + routed_scaling_factor: float = 1.0 + topk_group: int | None = None + topk_method: str | None = "greedy" + num_experts_per_tok: int | None = None + moe_intermediate_size: int = 1407 + + base_config_key = "text_config" + mlp_layer_types: list[str] | None = None + + def __post_init__(self, **kwargs): + self.head_dim = self.hidden_size // self.num_attention_heads + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + super().__post_init__(**kwargs) + + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.hidden_size % self.num_attention_heads != 0: + raise ValueError( + f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention " + f"heads ({self.num_attention_heads})." + ) + + +@auto_docstring(checkpoint="deepseek-community/DeepSeek-OCR-2") +@strict +class UnlimitedOcrConfig(PreTrainedConfig): + r""" + vision_config (`dict` or `UnlimitedOcrVisionConfig`, *optional*): + Configuration for the vision encoders. Defaults to `UnlimitedOcrVisionConfig()`. + """ + + model_type = "unlimited_ocr" + sub_configs = { + "vision_config": UnlimitedOcrVisionConfig, + "text_config": UnlimitedOcrTextConfig, + } + + vision_config: dict | PreTrainedConfig | None = None + text_config: dict | PreTrainedConfig | None = None + image_token_id: int = 128815 + tie_word_embeddings: bool = False + + def __post_init__(self, **kwargs): + if self.vision_config is None: + self.vision_config = UnlimitedOcrVisionConfig() + elif isinstance(self.vision_config, dict): + self.vision_config = UnlimitedOcrVisionConfig(**self.vision_config) + + if self.text_config is None: + self.text_config = UnlimitedOcrTextConfig() + elif isinstance(self.text_config, dict): + self.text_config = UnlimitedOcrTextConfig(**self.text_config) + + super().__post_init__(**kwargs) + + +__all__ = [ + "UnlimitedOcrConfig", + "UnlimitedOcrTextConfig", + "UnlimitedOcrVisionConfig", + "UnlimitedOcrVisionEncoderConfig", + "UnlimitedOcrSamVisionConfig", +] diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py new file mode 100644 index 000000000000..9cfa9042cff2 --- /dev/null +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -0,0 +1,1707 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/unlimited_ocr/modular_unlimited_ocr.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_unlimited_ocr.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_unlimited_ocr import ( + UnlimitedOcrConfig, + UnlimitedOcrSamVisionConfig, + UnlimitedOcrTextConfig, + UnlimitedOcrVisionConfig, +) + + +logger = logging.get_logger(__name__) + + +@dataclass +class UnlimitedOcrModelOutputWithPooling(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 UnlimitedOcrModelOutputWithPast(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 UnlimitedOcr causal language model (or autoregressive) outputs. + """ +) +@dataclass +class UnlimitedOcrCausalLMOutputWithPast(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 UnlimitedOcrPreTrainedModel(PreTrainedModel): + config: UnlimitedOcrConfig + base_model_prefix = "model" + input_modalities = ("image", "text") + supports_gradient_checkpointing = True + _no_split_modules = [ + "UnlimitedOcrSamVisionLayer", + "UnlimitedOcrVisionEncoderLayer", + "UnlimitedOcrTextDecoderLayer", + ] + _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, UnlimitedOcrSamVisionAttention): + if module.use_rel_pos: + init.zeros_(module.rel_pos_h) + init.zeros_(module.rel_pos_w) + elif isinstance(module, UnlimitedOcrSamVisionEncoder): + if module.pos_embed is not None: + init.zeros_(module.pos_embed) + elif isinstance(module, UnlimitedOcrModel): + embed_std = 1 / math.sqrt(self.config.text_config.hidden_size) + init.normal_(module.view_separator, mean=0.0, std=embed_std) + + +class UnlimitedOcrSamVisionAttention(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 UnlimitedOcrSamMLPBlock(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 UnlimitedOcrSamVisionSdpaAttention(UnlimitedOcrSamVisionAttention): + """ + 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 + + +UNLIMITED_OCR_SAM_VISION_ATTENTION_CLASSES = { + "eager": UnlimitedOcrSamVisionAttention, + "sdpa": UnlimitedOcrSamVisionSdpaAttention, +} + + +class UnlimitedOcrSamVisionLayer(GradientCheckpointingLayer): + def __init__(self, config, window_size): + super().__init__() + self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.attn = UNLIMITED_OCR_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 = UnlimitedOcrSamMLPBlock(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 UnlimitedOcrSamLayerNorm(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 UnlimitedOcrSamVisionNeck(nn.Module): + def __init__(self, config: UnlimitedOcrSamVisionConfig): + super().__init__() + self.config = config + + self.conv1 = nn.Conv2d(config.hidden_size, config.output_channels, kernel_size=1, bias=False) + self.layer_norm1 = UnlimitedOcrSamLayerNorm(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 = UnlimitedOcrSamLayerNorm(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 UnlimitedOcrSamPatchEmbeddings(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 UnlimitedOcrSamVisionProj(nn.Module): + """Neck and multi-scale downsampling for SAM ViT-B output.""" + + def __init__(self, config: UnlimitedOcrSamVisionConfig): + 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 UnlimitedOcrSamVisionEncoder(UnlimitedOcrPreTrainedModel): + _can_record_outputs = {"hidden_states": UnlimitedOcrSamVisionLayer, "attentions": UnlimitedOcrSamVisionAttention} + + def __init__(self, config: UnlimitedOcrSamVisionConfig): + super().__init__(config) + self.config = config + self.image_size = config.image_size + self.patch_embed = UnlimitedOcrSamPatchEmbeddings(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 = UnlimitedOcrSamVisionLayer( + config, + window_size=config.window_size if i not in config.global_attn_indexes else 0, + ) + self.layers.append(layer) + + self.neck = UnlimitedOcrSamVisionNeck(config) + + self.gradient_checkpointing = False + self.proj = UnlimitedOcrSamVisionProj(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 UnlimitedOcrVisionMLP(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 UnlimitedOcrVisionRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + UnlimitedOcrVisionRMSNorm 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 UnlimitedOcrVisionRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: UnlimitedOcrVisionConfig, 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: UnlimitedOcrVisionConfig | 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 UnlimitedOcrVisionAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: UnlimitedOcrVisionConfig, 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 UnlimitedOcrVisionEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: UnlimitedOcrVisionConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = UnlimitedOcrVisionAttention(config=config, layer_idx=layer_idx) + + self.mlp = UnlimitedOcrVisionMLP(config) + self.input_layernorm = UnlimitedOcrVisionRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = UnlimitedOcrVisionRMSNorm(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 UnlimitedOcrVisionEncoder(UnlimitedOcrPreTrainedModel): + _can_record_outputs = { + "hidden_states": UnlimitedOcrVisionEncoderLayer, + "attentions": UnlimitedOcrVisionAttention, + } + + def __init__(self, config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.layers = nn.ModuleList( + [UnlimitedOcrVisionEncoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = UnlimitedOcrVisionRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = UnlimitedOcrVisionRotaryEmbedding(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 UnlimitedOcrVisionModel(UnlimitedOcrPreTrainedModel): + """Vision pipeline: SAM ViT-B (with neck)""" + + def __init__(self, config: UnlimitedOcrVisionConfig): + super().__init__(config) + self.sam_encoder = UnlimitedOcrSamVisionEncoder(config.sam_config) + self.vision_encoder = UnlimitedOcrVisionEncoder(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 UnlimitedOcrTextRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: UnlimitedOcrTextConfig, 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: UnlimitedOcrTextConfig | 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 UnlimitedOcrTextAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: UnlimitedOcrTextConfig, 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 UnlimitedOcrTextMLP(nn.Module): + def __init__(self, config: UnlimitedOcrTextConfig, 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 UnlimitedOcrTextExperts(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 UnlimitedOcrTextMoe(nn.Module): + def __init__(self, config: UnlimitedOcrTextConfig): + super().__init__() + self.config = config + self.experts = UnlimitedOcrTextExperts(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 = UnlimitedOcrTextMLP(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 UnlimitedOcrTextRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + UnlimitedOcrTextRMSNorm 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 UnlimitedOcrTextDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = UnlimitedOcrTextAttention(config=config, layer_idx=layer_idx) + self.mlp = ( + UnlimitedOcrTextMoe(config) + if config.mlp_layer_types[layer_idx] == "sparse" + else UnlimitedOcrTextMLP(config) + ) + + self.input_layernorm = UnlimitedOcrTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = UnlimitedOcrTextRMSNorm(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 UnlimitedOcrTextPreTrainedModel(PreTrainedModel): + config: UnlimitedOcrTextConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["UnlimitedOcrTextDecoderLayer"] + _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": UnlimitedOcrTextDecoderLayer, + "attentions": UnlimitedOcrTextAttention, + } + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, UnlimitedOcrTextExperts): + 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 UnlimitedOcrTextModel(UnlimitedOcrTextPreTrainedModel): + def __init__(self, config: UnlimitedOcrTextConfig): + 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( + [UnlimitedOcrTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = UnlimitedOcrTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + # Use (cos/sin) RoPE instead of complex RoPE to match LlamaAttention (MHA) + self.rotary_emb = UnlimitedOcrTextRotaryEmbedding(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 UnlimitedOcrModel(UnlimitedOcrPreTrainedModel): + base_model_prefix = "model" + + def __init__(self, config: UnlimitedOcrConfig): + super().__init__(config) + + self.vision_tower = UnlimitedOcrVisionModel(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 = UnlimitedOcrTextModel(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 UnlimitedOcrModelOutputWithPooling( + 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 | UnlimitedOcrModelOutputWithPast: + 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 UnlimitedOcrModelOutputWithPast( + 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 UnlimitedOcrForConditionalGeneration(UnlimitedOcrPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} + + def __init__(self, config: UnlimitedOcrConfig): + super().__init__(config) + self.model = UnlimitedOcrModel(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("UnlimitedOcr 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 | UnlimitedOcrCausalLMOutputWithPast: + 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 UnlimitedOcrCausalLMOutputWithPast( + 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__ = [ + "UnlimitedOcrForConditionalGeneration", + "UnlimitedOcrModel", + "UnlimitedOcrPreTrainedModel", + "UnlimitedOcrTextModel", + "UnlimitedOcrTextPreTrainedModel", + "UnlimitedOcrVisionModel", +] diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py new file mode 100644 index 000000000000..c559d3664f14 --- /dev/null +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -0,0 +1,223 @@ +# 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 torch +from torch import nn + +from ..clip.configuration_clip import CLIPVisionConfig +from ..deepseek_ocr2.configuration_deepseek_ocr2 import ( + DeepseekOcr2Config, + DeepseekOcr2TextConfig, + DeepseekOcr2VisionConfig, +) +from ..deepseek_ocr2.modeling_deepseek_ocr2 import ( + DeepseekOcr2CausalLMOutputWithPast, + DeepseekOcr2ForConditionalGeneration, + DeepseekOcr2Model, + DeepseekOcr2ModelOutputWithPast, + DeepseekOcr2ModelOutputWithPooling, + DeepseekOcr2PreTrainedModel, + DeepseekOcr2SamLayerNorm, + DeepseekOcr2SamMLPBlock, + DeepseekOcr2SamPatchEmbeddings, + DeepseekOcr2SamVisionAttention, + DeepseekOcr2SamVisionEncoder, + DeepseekOcr2SamVisionLayer, + DeepseekOcr2SamVisionNeck, + DeepseekOcr2SamVisionProj, + DeepseekOcr2SamVisionSdpaAttention, + DeepseekOcr2TextAttention, + DeepseekOcr2TextDecoderLayer, + DeepseekOcr2TextExperts, + DeepseekOcr2TextMLP, + DeepseekOcr2TextModel, + DeepseekOcr2TextMoe, + DeepseekOcr2TextPreTrainedModel, + DeepseekOcr2TextRMSNorm, + DeepseekOcr2TextRotaryEmbedding, + DeepseekOcr2VisionAttention, + DeepseekOcr2VisionEncoder, + DeepseekOcr2VisionEncoderLayer, + DeepseekOcr2VisionMLP, + DeepseekOcr2VisionModel, + DeepseekOcr2VisionRMSNorm, + DeepseekOcr2VisionRotaryEmbedding, +) +from ..got_ocr2.configuration_got_ocr2 import GotOcr2VisionConfig + + +class UnlimitedOcrSamVisionConfig(GotOcr2VisionConfig): + pass + + +class UnlimitedOcrVisionEncoderConfig(CLIPVisionConfig): + pass + + +class UnlimitedOcrVisionConfig(DeepseekOcr2VisionConfig): + pass + + +class UnlimitedOcrTextConfig(DeepseekOcr2TextConfig): + pass + + +class UnlimitedOcrConfig(DeepseekOcr2Config): + pass + + +class UnlimitedOcrModelOutputWithPooling(DeepseekOcr2ModelOutputWithPooling): + pass + + +class UnlimitedOcrModelOutputWithPast(DeepseekOcr2ModelOutputWithPast): + pass + + +class UnlimitedOcrCausalLMOutputWithPast(DeepseekOcr2CausalLMOutputWithPast): + pass + + +class UnlimitedOcrPreTrainedModel(DeepseekOcr2PreTrainedModel): + pass + + +class UnlimitedOcrSamVisionAttention(DeepseekOcr2SamVisionAttention): + pass + + +class UnlimitedOcrSamMLPBlock(DeepseekOcr2SamMLPBlock): + pass + + +class UnlimitedOcrSamVisionSdpaAttention(DeepseekOcr2SamVisionSdpaAttention): + pass + + +class UnlimitedOcrSamVisionLayer(DeepseekOcr2SamVisionLayer): + pass + + +class UnlimitedOcrSamLayerNorm(DeepseekOcr2SamLayerNorm): + pass + + +class UnlimitedOcrSamVisionNeck(DeepseekOcr2SamVisionNeck): + pass + + +class UnlimitedOcrSamPatchEmbeddings(DeepseekOcr2SamPatchEmbeddings): + pass + + +class UnlimitedOcrSamVisionProj(DeepseekOcr2SamVisionProj): + pass + + +class UnlimitedOcrSamVisionEncoder(DeepseekOcr2SamVisionEncoder): + pass + + +class UnlimitedOcrVisionMLP(DeepseekOcr2VisionMLP): + pass + + +class UnlimitedOcrVisionRMSNorm(DeepseekOcr2VisionRMSNorm): + pass + + +class UnlimitedOcrVisionRotaryEmbedding(DeepseekOcr2VisionRotaryEmbedding): + pass + + +class UnlimitedOcrVisionAttention(DeepseekOcr2VisionAttention): + pass + + +class UnlimitedOcrVisionEncoderLayer(DeepseekOcr2VisionEncoderLayer): + pass + + +class UnlimitedOcrVisionEncoder(DeepseekOcr2VisionEncoder): + pass + + +class UnlimitedOcrVisionModel(DeepseekOcr2VisionModel): + pass + + +class UnlimitedOcrTextRotaryEmbedding(DeepseekOcr2TextRotaryEmbedding): + pass + + +class UnlimitedOcrTextAttention(DeepseekOcr2TextAttention): + pass + + +class UnlimitedOcrTextMLP(DeepseekOcr2TextMLP): + pass + + +class UnlimitedOcrTextExperts(DeepseekOcr2TextExperts): + pass + + +class UnlimitedOcrTextMoe(DeepseekOcr2TextMoe): + pass + + +class UnlimitedOcrTextRMSNorm(DeepseekOcr2TextRMSNorm): + pass + + +class UnlimitedOcrTextDecoderLayer(DeepseekOcr2TextDecoderLayer): + pass + + +class UnlimitedOcrTextPreTrainedModel(DeepseekOcr2TextPreTrainedModel): + pass + + +class UnlimitedOcrTextModel(DeepseekOcr2TextModel): + pass + + +class UnlimitedOcrModel(DeepseekOcr2Model): + def __init__(self, config: UnlimitedOcrConfig): + super().__init__(config) + n_embed = 1280 + self.multi_modal_projector = nn.Linear( + config.vision_config.sam_config.hidden_size + config.vision_config.encoder_config.hidden_size, n_embed + ) + embed_std = 1 / torch.sqrt(torch.tensor(n_embed, dtype=torch.float32)) + self.image_newline = nn.Parameter(torch.randn(n_embed) * embed_std) + self.view_seperator = nn.Parameter(torch.randn(n_embed) * embed_std) + + +class UnlimitedOcrForConditionalGeneration(DeepseekOcr2ForConditionalGeneration): + pass + + +__all__ = [ + "UnlimitedOcrConfig", + "UnlimitedOcrTextConfig", + "UnlimitedOcrVisionConfig", + "UnlimitedOcrVisionEncoderConfig", + "UnlimitedOcrSamVisionConfig", + "UnlimitedOcrForConditionalGeneration", + "UnlimitedOcrModel", + "UnlimitedOcrPreTrainedModel", + "UnlimitedOcrTextModel", + "UnlimitedOcrTextPreTrainedModel", + "UnlimitedOcrVisionModel", +] diff --git a/tests/models/unlimited_ocr/__init__.py b/tests/models/unlimited_ocr/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py new file mode 100644 index 000000000000..7cea1f9a196b --- /dev/null +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -0,0 +1,247 @@ +# 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. +"""Testing suite for the PyTorch UnlimitedOcr model.""" + +import unittest + +from transformers import ( + AutoProcessor, + UnlimitedOcrConfig, + is_torch_available, + is_vision_available, +) +from transformers.testing_utils import Expectations, cleanup, require_torch, slow, torch_device + +from ...test_processing_common import url_to_local_path +from ...vlm_tester import VLMModelTest, VLMModelTester + + +if is_torch_available(): + import torch + + from transformers import ( + UnlimitedOcrForConditionalGeneration, + UnlimitedOcrModel, + ) + from transformers.models.deepseek_ocr2.configuration_deepseek_ocr2 import ( + UnlimitedOcrTextConfig, + UnlimitedOcrVisionConfig, + ) + +if is_vision_available(): + from transformers.image_utils import load_image + + +class UnlimitedOcrVisionText2TextModelTester(VLMModelTester): + base_model_class = UnlimitedOcrModel + config_class = UnlimitedOcrConfig + conditional_generation_class = UnlimitedOcrForConditionalGeneration + text_config_class = UnlimitedOcrTextConfig + vision_config_class = UnlimitedOcrVisionConfig + + def __init__(self, parent, **kwargs): + # VisionModel always selects query_768_resolution (144 tokens) for small images + 1 separator + kwargs.setdefault("num_image_tokens", 145) + kwargs.setdefault("image_token_id", 1) + kwargs.setdefault("image_size", 16) + kwargs.setdefault("hidden_size", 128) + kwargs.setdefault("intermediate_size", 256) + kwargs.setdefault("num_hidden_layers", 2) + kwargs.setdefault("num_attention_heads", 4) + kwargs.setdefault("num_key_value_heads", 4) + kwargs.setdefault("hidden_act", "silu") + kwargs.setdefault("max_position_embeddings", 512) + kwargs.setdefault("tie_word_embeddings", False) + kwargs.setdefault("bos_token_id", 2) + kwargs.setdefault("eos_token_id", 3) + kwargs.setdefault("pad_token_id", 4) + kwargs.setdefault("n_routed_experts", 8) + kwargs.setdefault("n_shared_experts", 1) + kwargs.setdefault("mlp_layer_types", ["dense", "sparse"]) + kwargs.setdefault("moe_intermediate_size", 64) + kwargs.setdefault("num_experts_per_tok", 2) + super().__init__(parent, **kwargs) + + self.sam_config = { + "hidden_size": 32, + "output_channels": 16, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_channels": 3, + "image_size": 16, + "patch_size": 2, + "hidden_act": "gelu", + "mlp_ratio": 4.0, + "window_size": 4, + "global_attn_indexes": [1], + "downsample_channels": [32, 64], + } + self.encoder_config = { + "hidden_size": 64, + "intermediate_size": 128, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "hidden_act": "silu", + "max_position_embeddings": 512, + "rms_norm_eps": 1.0, + } + + def get_vision_config(self): + return UnlimitedOcrVisionConfig( + sam_config=self.sam_config, + encoder_config=self.encoder_config, + ) + + def get_config(self): + return self.config_class( + vision_config=self.get_vision_config(), + text_config=self.get_text_config(), + image_token_id=self.image_token_id, + ) + + +@require_torch +class UnlimitedOcrModelTest(VLMModelTest, unittest.TestCase): + model_tester_class = UnlimitedOcrVisionText2TextModelTester + test_all_params_have_gradient = False + test_torch_exportable = False + + @unittest.skip( + reason="UnlimitedOcrVisionModel builds a hybrid bidirectional+causal mask internally, so SDPA is always called with a non-null `attn_mask`." + ) + def test_sdpa_can_dispatch_on_flash(self): + pass + + @unittest.skip( + reason="UnlimitedOcrVisionModel uses `self.query_*.weight` directly, causing device mismatch when offloading." + ) + def test_cpu_offload(self): + pass + + @unittest.skip( + reason="UnlimitedOcrVisionModel uses `self.query_*.weight` directly, causing device mismatch when offloading." + ) + def test_disk_offload_bin(self): + pass + + @unittest.skip( + reason="UnlimitedOcrVisionModel uses `self.query_*.weight` directly, causing device mismatch when offloading." + ) + def test_disk_offload_safetensors(self): + pass + + def _image_features_prepare_config_and_inputs(self): + config, inputs_dict = super()._image_features_prepare_config_and_inputs() + # test_get_image_features_output expects vision_config.hidden_size, but ours is in encoder_config. + config.vision_config.hidden_size = config.vision_config.encoder_config.hidden_size + return config, inputs_dict + + +@require_torch +class UnlimitedOcrIntegrationTest(unittest.TestCase): + model_id = "deepseek-community/DeepSeek-OCR-2" + + def setUp(self): + self.processor = AutoProcessor.from_pretrained(self.model_id) + + def tearDown(self): + cleanup(torch_device, gc_collect=True) + + @slow + def test_small_model_integration_test_free_ocr(self): + model = UnlimitedOcrForConditionalGeneration.from_pretrained( + self.model_id, torch_dtype=torch.bfloat16, device_map=torch_device + ) + image = load_image( + url_to_local_path( + "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" + ) + ) + inputs = self.processor(images=image, text="\nFree OCR.", return_tensors="pt").to( + model.device, dtype=torch.bfloat16 + ) + generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) + decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True) + EXPECTED_DECODED_TEXT = Expectations( + { + ("cuda", None): "R&D QUALITY IMPROVEMENT SUGGESTION/SOLUTION FORM\n\nName/", + ("xpu", 5): "R&D QUALITY IMPROVEMENT SUGGESTION/SOLUTION FORM\n\nName/", + } + ).get_expectation() # fmt: skip + self.assertEqual(decoded, EXPECTED_DECODED_TEXT) + + @slow + def test_small_model_integration_test_grounding_markdown(self): + model = UnlimitedOcrForConditionalGeneration.from_pretrained( + self.model_id, torch_dtype=torch.bfloat16, device_map=torch_device + ) + image = load_image( + url_to_local_path( + "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" + ) + ) + inputs = self.processor( + images=image, + text="\n<|grounding|>Convert the document to markdown.", + return_tensors="pt", + ).to(model.device, dtype=torch.bfloat16) + generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) + decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=False) + EXPECTED_DECODED_TEXT = Expectations( + { + ("cuda", None): "<|ref|>title<|/ref|><|det|>[[330, 198, 559, 230]]<|/det|>\n# R", + ("xpu", 5): "<|ref|>title<|/ref|><|det|>[[330, 198, 558, 230]]<|/det|>\n# R", + } + ).get_expectation() # fmt: skip + self.assertEqual(decoded, EXPECTED_DECODED_TEXT) + + @slow + def test_small_model_integration_test_batched(self): + model = UnlimitedOcrForConditionalGeneration.from_pretrained( + self.model_id, torch_dtype=torch.bfloat16, device_map=torch_device + ) + image1 = load_image( + url_to_local_path( + "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" + ) + ) + image2 = load_image( + url_to_local_path( + "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png" + ) + ) + inputs = self.processor( + images=[image1, image2], + text=["\nFree OCR.", "\nFree OCR."], + return_tensors="pt", + padding=True, + ).to(model.device, dtype=torch.bfloat16) + generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) + decoded = self.processor.batch_decode( + generate_ids[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True + ) + EXPECTED_DECODED_TEXT = Expectations( + { + ("cuda", None): [ + "R&D QUALITY IMPROVEMENT SUGGESTION/SOLUTION FORM\n\nName/", + "# Reducing the number of images\n\nIt is also believed that the performance of a website is a critical", + ], + ("xpu", 5): [ + "R&D QUALITY IMPROVEMENT SUGGESTION/SOLUTION FORM\n\nName/", + "# Reducing the number of images\n\nIt is also believed that the performance of a website is a critical", + ], + } + ).get_expectation() # fmt: skip + self.assertEqual(decoded, EXPECTED_DECODED_TEXT) From 3d0a006adf79becd5629dddcc54b7c7e03c2f7a4 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 12:15:55 +0200 Subject: [PATCH 002/195] update --- src/transformers/models/auto/auto_mappings.py | 6 +- .../configuration_unlimited_ocr.py | 118 ++--- .../image_processing_unlimited_ocr.py | 319 +++++++++++++ .../unlimited_ocr/modeling_unlimited_ocr.py | 93 ++-- .../unlimited_ocr/modular_unlimited_ocr.py | 443 +++++++++++++++++- .../unlimited_ocr/processing_unlimited_ocr.py | 149 ++++++ 6 files changed, 998 insertions(+), 130 deletions(-) create mode 100644 src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py create mode 100644 src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py diff --git a/src/transformers/models/auto/auto_mappings.py b/src/transformers/models/auto/auto_mappings.py index 3154e30c53e0..8cf7acba456b 100644 --- a/src/transformers/models/auto/auto_mappings.py +++ b/src/transformers/models/auto/auto_mappings.py @@ -623,10 +623,10 @@ ("unispeech-sat", "UniSpeechSatConfig"), ("univnet", "UnivNetConfig"), ("unlimited_ocr", "UnlimitedOcrConfig"), - ("unlimited_ocr_encoder", "UnlimitedOcrVisionEncoderConfig"), ("unlimited_ocr_sam_vision_model", "UnlimitedOcrSamVisionConfig"), ("unlimited_ocr_text", "UnlimitedOcrTextConfig"), ("unlimited_ocr_vision", "UnlimitedOcrVisionConfig"), + ("unlimited_ocr_vision_encoder", "UnlimitedOcrVisionEncoderConfig"), ("upernet", "UperNetConfig"), ("uvdoc", "UVDocConfig"), ("uvdoc_backbone", "UVDocBackboneConfig"), @@ -896,10 +896,10 @@ ("t5gemma2_text", "t5gemma2"), ("table-transformer", "table_transformer"), ("unispeech-sat", "unispeech_sat"), - ("unlimited_ocr_encoder", "unlimited_ocr"), ("unlimited_ocr_sam_vision_model", "unlimited_ocr"), ("unlimited_ocr_text", "unlimited_ocr"), ("unlimited_ocr_vision", "unlimited_ocr"), + ("unlimited_ocr_vision_encoder", "unlimited_ocr"), ("uvdoc_backbone", "uvdoc"), ("video_llama_3_vision", "video_llama_3"), ("videoprism_text_model", "videoprism"), @@ -1099,6 +1099,7 @@ ("trocr", "TrOCRProcessor"), ("tvp", "TvpProcessor"), ("udop", "UdopProcessor"), + ("unlimited_ocr", "UnlimitedOcrProcessor"), ("vibevoice_asr", "VibeVoiceAsrProcessor"), ("video_llama_3", "VideoLlama3Processor"), ("video_llava", "VideoLlavaProcessor"), @@ -1224,6 +1225,7 @@ ("swin2sr", {"pil": "Swin2SRImageProcessorPil", "torchvision": "Swin2SRImageProcessor"}), ("textnet", {"pil": "TextNetImageProcessorPil", "torchvision": "TextNetImageProcessor"}), ("tvp", {"pil": "TvpImageProcessorPil", "torchvision": "TvpImageProcessor"}), + ("unlimited_ocr", {"torchvision": "UnlimitedOcrImageProcessor"}), ("uvdoc", {"torchvision": "UVDocImageProcessor"}), ("video_llama_3", {"pil": "VideoLlama3ImageProcessorPil", "torchvision": "VideoLlama3ImageProcessor"}), ("videomae", {"pil": "VideoMAEImageProcessorPil", "torchvision": "VideoMAEImageProcessor"}), diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index 54d2171bc010..803efdb0cf38 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -17,7 +17,6 @@ # 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. - from huggingface_hub.dataclasses import strict from ...configuration_utils import PreTrainedConfig @@ -25,25 +24,25 @@ from ...utils import auto_docstring -@auto_docstring(checkpoint="deepseek-community/DeepSeek-OCR-2") +@auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict class UnlimitedOcrSamVisionConfig(PreTrainedConfig): r""" output_channels (`int`, *optional*, defaults to 256): - The number of output channels in the SAM neck. + Dimensionality of the output channels in the Patch Encoder. + use_abs_pos (`bool`, *optional*, defaults to `True`): + Whether to use absolute position embedding. + use_rel_pos (`bool`, *optional*, defaults to `True`): + Whether to use relative position embedding. window_size (`int`, *optional*, defaults to 14): - Window size for windowed attention layers. + Window size for relative position. global_attn_indexes (`list[int]`, *optional*, defaults to `[2, 5, 8, 11]`): - Indices of encoder layers that use global (non-windowed) attention. - mlp_dim (`int`, *optional*): - Dimensionality of the MLP layer in each vision encoder block. Defaults to `hidden_size * mlp_ratio`. - downsample_channels (`list[int]`, *optional*): - The channel dimensions for the multi-scale downsampling neck layers. Defaults to `[512, 896]`. + The indexes of the global attention layers. + mlp_dim (`int`, *optional*, defaults to 3072): + The dimensionality of the MLP layer in the Transformer encoder. """ base_config_key = "sam_config" - model_type = "unlimited_ocr_sam_vision_model" - hidden_size: int = 768 output_channels: int = 256 num_hidden_layers: int = 12 @@ -56,24 +55,15 @@ class UnlimitedOcrSamVisionConfig(PreTrainedConfig): attention_dropout: float | int = 0.0 initializer_range: float = 1e-10 qkv_bias: bool = True - mlp_ratio: float = 4.0 use_abs_pos: bool = True use_rel_pos: bool = True window_size: int = 14 global_attn_indexes: list[int] | tuple[int, ...] = (2, 5, 8, 11) - mlp_dim: int | None = None - - downsample_channels: list[int] | None = None - - def __post_init__(self, **kwargs): - if self.downsample_channels is None: - self.downsample_channels = [512, 896] - self.mlp_dim = int(self.hidden_size * self.mlp_ratio) if self.mlp_dim is None else self.mlp_dim - self.scale = self.hidden_size // 2 - super().__post_init__(**kwargs) + mlp_dim: int = 3072 + model_type = "unlimited_ocr_sam_vision_model" -@auto_docstring(checkpoint="deepseek-community/DeepSeek-OCR-2") +@auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict class UnlimitedOcrVisionEncoderConfig(PreTrainedConfig): r""" @@ -86,66 +76,33 @@ class UnlimitedOcrVisionEncoderConfig(PreTrainedConfig): >>> encoder_config = config.vision_config.encoder_config ```""" - model_type = "unlimited_ocr_encoder" - keys_to_ignore_at_inference = ["past_key_values"] - - # Default tensor parallel plan for base model `UnlimitedOcrVisionEncoder` - base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - } - base_model_pp_plan = { - "embed_tokens": (["input_ids"], ["inputs_embeds"]), - "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), - "norm": (["hidden_states"], ["hidden_states"]), - } - - vocab_size: int = 151936 - hidden_size: int = 4096 - intermediate_size: int = 22016 - num_hidden_layers: int = 32 - num_attention_heads: int = 32 - num_key_value_heads: int | None = 32 - hidden_act: str = "silu" - max_position_embeddings: int = 32768 - initializer_range: float = 0.02 - rms_norm_eps: float = 1e-6 - use_cache: bool = True - tie_word_embeddings: bool = False - rope_parameters: RopeParameters | dict | None = None - use_sliding_window: bool = False - sliding_window: int | None = 4096 - max_window_layers: int = 28 - layer_types: list[str] | None = None - attention_dropout: float | int = 0.0 - pad_token_id: int | None = None - bos_token_id: int | None = None - eos_token_id: int | list[int] | None = None - + model_type = "unlimited_ocr_vision_encoder" base_config_key = "encoder_config" - def __post_init__(self, **kwargs): - self.sliding_window = self.sliding_window if self.use_sliding_window else None - if self.num_key_value_heads is None: - self.num_key_value_heads = self.num_attention_heads - - if self.layer_types is None: - self.layer_types = [ - "sliding_attention" - if self.sliding_window is not None and i >= self.max_window_layers - else "full_attention" - for i in range(self.num_hidden_layers) - ] + hidden_size: int = 768 + intermediate_size: int = 3072 + projection_dim: int = 512 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + num_channels: int = 3 + image_size: int | list[int] | tuple[int, int] | None = 224 + patch_size: int | list[int] | tuple[int, int] | None = 32 + hidden_act: str = "quick_gelu" + layer_norm_eps: float = 1e-5 + attention_dropout: int | float | None = 0.0 + initializer_range: float = 0.02 + initializer_factor: float = 1.0 - super().__post_init__(**kwargs) + def validate_architecture(self): + """Part of `@strict`-powered validation. Validates the architecture of the config.""" + if self.hidden_size % self.num_attention_heads != 0: + raise ValueError( + f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention " + f"heads ({self.num_attention_heads})." + ) -@auto_docstring(checkpoint="deepseek-community/DeepSeek-OCR-2") +@auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict class UnlimitedOcrVisionConfig(PreTrainedConfig): r""" @@ -179,7 +136,7 @@ def __post_init__(self, **kwargs): super().__post_init__(**kwargs) -@auto_docstring(checkpoint="deepseek-community/DeepSeek-OCR-2") +@auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict class UnlimitedOcrTextConfig(PreTrainedConfig): r""" @@ -245,7 +202,6 @@ class UnlimitedOcrTextConfig(PreTrainedConfig): topk_method: str | None = "greedy" num_experts_per_tok: int | None = None moe_intermediate_size: int = 1407 - base_config_key = "text_config" mlp_layer_types: list[str] | None = None @@ -264,7 +220,7 @@ def validate_architecture(self): ) -@auto_docstring(checkpoint="deepseek-community/DeepSeek-OCR-2") +@auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict class UnlimitedOcrConfig(PreTrainedConfig): r""" diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py new file mode 100644 index 000000000000..36f4f83048e9 --- /dev/null +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -0,0 +1,319 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/unlimited_ocr/modular_unlimited_ocr.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_unlimited_ocr.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. + +from functools import lru_cache + +import torch +from torchvision.transforms.v2 import functional as tvF + +from ...feature_extraction_utils import BatchFeature +from ...image_processing_backends import TorchvisionBackend +from ...image_utils import IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD, PILImageResampling, SizeDict +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import auto_docstring + + +class UnlimitedOcrImageProcessorKwargs(ImagesKwargs, total=False): + r""" + crop_to_patches (`bool`, *optional*, defaults to `self.crop_to_patches`): + Whether to crop the image to patches. Can be overridden by the `crop_to_patches` parameter in the + `preprocess` method. + min_patches (`int`, *optional*, defaults to `self.min_patches`): + The minimum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is + set to `True`. Can be overridden by the `min_patches` parameter in the `preprocess` method. + max_patches (`int`, *optional*, defaults to `self.max_patches`): + The maximum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is + set to `True`. Can be overridden by the `max_patches` parameter in the `preprocess` method. + tile_size (`int`, *optional*, defaults to `768`): + The size of each local tile. Must match the model's query embedding size. + background_color (`list[int]`, *optional*, defaults to `[127, 127, 127]`): + The background color for padding. + """ + + crop_to_patches: bool + min_patches: int + max_patches: int + + tile_size: int + background_color: list[int] + + +@lru_cache(maxsize=10) +def get_all_supported_aspect_ratios(min_image_tiles: int, max_image_tiles: int) -> list[tuple[int, int]]: + """ + Computes all allowed aspect ratios for a given minimum and maximum number of input tiles. + + This function calculates all possible arrangements of tiles that can be formed + within the constraint of the minimum and maximum number of tiles. Each arrangement is + represented by its aspect ratio (width/height) and the corresponding tile configuration. + + Args: + min_image_tiles (`int`): + The minimum number of tiles allowed. + max_image_tiles (`int`): + The maximum number of tiles allowed. + + Returns: + `list[tuple[int, int]]`: A list of tuples, each tuple representing a valid (width, height) + configuration in terms of number of tiles. + + Example: + >>> get_all_supported_aspect_ratios(1, 4) + [(1, 1), (1, 2), (2, 1), (1, 3), (3, 1), (1, 4), (2, 2), (4, 1)] + + """ + aspect_ratios = [] + for width in range(1, max_image_tiles + 1): + for height in range(1, max_image_tiles + 1): + if width * height <= max_image_tiles and width * height >= min_image_tiles: + aspect_ratios.append((width, height)) + + aspect_ratios = sorted(aspect_ratios, key=lambda x: x[0] * x[1]) + + return aspect_ratios + + +@lru_cache(maxsize=100) +def get_optimal_tiled_canvas( + original_image_size: tuple[int, int], + target_tile_size: tuple[int, int], + min_image_tiles: int, + max_image_tiles: int, +) -> tuple[int, int]: + """ + Given a minimum and maximum number of tiles, find the canvas with the closest aspect ratio to the + original image aspect ratio. + In case of tie-breaking condition when two canvases have the same aspect ratio difference, we favor the canvas with + more tiles, until the area covered by the tiles is more than twice the target area, in order to avoid unnecessarily + excessive tiling. + """ + possible_tile_arrangements = get_all_supported_aspect_ratios(min_image_tiles, max_image_tiles) + + original_height, original_width = original_image_size + target_tile_height, target_tile_width = target_tile_size + aspect_ratio = original_width / original_height + area = original_width * original_height + + # find the grid with the best aspect ratio + best_ratio_diff = float("inf") + best_grid = (1, 1) + for grid in possible_tile_arrangements: + grid_aspect_ratio = grid[0] / grid[1] + ratio_diff = abs(aspect_ratio - grid_aspect_ratio) + if ratio_diff < best_ratio_diff: + best_ratio_diff = ratio_diff + best_grid = grid + elif ratio_diff == best_ratio_diff: + # if the aspect ratio difference is the same, we favor the grid with more patches + # until the area covered by the patches is more than twice the original image area + if area > 0.5 * target_tile_height * target_tile_width * grid[0] * grid[1]: + best_grid = grid + + return best_grid + + +@auto_docstring +class UnlimitedOcrImageProcessor(TorchvisionBackend): + valid_kwargs = UnlimitedOcrImageProcessorKwargs + resample = PILImageResampling.BICUBIC + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 1024, "width": 1024} + do_resize = True + do_rescale = True + do_normalize = True + do_convert_rgb = True + crop_to_patches = True + min_patches = 2 + max_patches = 32 + tile_size = 640 + background_color = [127, 127, 127] + model_input_names = ["pixel_values", "num_local_patches", "image_spatial_crop"] + + def __init__(self, **kwargs: Unpack[UnlimitedOcrImageProcessorKwargs]): + super().__init__(**kwargs) + + def crop_image_to_patches( + self, + images: "torch.Tensor", + min_patches: int, + max_patches: int, + tile_size: int, + resample: PILImageResampling | None = None, + ) -> tuple["torch.Tensor", int]: + """ + Crop batched images to patches based on optimal tiling. + + Args: + images (`torch.Tensor`): + The images to crop, shape `(batch, channels, height, width)`. + min_patches (`int`): + Minimum number of patches. + max_patches (`int`): + Maximum number of patches. + tile_size (`int`): + The size of each tile. + resample (`PILImageResampling`, *optional*): + Resampling filter for resizing. + + Returns: + `tuple[torch.Tensor, int]`: Stacked patches `(batch, num_patches, channels, tile_size, tile_size)` + and number of patches per image. + """ + original_height, original_width = images.shape[-2:] + + num_columns, num_rows = get_optimal_tiled_canvas( + (original_height, original_width), (tile_size, tile_size), min_patches, max_patches + ) + + target_width = tile_size * num_columns + target_height = tile_size * num_rows + num_blocks = num_columns * num_rows + + resized = self.resize(images, SizeDict(height=target_height, width=target_width), resample=resample) + + patches = [] + for i in range(num_blocks): + col = i % num_columns + row = i // num_columns + patch = resized[ + ..., + row * tile_size : (row + 1) * tile_size, + col * tile_size : (col + 1) * tile_size, + ] + patches.append(patch) + + stacked_patches = torch.stack(patches, dim=1) + + return stacked_patches, num_blocks + + def _preprocess( + self, + images: list["torch.Tensor"], + size, + crop_to_patches: bool, + min_patches: int, + max_patches: int, + tile_size: int, + resample, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean, + image_std, + disable_grouping: bool | None, + return_tensors, + **kwargs, + ) -> BatchFeature: + batch_feature = super()._preprocess( + images, + size=size, + crop_to_patches=crop_to_patches, + min_patches=min_patches, + max_patches=max_patches, + tile_size=tile_size, + resample=resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + disable_grouping=disable_grouping, + return_tensors=return_tensors, + **kwargs, + ) + + image_spatial_crop = [] + for image in images: + height, width = image.shape[-2:] + if crop_to_patches and max(height, width) > tile_size: + num_columns, num_rows = get_optimal_tiled_canvas( + (height, width), (tile_size, tile_size), min_patches, max_patches + ) + else: + num_columns, num_rows = 1, 1 + image_spatial_crop.append([num_columns, num_rows]) + batch_feature["image_spatial_crop"] = torch.tensor(image_spatial_crop, dtype=torch.long) + return batch_feature + + def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None) -> int: + """ + Returns the number of image patches for a given image size (1 global + local patches). + """ + if images_kwargs is None: + images_kwargs = {} + min_patches = images_kwargs.get("min_patches", self.min_patches) + max_patches = images_kwargs.get("max_patches", self.max_patches) + tile_size = images_kwargs.get("tile_size", self.tile_size) + crop_to_patches = images_kwargs.get("crop_to_patches", self.crop_to_patches) + + num_patches = 1 # global view + if crop_to_patches and max(height, width) > tile_size: + num_columns, num_rows = get_optimal_tiled_canvas( + (height, width), (tile_size, tile_size), min_patches, max_patches + ) + num_patches += num_columns * num_rows + + return num_patches + + def pad_to_square( + self, + images: "torch.Tensor", + background_color: int | tuple[int, int, int] = 0, + ) -> "torch.Tensor": + """ + Pads an image to a square based on the longest edge. + + Args: + images (`torch.Tensor`): + The images to pad. Shape: (batch_size, num_channels, height, width) or (num_channels, height, width). + background_color (`int` or `tuple[int, int, int]`, *optional*, defaults to 0): + The color to use for the padding. Can be an integer for single channel or a + tuple of integers representing for multi-channel images. If passed as integer + in multi-channel mode, it will default to `0` in subsequent channels. + Returns: + `torch.Tensor`: The padded images. + """ + height, width = images.shape[-2:] + + if height == width: + return images + + num_channels = images.shape[1] if len(images.shape) == 4 else images.shape[0] + if isinstance(background_color, int): + background_color = [background_color] + [0] * (num_channels - 1) + elif len(background_color) != num_channels: + raise ValueError( + f"background_color must have no more than {num_channels} elements to match the number of channels" + ) + + max_dim = max(height, width) + paste_x_left = (max_dim - width) // 2 + paste_y_left = (max_dim - height) // 2 + paste_x_right = max_dim - width - paste_x_left + paste_y_right = max_dim - height - paste_y_left + padded_images = tvF.pad( + images, padding=[paste_x_left, paste_y_left, paste_x_right, paste_y_right], fill=background_color + ) + + return padded_images + + +__all__ = ["UnlimitedOcrImageProcessor"] diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 9cfa9042cff2..8d23f6f055ef 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -17,7 +17,6 @@ # 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 @@ -973,35 +972,23 @@ def __init__(self, config: UnlimitedOcrVisionConfig): super().__init__(config) self.sam_encoder = UnlimitedOcrSamVisionEncoder(config.sam_config) self.vision_encoder = UnlimitedOcrVisionEncoder(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 + vision_encoder_outputs = self.vision_encoder(pixel_values, **kwargs) - 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:, :] + sam_hidden_state = sam_encoder_outputs.last_hidden_state + vision_encoder_hidden_state = vision_encoder_outputs.last_hidden_state + hidden_state = torch.cat([vision_encoder_hidden_state[:, 1:], sam_hidden_state], dim=1) + # TODO: How to pass hidden_states and attentions? return BaseModelOutput( - last_hidden_state=query_features, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, + last_hidden_state=hidden_state, + # hidden_states=encoder_outputs.hidden_states, + # attentions=encoder_outputs.attentions, ) @@ -1423,7 +1410,8 @@ def __init__(self, config: UnlimitedOcrConfig): self.vision_tower = UnlimitedOcrVisionModel(config.vision_config) self.multi_modal_projector = nn.Linear( - config.vision_config.encoder_config.hidden_size, config.text_config.hidden_size + config.vision_config.sam_config.hidden_size + config.vision_config.encoder_config.hidden_size, + config.text_config.hidden_size, ) self.vocab_size = config.text_config.vocab_size @@ -1432,6 +1420,8 @@ def __init__(self, config: UnlimitedOcrConfig): # Learnable separator between local and global views (initialized in `_init_weights`). self.view_separator = nn.Parameter(torch.empty(config.text_config.hidden_size)) + embed_std = 1 / torch.sqrt(torch.tensor(config.text_config.hidden_size, dtype=torch.float32)) + self.image_newline = nn.Parameter(torch.randn(config.text_config.hidden_size) * embed_std) self.post_init() @can_return_tuple @@ -1441,15 +1431,18 @@ def get_image_features( pixel_values: torch.FloatTensor, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, + image_spatial_crop: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutputWithPooling: + ) -> "UnlimitedOcrModelOutputWithPooling": 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]`. + image_spatial_crop (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + The local crop grid `(num_columns, num_rows)` per image. """ - # torch.split requires list[int], not Tensor, for per-image variable-length splitting + # `torch.split` requires `list[int]`, not a Tensor, for per-image variable-length splitting. if isinstance(num_local_patches, torch.Tensor): num_local_patches = num_local_patches.tolist() @@ -1471,16 +1464,37 @@ def get_image_features( else: per_image_local = [None] * batch_size + hidden_size = global_features.shape[-1] + newline = self.image_newline[None, None, :] + view_separator = self.view_separator[None, :] + 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)) + num_queries_global = int(global_features.shape[1] ** 0.5) + global_grid = global_features[idx].reshape(num_queries_global, num_queries_global, hidden_size) + global_grid = torch.cat([global_grid, newline.expand(num_queries_global, 1, hidden_size)], dim=1) + global_flat = global_grid.reshape(-1, hidden_size) + + local_features = per_image_local[idx] + if local_features is not None and local_features.shape[0] > 0: + num_columns, num_rows = int(image_spatial_crop[idx][0]), int(image_spatial_crop[idx][1]) + num_queries_local = int(local_features.shape[1] ** 0.5) + local_grid = local_features.reshape( + num_rows, num_columns, num_queries_local, num_queries_local, hidden_size + ) + local_grid = local_grid.permute(0, 2, 1, 3, 4).reshape( + num_rows * num_queries_local, num_columns * num_queries_local, hidden_size + ) + local_grid = torch.cat( + [local_grid, newline.expand(num_rows * num_queries_local, 1, hidden_size)], dim=1 + ) + local_flat = local_grid.reshape(-1, hidden_size) + # NOTE: local-then-global ordering matches the reference `forward` (the feature order the weights + # were trained on), NOT the token order built by the reference `infer`. Verify against a reference + # generation on a cropped image before trusting crop mode. + all_features.append(torch.cat([local_flat, global_flat, view_separator], dim=0)) else: - all_features.append(torch.cat([global_flat, view_sep], dim=0)) + all_features.append(torch.cat([global_flat, view_separator], dim=0)) image_features = torch.cat(all_features, dim=0) return UnlimitedOcrModelOutputWithPooling( @@ -1522,6 +1536,7 @@ def forward( pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, + image_spatial_crop: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -1534,6 +1549,8 @@ def forward( 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. + image_spatial_crop (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + The local crop grid `(num_columns, num_rows)` per image. """ if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) @@ -1541,7 +1558,7 @@ def forward( 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 + pixel_values, pixel_values_local, num_local_patches, image_spatial_crop, return_dict=True ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) @@ -1615,6 +1632,7 @@ def forward( pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, + image_spatial_crop: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -1629,12 +1647,15 @@ def forward( 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. + image_spatial_crop (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + The local crop grid `(num_columns, num_rows)` per image. """ outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, pixel_values_local=pixel_values_local, num_local_patches=num_local_patches, + image_spatial_crop=image_spatial_crop, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, @@ -1674,6 +1695,7 @@ def prepare_inputs_for_generation( pixel_values=None, pixel_values_local=None, num_local_patches=None, + image_spatial_crop=None, attention_mask=None, logits_to_keep=None, is_first_iteration=False, @@ -1683,6 +1705,9 @@ def prepare_inputs_for_generation( input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_values_local=pixel_values_local, + num_local_patches=num_local_patches, attention_mask=attention_mask, logits_to_keep=logits_to_keep, is_first_iteration=is_first_iteration, @@ -1690,9 +1715,7 @@ def prepare_inputs_for_generation( ) 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 + model_inputs["image_spatial_crop"] = image_spatial_crop return model_inputs diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index c559d3664f14..58d024898a39 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -11,15 +11,26 @@ # 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 math + import torch +from huggingface_hub.dataclasses import strict from torch import nn +from ...cache_utils import Cache +from ...feature_extraction_utils import BatchFeature +from ...image_utils import ImageInput +from ...modeling_outputs import BaseModelOutput +from ...processing_utils import Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple from ..clip.configuration_clip import CLIPVisionConfig from ..deepseek_ocr2.configuration_deepseek_ocr2 import ( DeepseekOcr2Config, DeepseekOcr2TextConfig, DeepseekOcr2VisionConfig, ) +from ..deepseek_ocr2.image_processing_deepseek_ocr2 import DeepseekOcr2ImageProcessor, get_optimal_tiled_canvas from ..deepseek_ocr2.modeling_deepseek_ocr2 import ( DeepseekOcr2CausalLMOutputWithPast, DeepseekOcr2ForConditionalGeneration, @@ -53,27 +64,187 @@ DeepseekOcr2VisionRMSNorm, DeepseekOcr2VisionRotaryEmbedding, ) +from ..deepseek_ocr2.processing_deepseek_ocr2 import DeepseekOcr2Processor, DeepseekOcr2ProcessorKwargs from ..got_ocr2.configuration_got_ocr2 import GotOcr2VisionConfig +class UnlimitedOcrImageProcessor(DeepseekOcr2ImageProcessor): + tile_size = 640 + max_patches = 32 + model_input_names = ["pixel_values", "num_local_patches", "image_spatial_crop"] + + def _preprocess( + self, + images: list["torch.Tensor"], + size, + crop_to_patches: bool, + min_patches: int, + max_patches: int, + tile_size: int, + resample, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean, + image_std, + disable_grouping: bool | None, + return_tensors, + **kwargs, + ) -> BatchFeature: + batch_feature = super()._preprocess( + images, + size=size, + crop_to_patches=crop_to_patches, + min_patches=min_patches, + max_patches=max_patches, + tile_size=tile_size, + resample=resample, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + disable_grouping=disable_grouping, + return_tensors=return_tensors, + **kwargs, + ) + + image_spatial_crop = [] + for image in images: + height, width = image.shape[-2:] + if crop_to_patches and max(height, width) > tile_size: + num_columns, num_rows = get_optimal_tiled_canvas( + (height, width), (tile_size, tile_size), min_patches, max_patches + ) + else: + num_columns, num_rows = 1, 1 + image_spatial_crop.append([num_columns, num_rows]) + batch_feature["image_spatial_crop"] = torch.tensor(image_spatial_crop, dtype=torch.long) + return batch_feature + + +class UnlimitedOcrProcessorKwargs(DeepseekOcr2ProcessorKwargs, total=False): + _defaults = { + "text_kwargs": { + "padding": False, + }, + } + + +class UnlimitedOcrProcessor(DeepseekOcr2Processor): + def _expand_image_tokens( + self, + text: list[TextInput], + image_spatial_crop: torch.Tensor, + num_local_patches: list[int] | torch.Tensor, + ) -> list[str]: + size = self.image_processor.size["height"] + tile_size = self.image_processor.tile_size + + num_queries_global = math.ceil(size / self.patch_size / self.downsample_ratio) + num_queries_local = math.ceil(tile_size / self.patch_size / self.downsample_ratio) + + crop_index = 0 + for i in range(len(text)): + while self.image_token in text[i]: + num_columns = int(image_spatial_crop[crop_index][0]) + num_rows = int(image_spatial_crop[crop_index][1]) + num_tokens = num_queries_global * (num_queries_global + 1) + 1 + if int(num_local_patches[crop_index]) > 0: + num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) + text[i] = text[i].replace(self.image_token, "<|placeholder|>" * num_tokens, 1) + crop_index += 1 + text[i] = text[i].replace("<|placeholder|>", self.image_token) + return text + + @auto_docstring + def __call__( + self, + images: ImageInput | None = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + **kwargs: Unpack[UnlimitedOcrProcessorKwargs], + ) -> BatchFeature: + if images is None: + raise ValueError("`images` are expected as arguments to a `UnlimitedOcrProcessor` instance.") + if text is None: + raise ValueError("`text` is required for `UnlimitedOcrProcessor`. Example: `'\\nFree OCR.'`") + + output_kwargs = self._merge_kwargs( + UnlimitedOcrProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + if isinstance(text, str): + text = [text] + elif not (isinstance(text, (list, tuple)) and all(isinstance(t, str) for t in text)): + raise TypeError("Invalid input text. Please provide a string, or a list of strings") + + text = text.copy() # below lines change text in-place + + image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) + text = self._expand_image_tokens(text, image_inputs["image_spatial_crop"], image_inputs["num_local_patches"]) + + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) + self._check_special_mm_tokens(text, text_inputs, modalities=["image"]) + + return BatchFeature( + data={**text_inputs, **image_inputs}, + tensor_type=return_tensors, + ) + + +@auto_docstring(checkpoint="baidu/Unlimited-OCR") +@strict class UnlimitedOcrSamVisionConfig(GotOcr2VisionConfig): - pass + model_type = "unlimited_ocr_sam_vision_model" + base_config_key = "sam_config" +@auto_docstring(checkpoint="baidu/Unlimited-OCR") +@strict class UnlimitedOcrVisionEncoderConfig(CLIPVisionConfig): - pass + r""" + Example: + ```python + >>> from transformers import UnlimitedOcrConfig + >>> config = UnlimitedOcrConfig() + >>> encoder_config = config.vision_config.encoder_config + ```""" + + model_type = "unlimited_ocr_vision_encoder" + base_config_key = "encoder_config" + + +@auto_docstring(checkpoint="baidu/Unlimited-OCR") +@strict class UnlimitedOcrVisionConfig(DeepseekOcr2VisionConfig): - pass + model_type = "unlimited_ocr_vision" + base_config_key = "vision_config" + sub_configs = { + "sam_config": UnlimitedOcrSamVisionConfig, + "encoder_config": UnlimitedOcrVisionEncoderConfig, + } +@auto_docstring(checkpoint="baidu/Unlimited-OCR") +@strict class UnlimitedOcrTextConfig(DeepseekOcr2TextConfig): - pass + model_type = "unlimited_ocr_text" + base_config_key = "text_config" +@auto_docstring(checkpoint="baidu/Unlimited-OCR") +@strict class UnlimitedOcrConfig(DeepseekOcr2Config): - pass + model_type = "unlimited_ocr" + sub_configs = { + "vision_config": UnlimitedOcrVisionConfig, + "text_config": UnlimitedOcrTextConfig, + } class UnlimitedOcrModelOutputWithPooling(DeepseekOcr2ModelOutputWithPooling): @@ -153,7 +324,27 @@ class UnlimitedOcrVisionEncoder(DeepseekOcr2VisionEncoder): class UnlimitedOcrVisionModel(DeepseekOcr2VisionModel): - pass + def __init__(self, config: UnlimitedOcrVisionConfig): + super().__init__(config) + del self.query_768_resolution + del self.query_1024_resolution + + @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) + vision_encoder_outputs = self.vision_encoder(pixel_values, **kwargs) + + sam_hidden_state = sam_encoder_outputs.last_hidden_state + vision_encoder_hidden_state = vision_encoder_outputs.last_hidden_state + hidden_state = torch.cat([vision_encoder_hidden_state[:, 1:], sam_hidden_state], dim=1) + + # TODO: How to pass hidden_states and attentions? + return BaseModelOutput( + last_hidden_state=hidden_state, + # hidden_states=encoder_outputs.hidden_states, + # attentions=encoder_outputs.attentions, + ) class UnlimitedOcrTextRotaryEmbedding(DeepseekOcr2TextRotaryEmbedding): @@ -195,17 +386,243 @@ class UnlimitedOcrTextModel(DeepseekOcr2TextModel): class UnlimitedOcrModel(DeepseekOcr2Model): def __init__(self, config: UnlimitedOcrConfig): super().__init__(config) - n_embed = 1280 self.multi_modal_projector = nn.Linear( - config.vision_config.sam_config.hidden_size + config.vision_config.encoder_config.hidden_size, n_embed + config.vision_config.sam_config.hidden_size + config.vision_config.encoder_config.hidden_size, + config.text_config.hidden_size, + ) + embed_std = 1 / torch.sqrt(torch.tensor(config.text_config.hidden_size, dtype=torch.float32)) + self.image_newline = nn.Parameter(torch.randn(config.text_config.hidden_size) * embed_std) + + 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, + image_spatial_crop: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> "UnlimitedOcrModelOutputWithPooling": + 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]`. + image_spatial_crop (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + The local crop grid `(num_columns, num_rows)` per image. + """ + # `torch.split` requires `list[int]`, not a 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 + + hidden_size = global_features.shape[-1] + newline = self.image_newline[None, None, :] + view_separator = self.view_separator[None, :] + + all_features = [] + for idx in range(batch_size): + num_queries_global = int(global_features.shape[1] ** 0.5) + global_grid = global_features[idx].reshape(num_queries_global, num_queries_global, hidden_size) + global_grid = torch.cat([global_grid, newline.expand(num_queries_global, 1, hidden_size)], dim=1) + global_flat = global_grid.reshape(-1, hidden_size) + + local_features = per_image_local[idx] + if local_features is not None and local_features.shape[0] > 0: + num_columns, num_rows = int(image_spatial_crop[idx][0]), int(image_spatial_crop[idx][1]) + num_queries_local = int(local_features.shape[1] ** 0.5) + local_grid = local_features.reshape( + num_rows, num_columns, num_queries_local, num_queries_local, hidden_size + ) + local_grid = local_grid.permute(0, 2, 1, 3, 4).reshape( + num_rows * num_queries_local, num_columns * num_queries_local, hidden_size + ) + local_grid = torch.cat( + [local_grid, newline.expand(num_rows * num_queries_local, 1, hidden_size)], dim=1 + ) + local_flat = local_grid.reshape(-1, hidden_size) + # NOTE: local-then-global ordering matches the reference `forward` (the feature order the weights + # were trained on), NOT the token order built by the reference `infer`. Verify against a reference + # generation on a cropped image before trusting crop mode. + all_features.append(torch.cat([local_flat, global_flat, view_separator], dim=0)) + else: + all_features.append(torch.cat([global_flat, view_separator], dim=0)) + + image_features = torch.cat(all_features, dim=0) + return UnlimitedOcrModelOutputWithPooling( + 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, + ) + + @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, + image_spatial_crop: 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 | UnlimitedOcrModelOutputWithPast: + 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. + image_spatial_crop (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + The local crop grid `(num_columns, num_rows)` per image. + """ + 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, image_spatial_crop, 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 UnlimitedOcrModelOutputWithPast( + 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, ) - embed_std = 1 / torch.sqrt(torch.tensor(n_embed, dtype=torch.float32)) - self.image_newline = nn.Parameter(torch.randn(n_embed) * embed_std) - self.view_seperator = nn.Parameter(torch.randn(n_embed) * embed_std) class UnlimitedOcrForConditionalGeneration(DeepseekOcr2ForConditionalGeneration): - pass + @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, + image_spatial_crop: 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 | UnlimitedOcrCausalLMOutputWithPast: + 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. + image_spatial_crop (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + The local crop grid `(num_columns, num_rows)` per image. + """ + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + pixel_values_local=pixel_values_local, + num_local_patches=num_local_patches, + image_spatial_crop=image_spatial_crop, + 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 UnlimitedOcrCausalLMOutputWithPast( + 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, + image_spatial_crop=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, + pixel_values=pixel_values, + pixel_values_local=pixel_values_local, + num_local_patches=num_local_patches, + 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["image_spatial_crop"] = image_spatial_crop + + return model_inputs __all__ = [ @@ -215,8 +632,10 @@ class UnlimitedOcrForConditionalGeneration(DeepseekOcr2ForConditionalGeneration) "UnlimitedOcrVisionEncoderConfig", "UnlimitedOcrSamVisionConfig", "UnlimitedOcrForConditionalGeneration", + "UnlimitedOcrImageProcessor", "UnlimitedOcrModel", "UnlimitedOcrPreTrainedModel", + "UnlimitedOcrProcessor", "UnlimitedOcrTextModel", "UnlimitedOcrTextPreTrainedModel", "UnlimitedOcrVisionModel", diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py new file mode 100644 index 000000000000..8c0b571f6040 --- /dev/null +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -0,0 +1,149 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/unlimited_ocr/modular_unlimited_ocr.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_unlimited_ocr.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 math + +import torch + +from ...feature_extraction_utils import BatchFeature +from ...image_utils import ImageInput +from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import auto_docstring + + +class UnlimitedOcrProcessorKwargs(ProcessingKwargs, total=False): + _defaults = { + "text_kwargs": { + "padding": False, + }, + } + + +@auto_docstring +class UnlimitedOcrProcessor(ProcessorMixin): + def __init__( + self, + image_processor=None, + tokenizer=None, + chat_template=None, + patch_size=16, + downsample_ratio=4, + **kwargs, + ): + r""" + patch_size (`int`, *optional*, defaults to `16`): + The patch size used by the vision encoder (SAM ViT patch embedding size). + downsample_ratio (`int`, *optional*, defaults to `4`): + The downsampling ratio applied after the vision encoder. + """ + self.image_token = "" if not hasattr(tokenizer, "image_token") else tokenizer.image_token + self.patch_size = patch_size + self.downsample_ratio = downsample_ratio + self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token) + super().__init__(image_processor, tokenizer, chat_template=chat_template, **kwargs) + + def _expand_image_tokens( + self, + text: list[TextInput], + image_spatial_crop: torch.Tensor, + num_local_patches: list[int] | torch.Tensor, + ) -> list[str]: + """ + Expand each `` placeholder in the text to the correct number of image tokens. + + Args: + text (`list[str]`): + List of text strings, each potentially containing `` placeholders. + num_crops_list (`list[int]`): + Number of crops for each image, consumed in order as `` placeholders + are encountered across all text samples. + + Returns: + `list[str]`: Text with expanded image token placeholders. + """ + size = self.image_processor.size["height"] + tile_size = self.image_processor.tile_size + + num_queries_global = math.ceil(size / self.patch_size / self.downsample_ratio) + num_queries_local = math.ceil(tile_size / self.patch_size / self.downsample_ratio) + + crop_index = 0 + for i in range(len(text)): + while self.image_token in text[i]: + num_columns = int(image_spatial_crop[crop_index][0]) + num_rows = int(image_spatial_crop[crop_index][1]) + num_tokens = num_queries_global * (num_queries_global + 1) + 1 + if int(num_local_patches[crop_index]) > 0: + num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) + text[i] = text[i].replace(self.image_token, "<|placeholder|>" * num_tokens, 1) + crop_index += 1 + text[i] = text[i].replace("<|placeholder|>", self.image_token) + return text + + @auto_docstring + def __call__( + self, + images: ImageInput | None = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + **kwargs: Unpack[UnlimitedOcrProcessorKwargs], + ) -> BatchFeature: + r""" + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Global view pixel values. Returned when `images` is not `None`. + - **pixel_values_local** -- Local patch pixel values. Returned when `images` is not `None`. + """ + if images is None: + raise ValueError("`images` are expected as arguments to a `UnlimitedOcrProcessor` instance.") + if text is None: + raise ValueError("`text` is required for `UnlimitedOcrProcessor`. Example: `'\\nFree OCR.'`") + + output_kwargs = self._merge_kwargs( + UnlimitedOcrProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + if isinstance(text, str): + text = [text] + elif not (isinstance(text, (list, tuple)) and all(isinstance(t, str) for t in text)): + raise TypeError("Invalid input text. Please provide a string, or a list of strings") + + text = text.copy() # below lines change text in-place + + image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) + text = self._expand_image_tokens(text, image_inputs["image_spatial_crop"], image_inputs["num_local_patches"]) + + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) + self._check_special_mm_tokens(text, text_inputs, modalities=["image"]) + + return BatchFeature( + data={**text_inputs, **image_inputs}, + tensor_type=return_tensors, + ) + + +__all__ = ["UnlimitedOcrProcessor"] From 16a7260c70a6e2d3fe38a402d71bdf6950c505bd Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 12:41:37 +0200 Subject: [PATCH 003/195] add CLIP --- .../unlimited_ocr/modeling_unlimited_ocr.py | 318 ++++++++++++++---- .../unlimited_ocr/modular_unlimited_ocr.py | 96 +++++- 2 files changed, 341 insertions(+), 73 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 8d23f6f055ef..9f193a2fc571 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -40,11 +40,19 @@ 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_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling 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 import ( + ModelOutput, + TransformersKwargs, + auto_docstring, + can_return_tuple, + logging, + torch_compilable_check, + torch_int, +) from ...utils.generic import maybe_autocast, merge_with_config_defaults from ...utils.output_capturing import capture_outputs from .configuration_unlimited_ocr import ( @@ -898,71 +906,261 @@ def forward( return hidden_states -@auto_docstring(custom_intro="Vision encoder for DeepSeek-OCR-2.") -class UnlimitedOcrVisionEncoder(UnlimitedOcrPreTrainedModel): - _can_record_outputs = { - "hidden_states": UnlimitedOcrVisionEncoderLayer, - "attentions": UnlimitedOcrVisionAttention, - } +class UnlimitedOcrVisionEmbeddings(nn.Module): + def __init__(self, config: UnlimitedOcrVisionConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size - def __init__(self, config): - super().__init__(config) - self.padding_idx = config.pad_token_id - self.vocab_size = config.vocab_size - self.layers = nn.ModuleList( - [UnlimitedOcrVisionEncoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + self.class_embedding = nn.Parameter(torch.randn(self.embed_dim)) + + self.patch_embedding = nn.Conv2d( + in_channels=config.num_channels, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + bias=False, ) - self.norm = UnlimitedOcrVisionRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.rotary_emb = UnlimitedOcrVisionRotaryEmbedding(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() + self.num_patches = (self.image_size // self.patch_size) ** 2 + self.num_positions = self.num_patches + 1 + self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False) + + def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: + """ + This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution + images. This method is also adapted to support torch.jit tracing. + + Adapted from: + - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and + - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 + """ + num_patches = embeddings.shape[1] - 1 + position_embedding = self.position_embedding.weight.unsqueeze(0) + num_positions = position_embedding.shape[1] - 1 + + if not torch.jit.is_tracing() and num_patches == num_positions and height == width: + return self.position_embedding(self.position_ids) + + class_pos_embed = position_embedding[:, :1] + patch_pos_embed = position_embedding[:, 1:] + + dim = embeddings.shape[-1] + + new_height = height // self.patch_size + new_width = width // self.patch_size + + sqrt_num_positions = torch_int(num_positions**0.5) + patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) + patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) + + target_dtype = patch_pos_embed.dtype + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed.to(torch.float32), + size=(new_height, new_width), + mode="bicubic", + antialias=True, + align_corners=False, + ).to(target_dtype) + + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + + return torch.cat((class_pos_embed, patch_pos_embed), dim=1) + + def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: + r""" + patch_embeds (`torch.Tensor` of shape `(batch_size, hidden_size, grid_height, grid_width)`): + The SAM feature map, injected directly as the CLIP patch embeddings. The CLIP patch convolution + (`self.patch_embedding`) is intentionally bypassed. + """ + batch_size, _, grid_height, grid_width = patch_embeds.shape + patch_embeds = patch_embeds.flatten(2).transpose(1, 2) + class_embeds = self.class_embedding.expand(batch_size, 1, -1) + embeddings = torch.cat([class_embeds, patch_embeds], dim=1) + embeddings = embeddings + self.interpolate_pos_encoding( + embeddings, grid_height * self.patch_size, grid_width * self.patch_size + ) + return embeddings + + +class UnlimitedOcrAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: UnlimitedOcrVisionConfig | UnlimitedOcrTextConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + self.scale = self.head_dim**-0.5 + self.dropout = config.attention_dropout + self.is_causal = False + + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) - @merge_with_config_defaults - @capture_outputs - @auto_docstring def forward( self, - inputs_embeds: torch.FloatTensor, - num_patches: int, - position_ids: torch.LongTensor | None = None, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | 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) + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Input shape: Batch x Time x Channel""" - 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, + input_shape = hidden_states.shape[:-1] + + hidden_shape = (*input_shape, -1, self.head_dim) + queries = self.q_proj(hidden_states) + keys = self.k_proj(hidden_states) + values = self.v_proj(hidden_states) + + queries = queries.view(hidden_shape).transpose(1, 2) + keys = keys.view(hidden_shape).transpose(1, 2) + values = values.view(hidden_shape).transpose(1, 2) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward ) - hidden_states = inputs_embeds - position_embeddings = self.rotary_emb(hidden_states, position_ids) + attn_output, attn_weights = attention_interface( + self, + queries, + keys, + values, + attention_mask, + scaling=self.scale, + dropout=0.0 if not self.training else self.dropout, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.out_proj(attn_output) - for encoder_layer in self.layers[: self.config.num_hidden_layers]: + return attn_output, attn_weights + + +class UnlimitedOcrMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + + +class UnlimitedOcrEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: UnlimitedOcrVisionConfig | UnlimitedOcrTextConfig): + super().__init__() + self.embed_dim = config.hidden_size + self.self_attn = UnlimitedOcrAttention(config) + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.mlp = UnlimitedOcrMLP(config) + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.FloatTensor: + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + **kwargs, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class UnlimitedOcrEncoder(nn.Module): + """ + Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a + [`UnlimitedOcrEncoderLayer`]. + + Args: + config: UnlimitedOcrConfig + """ + + def __init__(self, config: UnlimitedOcrConfig): + super().__init__() + self.config = config + self.layers = nn.ModuleList([UnlimitedOcrEncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + def forward( + self, + inputs_embeds, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutput: + hidden_states = inputs_embeds + for encoder_layer in self.layers: hidden_states = encoder_layer( hidden_states, - attention_mask=attention_mask, - position_embeddings=position_embeddings, - position_ids=position_ids, + attention_mask, **kwargs, ) - hidden_states = self.norm(hidden_states) - return BaseModelOutputWithPast(last_hidden_state=hidden_states) + return BaseModelOutput( + last_hidden_state=hidden_states, + ) + + +@auto_docstring( + custom_intro=""" + The vision model from UNLIMITED_OCR without any head or projection on top. + """ +) +class UnlimitedOcrVisionEncoder(UnlimitedOcrPreTrainedModel): + config: UnlimitedOcrVisionConfig + main_input_name = "patch_embeds" + input_modalities = ("image",) + _input_embed_layer = "patch_embedding" + + def __init__(self, config: "UnlimitedOcrVisionEncoderConfig"): + super().__init__(config) + embed_dim = config.hidden_size + + self.embeddings = UnlimitedOcrVisionEmbeddings(config) + self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + self.encoder = UnlimitedOcrEncoder(config) + self.post_init() + + @can_return_tuple + @auto_docstring + def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: + r""" + patch_embeds (`torch.Tensor` of shape `(batch_size, hidden_size, grid_height, grid_width)`): + The SAM feature map used in place of the CLIP patch embeddings. + """ + hidden_states = self.embeddings(patch_embeds) + hidden_states = self.pre_layrnorm(hidden_states) + encoder_outputs: BaseModelOutput = self.encoder(inputs_embeds=hidden_states, **kwargs) + return BaseModelOutput( + last_hidden_state=encoder_outputs.last_hidden_state, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) class UnlimitedOcrVisionModel(UnlimitedOcrPreTrainedModel): @@ -978,17 +1176,18 @@ def __init__(self, config: UnlimitedOcrVisionConfig): @auto_docstring def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: sam_encoder_outputs = self.sam_encoder(pixel_values, **kwargs) - vision_encoder_outputs = self.vision_encoder(pixel_values, **kwargs) + sam_feature_map = sam_encoder_outputs.last_hidden_state - sam_hidden_state = sam_encoder_outputs.last_hidden_state + vision_encoder_outputs = self.vision_encoder(sam_feature_map, **kwargs) vision_encoder_hidden_state = vision_encoder_outputs.last_hidden_state - hidden_state = torch.cat([vision_encoder_hidden_state[:, 1:], sam_hidden_state], dim=1) - # TODO: How to pass hidden_states and attentions? + sam_hidden_state = sam_feature_map.flatten(2).transpose(1, 2) + hidden_state = torch.cat([vision_encoder_hidden_state[:, 1:], sam_hidden_state], dim=-1) + return BaseModelOutput( last_hidden_state=hidden_state, - # hidden_states=encoder_outputs.hidden_states, - # attentions=encoder_outputs.attentions, + hidden_states=vision_encoder_outputs.hidden_states, + attentions=vision_encoder_outputs.attentions, ) @@ -1442,7 +1641,6 @@ def get_image_features( image_spatial_crop (`torch.Tensor` of shape `(num_images, 2)`, *optional*): The local crop grid `(num_columns, num_rows)` per image. """ - # `torch.split` requires `list[int]`, not a Tensor, for per-image variable-length splitting. if isinstance(num_local_patches, torch.Tensor): num_local_patches = num_local_patches.tolist() @@ -1490,8 +1688,8 @@ def get_image_features( ) local_flat = local_grid.reshape(-1, hidden_size) # NOTE: local-then-global ordering matches the reference `forward` (the feature order the weights - # were trained on), NOT the token order built by the reference `infer`. Verify against a reference - # generation on a cropped image before trusting crop mode. + # were trained on), NOT the token order built by the reference `infer`. + # TODO: verify correctness all_features.append(torch.cat([local_flat, global_flat, view_separator], dim=0)) else: all_features.append(torch.cat([global_flat, view_separator], dim=0)) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 58d024898a39..3267c86bc9c6 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -23,8 +23,9 @@ from ...modeling_outputs import BaseModelOutput from ...processing_utils import Unpack from ...tokenization_utils_base import PreTokenizedInput, TextInput -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_int from ..clip.configuration_clip import CLIPVisionConfig +from ..clip.modeling_clip import CLIPVisionEmbeddings, CLIPVisionModel from ..deepseek_ocr2.configuration_deepseek_ocr2 import ( DeepseekOcr2Config, DeepseekOcr2TextConfig, @@ -57,7 +58,6 @@ DeepseekOcr2TextRMSNorm, DeepseekOcr2TextRotaryEmbedding, DeepseekOcr2VisionAttention, - DeepseekOcr2VisionEncoder, DeepseekOcr2VisionEncoderLayer, DeepseekOcr2VisionMLP, DeepseekOcr2VisionModel, @@ -319,8 +319,78 @@ class UnlimitedOcrVisionEncoderLayer(DeepseekOcr2VisionEncoderLayer): pass -class UnlimitedOcrVisionEncoder(DeepseekOcr2VisionEncoder): - pass +class UnlimitedOcrVisionEmbeddings(CLIPVisionEmbeddings): + def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: + num_patches = embeddings.shape[1] - 1 + position_embedding = self.position_embedding.weight.unsqueeze(0) + num_positions = position_embedding.shape[1] - 1 + + if not torch.jit.is_tracing() and num_patches == num_positions and height == width: + return self.position_embedding(self.position_ids) + + class_pos_embed = position_embedding[:, :1] + patch_pos_embed = position_embedding[:, 1:] + + dim = embeddings.shape[-1] + + new_height = height // self.patch_size + new_width = width // self.patch_size + + sqrt_num_positions = torch_int(num_positions**0.5) + patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) + patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) + + target_dtype = patch_pos_embed.dtype + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed.to(torch.float32), + size=(new_height, new_width), + mode="bicubic", + antialias=True, + align_corners=False, + ).to(target_dtype) + + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + + return torch.cat((class_pos_embed, patch_pos_embed), dim=1) + + def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: + r""" + patch_embeds (`torch.Tensor` of shape `(batch_size, hidden_size, grid_height, grid_width)`): + The SAM feature map, injected directly as the CLIP patch embeddings. The CLIP patch convolution + (`self.patch_embedding`) is intentionally bypassed. + """ + batch_size, _, grid_height, grid_width = patch_embeds.shape + patch_embeds = patch_embeds.flatten(2).transpose(1, 2) + class_embeds = self.class_embedding.expand(batch_size, 1, -1) + embeddings = torch.cat([class_embeds, patch_embeds], dim=1) + embeddings = embeddings + self.interpolate_pos_encoding( + embeddings, grid_height * self.patch_size, grid_width * self.patch_size + ) + return embeddings + + +class UnlimitedOcrVisionEncoder(CLIPVisionModel): + main_input_name = "patch_embeds" + + def __init__(self, config: "UnlimitedOcrVisionEncoderConfig"): + super().__init__(config) + del self.post_layernorm + + @can_return_tuple + @auto_docstring + def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: + r""" + patch_embeds (`torch.Tensor` of shape `(batch_size, hidden_size, grid_height, grid_width)`): + The SAM feature map used in place of the CLIP patch embeddings. + """ + hidden_states = self.embeddings(patch_embeds) + hidden_states = self.pre_layrnorm(hidden_states) + encoder_outputs: BaseModelOutput = self.encoder(inputs_embeds=hidden_states, **kwargs) + return BaseModelOutput( + last_hidden_state=encoder_outputs.last_hidden_state, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) class UnlimitedOcrVisionModel(DeepseekOcr2VisionModel): @@ -333,17 +403,18 @@ def __init__(self, config: UnlimitedOcrVisionConfig): @auto_docstring def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: sam_encoder_outputs = self.sam_encoder(pixel_values, **kwargs) - vision_encoder_outputs = self.vision_encoder(pixel_values, **kwargs) + sam_feature_map = sam_encoder_outputs.last_hidden_state - sam_hidden_state = sam_encoder_outputs.last_hidden_state + vision_encoder_outputs = self.vision_encoder(sam_feature_map, **kwargs) vision_encoder_hidden_state = vision_encoder_outputs.last_hidden_state - hidden_state = torch.cat([vision_encoder_hidden_state[:, 1:], sam_hidden_state], dim=1) - # TODO: How to pass hidden_states and attentions? + sam_hidden_state = sam_feature_map.flatten(2).transpose(1, 2) + hidden_state = torch.cat([vision_encoder_hidden_state[:, 1:], sam_hidden_state], dim=-1) + return BaseModelOutput( last_hidden_state=hidden_state, - # hidden_states=encoder_outputs.hidden_states, - # attentions=encoder_outputs.attentions, + hidden_states=vision_encoder_outputs.hidden_states, + attentions=vision_encoder_outputs.attentions, ) @@ -409,7 +480,6 @@ def get_image_features( image_spatial_crop (`torch.Tensor` of shape `(num_images, 2)`, *optional*): The local crop grid `(num_columns, num_rows)` per image. """ - # `torch.split` requires `list[int]`, not a Tensor, for per-image variable-length splitting. if isinstance(num_local_patches, torch.Tensor): num_local_patches = num_local_patches.tolist() @@ -457,8 +527,8 @@ def get_image_features( ) local_flat = local_grid.reshape(-1, hidden_size) # NOTE: local-then-global ordering matches the reference `forward` (the feature order the weights - # were trained on), NOT the token order built by the reference `infer`. Verify against a reference - # generation on a cropped image before trusting crop mode. + # were trained on), NOT the token order built by the reference `infer`. + # TODO: verify correctness all_features.append(torch.cat([local_flat, global_flat, view_separator], dim=0)) else: all_features.append(torch.cat([global_flat, view_separator], dim=0)) From ebf9b3b67192d759fe1d8cba12834e765297e7eb Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 14:07:38 +0200 Subject: [PATCH 004/195] add r-swa --- .../configuration_unlimited_ocr.py | 14 ++ .../unlimited_ocr/modeling_unlimited_ocr.py | 51 ++++++-- .../unlimited_ocr/modular_unlimited_ocr.py | 122 +++++++++++++++++- .../test_modeling_unlimited_ocr.py | 2 +- 4 files changed, 169 insertions(+), 20 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index 803efdb0cf38..8c61d7a377a3 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -146,6 +146,14 @@ class UnlimitedOcrTextConfig(PreTrainedConfig): Method for selecting top-k experts in MoE layers. mlp_layer_types (`list[str]`, *optional*): MLP type (`"dense"` or `"sparse"`) for each decoder layer, e.g. `["dense", "sparse", "sparse", ...]`. + layer_types (`list[str]`, *optional*): + Attention type for each decoder layer. Defaults to `"full_attention"` on every layer so the KV cache + retains all tokens; the sliding window (`sliding_window`) is applied as a mask over generated tokens + only, not by truncating the cache. + sliding_window (`int`, *optional*, defaults to 128): + If set, each token additionally attends only to the last `sliding_window` tokens. The image and prompt + tokens processed during prefill stay fully visible (they are never evicted); the window only applies + across generated tokens. Set to `None` for full causal attention. """ model_type = "unlimited_ocr_text" @@ -204,8 +212,14 @@ class UnlimitedOcrTextConfig(PreTrainedConfig): moe_intermediate_size: int = 1407 base_config_key = "text_config" mlp_layer_types: list[str] | None = None + layer_types: list[str] | None = None + sliding_window: int | None = 128 def __post_init__(self, **kwargs): + if self.layer_types is None: + # Full attention on every layer keeps the KV cache complete: the sliding window is realized as a + # mask over generated tokens, while the image/prompt prefill is always retained. + self.layer_types = ["full_attention"] * self.num_hidden_layers self.head_dim = self.hidden_size // self.num_attention_heads if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 9f193a2fc571..b8f2dd3a1cd0 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -37,7 +37,7 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...masking_utils import create_causal_mask +from ...masking_utils import and_masks, causal_mask_function, create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling @@ -60,6 +60,7 @@ UnlimitedOcrSamVisionConfig, UnlimitedOcrTextConfig, UnlimitedOcrVisionConfig, + UnlimitedOcrVisionEncoderConfig, ) @@ -1137,7 +1138,7 @@ class UnlimitedOcrVisionEncoder(UnlimitedOcrPreTrainedModel): input_modalities = ("image",) _input_embed_layer = "patch_embedding" - def __init__(self, config: "UnlimitedOcrVisionEncoderConfig"): + def __init__(self, config: UnlimitedOcrVisionEncoderConfig): super().__init__(config) embed_dim = config.hidden_size @@ -1557,7 +1558,7 @@ def forward( 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) + inputs_embeds = self.embed_tokens(input_ids) if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) @@ -1567,13 +1568,7 @@ def forward( 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, - ) + causal_mask = self._create_attention_mask(inputs_embeds, attention_mask, past_key_values, position_ids) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) @@ -1595,6 +1590,36 @@ def forward( past_key_values=past_key_values, ) + def _create_attention_mask(self, inputs_embeds, attention_mask, past_key_values, position_ids): + mask_kwargs = { + "config": self.config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + if self.config.sliding_window is None: + return create_causal_mask(**mask_kwargs) + + # The sliding window only spans the generated tokens: the prefill (image + prompt) stays fully + # visible. We track the prefill length on the cache so it is stable across decode steps, then keep + # those positions attendable on top of the sliding-window-causal mask (while preserving causality). + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + if past_seen_tokens == 0: + prefill_length = inputs_embeds.shape[1] + if past_key_values is not None: + past_key_values.prefill_length = prefill_length + else: + prefill_length = getattr(past_key_values, "prefill_length", past_seen_tokens) + + def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): + return kv_idx < prefill_length + + return create_sliding_window_causal_mask( + **mask_kwargs, + or_mask_function=and_masks(prefill_overlay, causal_mask_function), + ) + @auto_docstring( custom_intro=""" @@ -1609,16 +1634,14 @@ def __init__(self, config: UnlimitedOcrConfig): self.vision_tower = UnlimitedOcrVisionModel(config.vision_config) self.multi_modal_projector = nn.Linear( - config.vision_config.sam_config.hidden_size + config.vision_config.encoder_config.hidden_size, + config.vision_config.sam_config.downsample_channels[-1] + config.vision_config.encoder_config.hidden_size, config.text_config.hidden_size, ) self.vocab_size = config.text_config.vocab_size self.language_model = UnlimitedOcrTextModel(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.view_separator = nn.Parameter(torch.randn(config.text_config.hidden_size) * embed_std) embed_std = 1 / torch.sqrt(torch.tensor(config.text_config.hidden_size, dtype=torch.float32)) self.image_newline = nn.Parameter(torch.randn(config.text_config.hidden_size) * embed_std) self.post_init() diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 3267c86bc9c6..e27c06253dc2 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -17,13 +17,21 @@ from huggingface_hub.dataclasses import strict from torch import nn -from ...cache_utils import Cache +from ...cache_utils import Cache, DynamicCache from ...feature_extraction_utils import BatchFeature from ...image_utils import ImageInput -from ...modeling_outputs import BaseModelOutput +from ...masking_utils import ( + and_masks, + causal_mask_function, + create_causal_mask, + create_sliding_window_causal_mask, +) +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast from ...processing_utils import Unpack from ...tokenization_utils_base import PreTokenizedInput, TextInput from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_int +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs from ..clip.configuration_clip import CLIPVisionConfig from ..clip.modeling_clip import CLIPVisionEmbeddings, CLIPVisionModel from ..deepseek_ocr2.configuration_deepseek_ocr2 import ( @@ -233,8 +241,34 @@ class UnlimitedOcrVisionConfig(DeepseekOcr2VisionConfig): @auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict class UnlimitedOcrTextConfig(DeepseekOcr2TextConfig): + r""" + n_group (`int`, *optional*): + Number of groups for grouped top-k expert routing. + topk_method (`str`, *optional*, defaults to `"greedy"`): + Method for selecting top-k experts in MoE layers. + mlp_layer_types (`list[str]`, *optional*): + MLP type (`"dense"` or `"sparse"`) for each decoder layer, e.g. `["dense", "sparse", "sparse", ...]`. + layer_types (`list[str]`, *optional*): + Attention type for each decoder layer. Defaults to `"full_attention"` on every layer so the KV cache + retains all tokens; the sliding window (`sliding_window`) is applied as a mask over generated tokens + only, not by truncating the cache. + sliding_window (`int`, *optional*, defaults to 128): + If set, each token additionally attends only to the last `sliding_window` tokens. The image and prompt + tokens processed during prefill stay fully visible (they are never evicted); the window only applies + across generated tokens. Set to `None` for full causal attention. + """ + model_type = "unlimited_ocr_text" base_config_key = "text_config" + layer_types: list[str] | None = None + sliding_window: int | None = 128 + + def __post_init__(self, **kwargs): + if self.layer_types is None: + # Full attention on every layer keeps the KV cache complete: the sliding window is realized as a + # mask over generated tokens, while the image/prompt prefill is always retained. + self.layer_types = ["full_attention"] * self.num_hidden_layers + super().__post_init__(**kwargs) @auto_docstring(checkpoint="baidu/Unlimited-OCR") @@ -372,7 +406,7 @@ def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: class UnlimitedOcrVisionEncoder(CLIPVisionModel): main_input_name = "patch_embeds" - def __init__(self, config: "UnlimitedOcrVisionEncoderConfig"): + def __init__(self, config: UnlimitedOcrVisionEncoderConfig): super().__init__(config) del self.post_layernorm @@ -451,18 +485,96 @@ class UnlimitedOcrTextPreTrainedModel(DeepseekOcr2TextPreTrainedModel): class UnlimitedOcrTextModel(DeepseekOcr2TextModel): - pass + def _create_attention_mask(self, inputs_embeds, attention_mask, past_key_values, position_ids): + mask_kwargs = { + "config": self.config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + if self.config.sliding_window is None: + return create_causal_mask(**mask_kwargs) + + # The sliding window only spans the generated tokens: the prefill (image + prompt) stays fully + # visible. We track the prefill length on the cache so it is stable across decode steps, then keep + # those positions attendable on top of the sliding-window-causal mask (while preserving causality). + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + if past_seen_tokens == 0: + prefill_length = inputs_embeds.shape[1] + if past_key_values is not None: + past_key_values.prefill_length = prefill_length + else: + prefill_length = getattr(past_key_values, "prefill_length", past_seen_tokens) + + def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): + return kv_idx < prefill_length + + return create_sliding_window_causal_mask( + **mask_kwargs, + or_mask_function=and_masks(prefill_overlay, causal_mask_function), + ) + + @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 = 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 = self._create_attention_mask(inputs_embeds, attention_mask, past_key_values, 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, + ) class UnlimitedOcrModel(DeepseekOcr2Model): def __init__(self, config: UnlimitedOcrConfig): super().__init__(config) self.multi_modal_projector = nn.Linear( - config.vision_config.sam_config.hidden_size + config.vision_config.encoder_config.hidden_size, + config.vision_config.sam_config.downsample_channels[-1] + config.vision_config.encoder_config.hidden_size, config.text_config.hidden_size, ) embed_std = 1 / torch.sqrt(torch.tensor(config.text_config.hidden_size, dtype=torch.float32)) self.image_newline = nn.Parameter(torch.randn(config.text_config.hidden_size) * embed_std) + self.view_separator = nn.Parameter(torch.randn(config.text_config.hidden_size) * embed_std) def get_image_features( self, diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 7cea1f9a196b..91e95eeb104e 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -34,7 +34,7 @@ UnlimitedOcrForConditionalGeneration, UnlimitedOcrModel, ) - from transformers.models.deepseek_ocr2.configuration_deepseek_ocr2 import ( + from transformers.models.unlimited_ocr.configuration_unlimited_ocr import ( UnlimitedOcrTextConfig, UnlimitedOcrVisionConfig, ) From a531dd839588d29a0fdf10cd162f1dea1c1a31a7 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 15:22:17 +0200 Subject: [PATCH 005/195] fix tests --- .../unlimited_ocr/modeling_unlimited_ocr.py | 206 ++++++++++-------- .../unlimited_ocr/modular_unlimited_ocr.py | 56 +++-- .../test_modeling_unlimited_ocr.py | 12 +- 3 files changed, 163 insertions(+), 111 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index b8f2dd3a1cd0..30a49dc1f81b 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -149,8 +149,10 @@ class UnlimitedOcrPreTrainedModel(PreTrainedModel): # SAM uses rel-pos bias, incompatible with flash attention. _supports_flash_attn = False _supports_sdpa = True - - _can_compile_fullgraph = True + # The text model keeps the image/prompt prefill fully visible by tracking its length on the cache and + # building the attention mask from it. This stateful, data-dependent masking cannot be traced into a single + # graph, so full-graph compilation is unsupported. + _can_compile_fullgraph = False _supports_flex_attn = True _supports_attention_backend = True @@ -167,6 +169,15 @@ def _init_weights(self, module): elif isinstance(module, UnlimitedOcrModel): embed_std = 1 / math.sqrt(self.config.text_config.hidden_size) init.normal_(module.view_separator, mean=0.0, std=embed_std) + if isinstance(module, UnlimitedOcrModel): + embed_std = 1 / math.sqrt(self.config.text_config.hidden_size) + init.normal_(module.image_newline, mean=0.0, std=embed_std) + elif isinstance(module, UnlimitedOcrVisionEmbeddings): + factor = module.config.initializer_factor + init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor) + init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor) + init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor) + init.copy_(module.position_ids, torch.arange(module.num_positions).expand((1, -1))) class UnlimitedOcrSamVisionAttention(nn.Module): @@ -907,90 +918,10 @@ def forward( return hidden_states -class UnlimitedOcrVisionEmbeddings(nn.Module): - def __init__(self, config: UnlimitedOcrVisionConfig): - super().__init__() - self.config = config - self.embed_dim = config.hidden_size - self.image_size = config.image_size - self.patch_size = config.patch_size - - self.class_embedding = nn.Parameter(torch.randn(self.embed_dim)) - - self.patch_embedding = nn.Conv2d( - in_channels=config.num_channels, - out_channels=self.embed_dim, - kernel_size=self.patch_size, - stride=self.patch_size, - bias=False, - ) - - self.num_patches = (self.image_size // self.patch_size) ** 2 - self.num_positions = self.num_patches + 1 - self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) - self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False) - - def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: - """ - This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution - images. This method is also adapted to support torch.jit tracing. - - Adapted from: - - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and - - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 - """ - num_patches = embeddings.shape[1] - 1 - position_embedding = self.position_embedding.weight.unsqueeze(0) - num_positions = position_embedding.shape[1] - 1 - - if not torch.jit.is_tracing() and num_patches == num_positions and height == width: - return self.position_embedding(self.position_ids) - - class_pos_embed = position_embedding[:, :1] - patch_pos_embed = position_embedding[:, 1:] - - dim = embeddings.shape[-1] - - new_height = height // self.patch_size - new_width = width // self.patch_size - - sqrt_num_positions = torch_int(num_positions**0.5) - patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) - patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) - - target_dtype = patch_pos_embed.dtype - patch_pos_embed = nn.functional.interpolate( - patch_pos_embed.to(torch.float32), - size=(new_height, new_width), - mode="bicubic", - antialias=True, - align_corners=False, - ).to(target_dtype) - - patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) - - return torch.cat((class_pos_embed, patch_pos_embed), dim=1) - - def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: - r""" - patch_embeds (`torch.Tensor` of shape `(batch_size, hidden_size, grid_height, grid_width)`): - The SAM feature map, injected directly as the CLIP patch embeddings. The CLIP patch convolution - (`self.patch_embedding`) is intentionally bypassed. - """ - batch_size, _, grid_height, grid_width = patch_embeds.shape - patch_embeds = patch_embeds.flatten(2).transpose(1, 2) - class_embeds = self.class_embedding.expand(batch_size, 1, -1) - embeddings = torch.cat([class_embeds, patch_embeds], dim=1) - embeddings = embeddings + self.interpolate_pos_encoding( - embeddings, grid_height * self.patch_size, grid_width * self.patch_size - ) - return embeddings - - class UnlimitedOcrAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" - def __init__(self, config: UnlimitedOcrVisionConfig | UnlimitedOcrTextConfig): + def __init__(self, config: UnlimitedOcrVisionEncoderConfig): super().__init__() self.config = config self.embed_dim = config.hidden_size @@ -1004,6 +935,9 @@ def __init__(self, config: UnlimitedOcrVisionConfig | UnlimitedOcrTextConfig): self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) + # The shared `eager_attention_forward` calls `repeat_kv(..., num_key_value_groups)`; CLIP attention is + # plain multi-head attention, so the group count is 1 and `repeat_kv` becomes a no-op. + self.num_key_value_groups = 1 def forward( self, @@ -1093,6 +1027,86 @@ def forward( return hidden_states +class UnlimitedOcrVisionEmbeddings(nn.Module): + def __init__(self, config: UnlimitedOcrVisionConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size + + self.class_embedding = nn.Parameter(torch.randn(self.embed_dim)) + + self.patch_embedding = nn.Conv2d( + in_channels=config.num_channels, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + bias=False, + ) + + self.num_patches = (self.image_size // self.patch_size) ** 2 + self.num_positions = self.num_patches + 1 + self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False) + + def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: + """ + This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution + images. This method is also adapted to support torch.jit tracing. + + Adapted from: + - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and + - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 + """ + num_patches = embeddings.shape[1] - 1 + position_embedding = self.position_embedding.weight.unsqueeze(0) + num_positions = position_embedding.shape[1] - 1 + + if not torch.jit.is_tracing() and num_patches == num_positions and height == width: + return self.position_embedding(self.position_ids) + + class_pos_embed = position_embedding[:, :1] + patch_pos_embed = position_embedding[:, 1:] + + dim = embeddings.shape[-1] + + new_height = height // self.patch_size + new_width = width // self.patch_size + + sqrt_num_positions = torch_int(num_positions**0.5) + patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) + patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) + + target_dtype = patch_pos_embed.dtype + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed.to(torch.float32), + size=(new_height, new_width), + mode="bicubic", + antialias=True, + align_corners=False, + ).to(target_dtype) + + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + + return torch.cat((class_pos_embed, patch_pos_embed), dim=1) + + def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: + r""" + patch_embeds (`torch.Tensor` of shape `(batch_size, hidden_size, grid_height, grid_width)`): + The SAM feature map, injected directly as the CLIP patch embeddings. The CLIP patch convolution + (`self.patch_embedding`) is intentionally bypassed. + """ + batch_size, _, grid_height, grid_width = patch_embeds.shape + patch_embeds = patch_embeds.flatten(2).transpose(1, 2) + class_embeds = self.class_embedding.expand(batch_size, 1, -1) + embeddings = torch.cat([class_embeds, patch_embeds], dim=1) + embeddings = embeddings + self.interpolate_pos_encoding( + embeddings, grid_height * self.patch_size, grid_width * self.patch_size + ) + return embeddings + + class UnlimitedOcrEncoder(nn.Module): """ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a @@ -1137,6 +1151,10 @@ class UnlimitedOcrVisionEncoder(UnlimitedOcrPreTrainedModel): main_input_name = "patch_embeds" input_modalities = ("image",) _input_embed_layer = "patch_embedding" + _can_record_outputs = { + "hidden_states": UnlimitedOcrEncoderLayer, + "attentions": UnlimitedOcrAttention, + } def __init__(self, config: UnlimitedOcrVisionEncoderConfig): super().__init__(config) @@ -1148,6 +1166,7 @@ def __init__(self, config: UnlimitedOcrVisionEncoderConfig): self.post_init() @can_return_tuple + @capture_outputs @auto_docstring def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: r""" @@ -1157,11 +1176,7 @@ def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwarg hidden_states = self.embeddings(patch_embeds) hidden_states = self.pre_layrnorm(hidden_states) encoder_outputs: BaseModelOutput = self.encoder(inputs_embeds=hidden_states, **kwargs) - return BaseModelOutput( - last_hidden_state=encoder_outputs.last_hidden_state, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - ) + return BaseModelOutput(last_hidden_state=encoder_outputs.last_hidden_state) class UnlimitedOcrVisionModel(UnlimitedOcrPreTrainedModel): @@ -1641,9 +1656,10 @@ def __init__(self, config: UnlimitedOcrConfig): self.vocab_size = config.text_config.vocab_size self.language_model = UnlimitedOcrTextModel(config.text_config) - self.view_separator = nn.Parameter(torch.randn(config.text_config.hidden_size) * embed_std) - embed_std = 1 / torch.sqrt(torch.tensor(config.text_config.hidden_size, dtype=torch.float32)) - self.image_newline = nn.Parameter(torch.randn(config.text_config.hidden_size) * embed_std) + + # Learnable separator between local and global views (initialized in `_init_weights`). + self.view_separator = nn.Parameter(torch.empty(config.text_config.hidden_size)) + self.image_newline = nn.Parameter(torch.empty(config.text_config.hidden_size)) self.post_init() @can_return_tuple @@ -1926,16 +1942,18 @@ def prepare_inputs_for_generation( input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, - pixel_values=pixel_values, - pixel_values_local=pixel_values_local, - num_local_patches=num_local_patches, attention_mask=attention_mask, logits_to_keep=logits_to_keep, is_first_iteration=is_first_iteration, **kwargs, ) + # Image inputs are only needed during prefill (or when the cache is disabled); once the image tokens + # have been embedded they must be dropped so later decode steps don't reprocess the pixel values. 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 model_inputs["image_spatial_crop"] = image_spatial_crop return model_inputs diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index e27c06253dc2..9942f50415c0 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -17,6 +17,7 @@ from huggingface_hub.dataclasses import strict from torch import nn +from ... import initialization as init from ...cache_utils import Cache, DynamicCache from ...feature_extraction_utils import BatchFeature from ...image_utils import ImageInput @@ -33,7 +34,7 @@ from ...utils.generic import merge_with_config_defaults from ...utils.output_capturing import capture_outputs from ..clip.configuration_clip import CLIPVisionConfig -from ..clip.modeling_clip import CLIPVisionEmbeddings, CLIPVisionModel +from ..clip.modeling_clip import CLIPAttention, CLIPEncoderLayer, CLIPVisionEmbeddings, CLIPVisionModel from ..deepseek_ocr2.configuration_deepseek_ocr2 import ( DeepseekOcr2Config, DeepseekOcr2TextConfig, @@ -294,7 +295,23 @@ class UnlimitedOcrCausalLMOutputWithPast(DeepseekOcr2CausalLMOutputWithPast): class UnlimitedOcrPreTrainedModel(DeepseekOcr2PreTrainedModel): - pass + # The text model keeps the image/prompt prefill fully visible by tracking its length on the cache and + # building the attention mask from it. This stateful, data-dependent masking cannot be traced into a single + # graph, so full-graph compilation is unsupported. + _can_compile_fullgraph = False + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, UnlimitedOcrModel): + embed_std = 1 / math.sqrt(self.config.text_config.hidden_size) + init.normal_(module.image_newline, mean=0.0, std=embed_std) + elif isinstance(module, UnlimitedOcrVisionEmbeddings): + factor = module.config.initializer_factor + init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor) + init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor) + init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor) + init.copy_(module.position_ids, torch.arange(module.num_positions).expand((1, -1))) class UnlimitedOcrSamVisionAttention(DeepseekOcr2SamVisionAttention): @@ -353,6 +370,18 @@ class UnlimitedOcrVisionEncoderLayer(DeepseekOcr2VisionEncoderLayer): pass +class UnlimitedOcrAttention(CLIPAttention): + def __init__(self, config: UnlimitedOcrVisionEncoderConfig): + super().__init__(config) + # The shared `eager_attention_forward` calls `repeat_kv(..., num_key_value_groups)`; CLIP attention is + # plain multi-head attention, so the group count is 1 and `repeat_kv` becomes a no-op. + self.num_key_value_groups = 1 + + +class UnlimitedOcrEncoderLayer(CLIPEncoderLayer): + pass + + class UnlimitedOcrVisionEmbeddings(CLIPVisionEmbeddings): def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: num_patches = embeddings.shape[1] - 1 @@ -405,12 +434,17 @@ def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: class UnlimitedOcrVisionEncoder(CLIPVisionModel): main_input_name = "patch_embeds" + _can_record_outputs = { + "hidden_states": UnlimitedOcrEncoderLayer, + "attentions": UnlimitedOcrAttention, + } def __init__(self, config: UnlimitedOcrVisionEncoderConfig): super().__init__(config) del self.post_layernorm @can_return_tuple + @capture_outputs @auto_docstring def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: r""" @@ -420,11 +454,7 @@ def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwarg hidden_states = self.embeddings(patch_embeds) hidden_states = self.pre_layrnorm(hidden_states) encoder_outputs: BaseModelOutput = self.encoder(inputs_embeds=hidden_states, **kwargs) - return BaseModelOutput( - last_hidden_state=encoder_outputs.last_hidden_state, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - ) + return BaseModelOutput(last_hidden_state=encoder_outputs.last_hidden_state) class UnlimitedOcrVisionModel(DeepseekOcr2VisionModel): @@ -572,9 +602,7 @@ def __init__(self, config: UnlimitedOcrConfig): config.vision_config.sam_config.downsample_channels[-1] + config.vision_config.encoder_config.hidden_size, config.text_config.hidden_size, ) - embed_std = 1 / torch.sqrt(torch.tensor(config.text_config.hidden_size, dtype=torch.float32)) - self.image_newline = nn.Parameter(torch.randn(config.text_config.hidden_size) * embed_std) - self.view_separator = nn.Parameter(torch.randn(config.text_config.hidden_size) * embed_std) + self.image_newline = nn.Parameter(torch.empty(config.text_config.hidden_size)) def get_image_features( self, @@ -792,16 +820,18 @@ def prepare_inputs_for_generation( input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, - pixel_values=pixel_values, - pixel_values_local=pixel_values_local, - num_local_patches=num_local_patches, attention_mask=attention_mask, logits_to_keep=logits_to_keep, is_first_iteration=is_first_iteration, **kwargs, ) + # Image inputs are only needed during prefill (or when the cache is disabled); once the image tokens + # have been embedded they must be dropped so later decode steps don't reprocess the pixel values. 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 model_inputs["image_spatial_crop"] = image_spatial_crop return model_inputs diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 91e95eeb104e..b20ef3e5a4fb 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -51,8 +51,7 @@ class UnlimitedOcrVisionText2TextModelTester(VLMModelTester): vision_config_class = UnlimitedOcrVisionConfig def __init__(self, parent, **kwargs): - # VisionModel always selects query_768_resolution (144 tokens) for small images + 1 separator - kwargs.setdefault("num_image_tokens", 145) + kwargs.setdefault("num_image_tokens", 7) kwargs.setdefault("image_token_id", 1) kwargs.setdefault("image_size", 16) kwargs.setdefault("hidden_size", 128) @@ -83,6 +82,7 @@ def __init__(self, parent, **kwargs): "patch_size": 2, "hidden_act": "gelu", "mlp_ratio": 4.0, + "mlp_dim": 128, "window_size": 4, "global_attn_indexes": [1], "downsample_channels": [32, 64], @@ -144,8 +144,12 @@ def test_disk_offload_safetensors(self): def _image_features_prepare_config_and_inputs(self): config, inputs_dict = super()._image_features_prepare_config_and_inputs() - # test_get_image_features_output expects vision_config.hidden_size, but ours is in encoder_config. - config.vision_config.hidden_size = config.vision_config.encoder_config.hidden_size + # `get_image_features` returns the concatenation of the SAM feature map and the CLIP encoder output as + # `last_hidden_state`, so its hidden size is the sum of the two. `vision_config` has no `hidden_size` of + # its own, so set it here for `test_get_image_features_output`. + config.vision_config.hidden_size = ( + config.vision_config.sam_config.downsample_channels[-1] + config.vision_config.encoder_config.hidden_size + ) return config, inputs_dict From 8cd0c325e23f5e0e2abd3029584a3ebeb314379b Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 13:23:49 +0000 Subject: [PATCH 006/195] update expected values --- .../test_modeling_unlimited_ocr.py | 26 +++++++------------ 1 file changed, 10 insertions(+), 16 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 91e95eeb104e..fd8b20c660b4 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -151,7 +151,7 @@ def _image_features_prepare_config_and_inputs(self): @require_torch class UnlimitedOcrIntegrationTest(unittest.TestCase): - model_id = "deepseek-community/DeepSeek-OCR-2" + model_id = "baidu/Unlimited-OCR" def setUp(self): self.processor = AutoProcessor.from_pretrained(self.model_id) @@ -160,7 +160,7 @@ def tearDown(self): cleanup(torch_device, gc_collect=True) @slow - def test_small_model_integration_test_free_ocr(self): + def test_small_model_integration_test_document_parsing(self): model = UnlimitedOcrForConditionalGeneration.from_pretrained( self.model_id, torch_dtype=torch.bfloat16, device_map=torch_device ) @@ -169,21 +169,20 @@ def test_small_model_integration_test_free_ocr(self): "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" ) ) - inputs = self.processor(images=image, text="\nFree OCR.", return_tensors="pt").to( + inputs = self.processor(images=image, text="\ndocument parsing.", return_tensors="pt").to( model.device, dtype=torch.bfloat16 ) generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True) EXPECTED_DECODED_TEXT = Expectations( { - ("cuda", None): "R&D QUALITY IMPROVEMENT SUGGESTION/SOLUTION FORM\n\nName/", - ("xpu", 5): "R&D QUALITY IMPROVEMENT SUGGESTION/SOLUTION FORM\n\nName/", + ("cuda", None): "image [383, 87, 497, 171]\ntitle [333", } ).get_expectation() # fmt: skip self.assertEqual(decoded, EXPECTED_DECODED_TEXT) @slow - def test_small_model_integration_test_grounding_markdown(self): + def test_small_model_integration_test_document_parsing_grounding(self): model = UnlimitedOcrForConditionalGeneration.from_pretrained( self.model_id, torch_dtype=torch.bfloat16, device_map=torch_device ) @@ -194,15 +193,14 @@ def test_small_model_integration_test_grounding_markdown(self): ) inputs = self.processor( images=image, - text="\n<|grounding|>Convert the document to markdown.", + text="\ndocument parsing.", return_tensors="pt", ).to(model.device, dtype=torch.bfloat16) generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=False) EXPECTED_DECODED_TEXT = Expectations( { - ("cuda", None): "<|ref|>title<|/ref|><|det|>[[330, 198, 559, 230]]<|/det|>\n# R", - ("xpu", 5): "<|ref|>title<|/ref|><|det|>[[330, 198, 558, 230]]<|/det|>\n# R", + ("cuda", None): "<|det|>image [383, 87, 497, 171]<|/det|>\n<|det|>title [333", } ).get_expectation() # fmt: skip self.assertEqual(decoded, EXPECTED_DECODED_TEXT) @@ -224,7 +222,7 @@ def test_small_model_integration_test_batched(self): ) inputs = self.processor( images=[image1, image2], - text=["\nFree OCR.", "\nFree OCR."], + text=["\ndocument parsing.", "\ndocument parsing."], return_tensors="pt", padding=True, ).to(model.device, dtype=torch.bfloat16) @@ -235,12 +233,8 @@ def test_small_model_integration_test_batched(self): EXPECTED_DECODED_TEXT = Expectations( { ("cuda", None): [ - "R&D QUALITY IMPROVEMENT SUGGESTION/SOLUTION FORM\n\nName/", - "# Reducing the number of images\n\nIt is also believed that the performance of a website is a critical", - ], - ("xpu", 5): [ - "R&D QUALITY IMPROVEMENT SUGGESTION/SOLUTION FORM\n\nName/", - "# Reducing the number of images\n\nIt is also believed that the performance of a website is a critical", + "image [383, 87, 497, 171]\ntitle [333", + "header [53, 23, 365, 41]Advanced Template and Styl", ], } ).get_expectation() # fmt: skip From b6a226e8770c5460f5d8d6de7454b55d22e62867 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 15:47:13 +0200 Subject: [PATCH 007/195] enable _can_compile_fullgraph --- .../unlimited_ocr/modeling_unlimited_ocr.py | 17 +++++------------ .../unlimited_ocr/modular_unlimited_ocr.py | 16 +++------------- 2 files changed, 8 insertions(+), 25 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 30a49dc1f81b..4574979c19f7 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -149,10 +149,8 @@ class UnlimitedOcrPreTrainedModel(PreTrainedModel): # SAM uses rel-pos bias, incompatible with flash attention. _supports_flash_attn = False _supports_sdpa = True - # The text model keeps the image/prompt prefill fully visible by tracking its length on the cache and - # building the attention mask from it. This stateful, data-dependent masking cannot be traced into a single - # graph, so full-graph compilation is unsupported. - _can_compile_fullgraph = False + + _can_compile_fullgraph = True _supports_flex_attn = True _supports_attention_backend = True @@ -1616,16 +1614,11 @@ def _create_attention_mask(self, inputs_embeds, attention_mask, past_key_values, if self.config.sliding_window is None: return create_causal_mask(**mask_kwargs) - # The sliding window only spans the generated tokens: the prefill (image + prompt) stays fully - # visible. We track the prefill length on the cache so it is stable across decode steps, then keep - # those positions attendable on top of the sliding-window-causal mask (while preserving causality). - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - if past_seen_tokens == 0: - prefill_length = inputs_embeds.shape[1] + prefill_length = getattr(past_key_values, "prefill_length", None) + if prefill_length is None: + prefill_length = torch.tensor(inputs_embeds.shape[1], device=inputs_embeds.device) if past_key_values is not None: past_key_values.prefill_length = prefill_length - else: - prefill_length = getattr(past_key_values, "prefill_length", past_seen_tokens) def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): return kv_idx < prefill_length diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 9942f50415c0..ab54aad0e1ff 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -295,11 +295,6 @@ class UnlimitedOcrCausalLMOutputWithPast(DeepseekOcr2CausalLMOutputWithPast): class UnlimitedOcrPreTrainedModel(DeepseekOcr2PreTrainedModel): - # The text model keeps the image/prompt prefill fully visible by tracking its length on the cache and - # building the attention mask from it. This stateful, data-dependent masking cannot be traced into a single - # graph, so full-graph compilation is unsupported. - _can_compile_fullgraph = False - @torch.no_grad() def _init_weights(self, module): super()._init_weights(module) @@ -526,16 +521,11 @@ def _create_attention_mask(self, inputs_embeds, attention_mask, past_key_values, if self.config.sliding_window is None: return create_causal_mask(**mask_kwargs) - # The sliding window only spans the generated tokens: the prefill (image + prompt) stays fully - # visible. We track the prefill length on the cache so it is stable across decode steps, then keep - # those positions attendable on top of the sliding-window-causal mask (while preserving causality). - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - if past_seen_tokens == 0: - prefill_length = inputs_embeds.shape[1] + prefill_length = getattr(past_key_values, "prefill_length", None) + if prefill_length is None: + prefill_length = torch.tensor(inputs_embeds.shape[1], device=inputs_embeds.device) if past_key_values is not None: past_key_values.prefill_length = prefill_length - else: - prefill_length = getattr(past_key_values, "prefill_length", past_seen_tokens) def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): return kv_idx < prefill_length From e2d8d50f00b7bbf3d8ccbdb3ff1b06f10318e9fa Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 16:05:26 +0200 Subject: [PATCH 008/195] fix downsample outputs --- .../configuration_unlimited_ocr.py | 12 ++++++++- .../unlimited_ocr/modular_unlimited_ocr.py | 25 +++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index 8c61d7a377a3..dd41a5170eaa 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -19,7 +19,7 @@ # limitations under the License. from huggingface_hub.dataclasses import strict -from ...configuration_utils import PreTrainedConfig +from ...configuration_utils import PreTrainedConfig, PretrainedConfig from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -40,6 +40,8 @@ class UnlimitedOcrSamVisionConfig(PreTrainedConfig): The indexes of the global attention layers. mlp_dim (`int`, *optional*, defaults to 3072): The dimensionality of the MLP layer in the Transformer encoder. + downsample_channels (`list[int]`, *optional*, defaults to `(512, 896)`): + The channel dimensions for the multi-scale downsampling neck layers. """ base_config_key = "sam_config" @@ -60,8 +62,16 @@ class UnlimitedOcrSamVisionConfig(PreTrainedConfig): window_size: int = 14 global_attn_indexes: list[int] | tuple[int, ...] = (2, 5, 8, 11) mlp_dim: int = 3072 + model_type = "unlimited_ocr_sam_vision_model" + downsample_channels: list[int] | tuple[int, ...] | None = None + + def __post_init__(self, **kwargs): + if self.downsample_channels is None: + self.downsample_channels = [512, 896] + return PretrainedConfig.__post_init__(self, **kwargs) + @auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index ab54aad0e1ff..733584a968dc 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -19,6 +19,7 @@ from ... import initialization as init from ...cache_utils import Cache, DynamicCache +from ...configuration_utils import PretrainedConfig from ...feature_extraction_utils import BatchFeature from ...image_utils import ImageInput from ...masking_utils import ( @@ -207,9 +208,33 @@ def __call__( @auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict class UnlimitedOcrSamVisionConfig(GotOcr2VisionConfig): + r""" + output_channels (`int`, *optional*, defaults to 256): + Dimensionality of the output channels in the Patch Encoder. + use_abs_pos (`bool`, *optional*, defaults to `True`): + Whether to use absolute position embedding. + use_rel_pos (`bool`, *optional*, defaults to `True`): + Whether to use relative position embedding. + window_size (`int`, *optional*, defaults to 14): + Window size for relative position. + global_attn_indexes (`list[int]`, *optional*, defaults to `[2, 5, 8, 11]`): + The indexes of the global attention layers. + mlp_dim (`int`, *optional*, defaults to 3072): + The dimensionality of the MLP layer in the Transformer encoder. + downsample_channels (`list[int]`, *optional*, defaults to `(512, 896)`): + The channel dimensions for the multi-scale downsampling neck layers. + """ + model_type = "unlimited_ocr_sam_vision_model" base_config_key = "sam_config" + downsample_channels: list[int] | tuple[int, ...] | None = None + + def __post_init__(self, **kwargs): + if self.downsample_channels is None: + self.downsample_channels = [512, 896] + return PretrainedConfig.__post_init__(self, **kwargs) + @auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict From e5f2c46be908e7193f5cdaeae65a8a576ca89a1d Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 14:09:51 +0000 Subject: [PATCH 009/195] use temp repo for tests --- tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index ea0d06727178..d72300ca015b 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -155,7 +155,7 @@ def _image_features_prepare_config_and_inputs(self): @require_torch class UnlimitedOcrIntegrationTest(unittest.TestCase): - model_id = "baidu/Unlimited-OCR" + model_id = "guarin/Unlimited-OCR" def setUp(self): self.processor = AutoProcessor.from_pretrained(self.model_id) From 0e473433a28db52b362f2d3f067ff9253d64d2d4 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 16:28:25 +0200 Subject: [PATCH 010/195] fix mlp_layer_types None --- .../models/unlimited_ocr/configuration_unlimited_ocr.py | 3 +++ src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index dd41a5170eaa..f48e1d47d493 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -230,6 +230,9 @@ def __post_init__(self, **kwargs): # Full attention on every layer keeps the KV cache complete: the sliding window is realized as a # mask over generated tokens, while the image/prompt prefill is always retained. self.layer_types = ["full_attention"] * self.num_hidden_layers + if self.mlp_layer_types is None: + # TODO: double check + self.mlp_layer_types = ["dense"] * self.num_hidden_layers self.head_dim = self.hidden_size // self.num_attention_heads if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 733584a968dc..602f34f8aaa1 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -294,6 +294,9 @@ def __post_init__(self, **kwargs): # Full attention on every layer keeps the KV cache complete: the sliding window is realized as a # mask over generated tokens, while the image/prompt prefill is always retained. self.layer_types = ["full_attention"] * self.num_hidden_layers + if self.mlp_layer_types is None: + # TODO: double check + self.mlp_layer_types = ["dense"] * self.num_hidden_layers super().__post_init__(**kwargs) From 38f3199d6255a6e920558ccc1e5b2321972db663 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 15:33:08 +0000 Subject: [PATCH 011/195] add partial conversion mapping --- src/transformers/conversion_mapping.py | 60 ++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index a0cff222273a..aa936113b02d 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -1380,6 +1380,66 @@ def _build_checkpoint_conversion_mapping(): ), WeightRenaming(r"decode_head\.conv_matting\.", r"decode_head.predictor."), ], + "UnlimitedOcrModel": [ + WeightRenaming(r"model\.sam_model", "vision_tower.sam_encoder"), + WeightRenaming(r"model\.vision_model", "vision_tower.vision_encoder"), + WeightRenaming(r"model\.view_seperator", "view_separator"), + WeightRenaming(r"model\.image_newline", "image_newline"), + WeightRenaming(r"model\.projector\.layers\.(weight|bias)", r"multi_modal_projector.\1"), + WeightRenaming(r"model\.embed_tokens\.", r"language_model.embed_tokens."), + WeightRenaming(r"model\.layers\.", r"language_model.layers."), + WeightRenaming(r"model\.norm\.", r"language_model.norm."), + ], + "UnlimitedOcrSamVisionEncoder": [ + WeightRenaming(r"blocks\.(\d+)\.norm1\.", r"layers.\1.layer_norm1."), + WeightRenaming(r"blocks\.(\d+)\.norm2\.", r"layers.\1.layer_norm2."), + WeightRenaming(r"blocks\.(\d+)\.attn\.", r"layers.\1.attn."), + WeightRenaming(r"blocks\.(\d+)\.mlp\.", r"layers.\1.mlp."), + WeightRenaming(r"patch_embed\.proj\.", "patch_embed.projection."), + WeightRenaming(r"pos_embed", "pos_embed"), + WeightRenaming(r"neck\.0\.", "neck.conv1."), + WeightRenaming(r"neck\.1\.", "neck.layer_norm1."), + WeightRenaming(r"neck\.2\.", "neck.conv2."), + WeightRenaming(r"neck\.3\.", "neck.layer_norm2."), + WeightRenaming(r"net_2\.", "proj.conv1."), + WeightRenaming(r"net_3\.", "proj.conv2."), + ], + "UnlimitedOcrVisionEncoder": [ + WeightRenaming(r"transformer", "encoder"), + WeightConverter( + source_patterns="qkv_proj.weight", + target_patterns=[ + "q_proj.weight", + "k_proj.weight", + "v_proj.weight", + ], + operations=[Chunk(dim=0)], + ), + WeightConverter( + source_patterns="qkv_proj.bias", + target_patterns=[ + "q_proj.bias", + "k_proj.bias", + "v_proj.bias", + ], + operations=[Chunk(dim=0)], + ), + ], + "UnlimitedOcrTextModel": [ + # WeightConverter( + # source_patterns=[ + # "mlp.experts.*.gate_proj.weight", + # "mlp.experts.*.up_proj.weight", + # ], + # target_patterns="mlp.experts.gate_up_proj", + # operations=[MergeModulelist(dim=0), Concatenate(dim=1)], + # ), + # WeightConverter( + # source_patterns="mlp.experts.*.down_proj.weight", + # target_patterns="mlp.experts.down_proj", + # operations=[MergeModulelist(dim=0)], + # ), + ] } # The legacy mapping is added to the esm model here since the extra weight renaming do not apply to the esm model. mapping["esm"] += mapping["legacy"].copy() From 6c42faa7495be07754bd84fb8c1df9c274fdace4 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 17:44:18 +0200 Subject: [PATCH 012/195] handle sparse layers --- src/transformers/conversion_mapping.py | 2 +- .../models/unlimited_ocr/configuration_unlimited_ocr.py | 5 +++-- .../models/unlimited_ocr/modular_unlimited_ocr.py | 5 +++-- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index aa936113b02d..1baad1d8d938 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -1439,7 +1439,7 @@ def _build_checkpoint_conversion_mapping(): # target_patterns="mlp.experts.down_proj", # operations=[MergeModulelist(dim=0)], # ), - ] + ], } # The legacy mapping is added to the esm model here since the extra weight renaming do not apply to the esm model. mapping["esm"] += mapping["legacy"].copy() diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index f48e1d47d493..57fc918378dd 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -231,8 +231,9 @@ def __post_init__(self, **kwargs): # mask over generated tokens, while the image/prompt prefill is always retained. self.layer_types = ["full_attention"] * self.num_hidden_layers if self.mlp_layer_types is None: - # TODO: double check - self.mlp_layer_types = ["dense"] * self.num_hidden_layers + self.mlp_layer_types = [ + "sparse" if layer_idx >= 1 else "dense" for layer_idx in range(self.num_hidden_layers) + ] self.head_dim = self.hidden_size // self.num_attention_heads if self.num_key_value_heads is None: self.num_key_value_heads = self.num_attention_heads diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 602f34f8aaa1..c11d19f9d050 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -295,8 +295,9 @@ def __post_init__(self, **kwargs): # mask over generated tokens, while the image/prompt prefill is always retained. self.layer_types = ["full_attention"] * self.num_hidden_layers if self.mlp_layer_types is None: - # TODO: double check - self.mlp_layer_types = ["dense"] * self.num_hidden_layers + self.mlp_layer_types = [ + "sparse" if layer_idx >= 1 else "dense" for layer_idx in range(self.num_hidden_layers) + ] super().__post_init__(**kwargs) From 126366486e37ab2e8fa851af52d04e84004542dd Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 15:53:36 +0000 Subject: [PATCH 013/195] update conversion mapping --- src/transformers/conversion_mapping.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 1baad1d8d938..92d4461c9a3d 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -1426,19 +1426,19 @@ def _build_checkpoint_conversion_mapping(): ), ], "UnlimitedOcrTextModel": [ - # WeightConverter( - # source_patterns=[ - # "mlp.experts.*.gate_proj.weight", - # "mlp.experts.*.up_proj.weight", - # ], - # target_patterns="mlp.experts.gate_up_proj", - # operations=[MergeModulelist(dim=0), Concatenate(dim=1)], - # ), - # WeightConverter( - # source_patterns="mlp.experts.*.down_proj.weight", - # target_patterns="mlp.experts.down_proj", - # operations=[MergeModulelist(dim=0)], - # ), + WeightConverter( + source_patterns=[ + "mlp.experts.*.gate_proj.weight", + "mlp.experts.*.up_proj.weight", + ], + target_patterns="mlp.experts.gate_up_proj", + operations=[MergeModulelist(dim=0), Concatenate(dim=1)], + ), + WeightConverter( + source_patterns="mlp.experts.*.down_proj.weight", + target_patterns="mlp.experts.down_proj", + operations=[MergeModulelist(dim=0)], + ), ], } # The legacy mapping is added to the esm model here since the extra weight renaming do not apply to the esm model. From b7c027687fa82291e7600df96906a8c326c9a96b Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 17:57:51 +0200 Subject: [PATCH 014/195] update num_hidden_layers --- .../models/unlimited_ocr/configuration_unlimited_ocr.py | 4 ++-- .../models/unlimited_ocr/modular_unlimited_ocr.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index 57fc918378dd..5dd0486f2956 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -92,7 +92,7 @@ class UnlimitedOcrVisionEncoderConfig(PreTrainedConfig): hidden_size: int = 768 intermediate_size: int = 3072 projection_dim: int = 512 - num_hidden_layers: int = 12 + num_hidden_layers = 24 num_attention_heads: int = 12 num_channels: int = 3 image_size: int | list[int] | tuple[int, int] | None = 224 @@ -194,7 +194,7 @@ class UnlimitedOcrTextConfig(PreTrainedConfig): vocab_size: int = 32000 hidden_size: int = 4096 intermediate_size: int = 11008 - num_hidden_layers: int = 32 + num_hidden_layers: int = 12 num_attention_heads: int = 32 num_key_value_heads: int | None = None hidden_act: str = "silu" diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index c11d19f9d050..f091ae9d0723 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -251,6 +251,7 @@ class UnlimitedOcrVisionEncoderConfig(CLIPVisionConfig): model_type = "unlimited_ocr_vision_encoder" base_config_key = "encoder_config" + num_hidden_layers = 24 @auto_docstring(checkpoint="baidu/Unlimited-OCR") @@ -287,6 +288,7 @@ class UnlimitedOcrTextConfig(DeepseekOcr2TextConfig): model_type = "unlimited_ocr_text" base_config_key = "text_config" layer_types: list[str] | None = None + num_hidden_layers: int = 12 sliding_window: int | None = 128 def __post_init__(self, **kwargs): From 34abc9f02477954e049205c570fc1c8de75874ee Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 17:59:11 +0200 Subject: [PATCH 015/195] ignore lm_head --- src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py | 1 + src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 4574979c19f7..0a94901cdc67 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1636,6 +1636,7 @@ def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): ) class UnlimitedOcrModel(UnlimitedOcrPreTrainedModel): base_model_prefix = "model" + _keys_to_ignore_on_load_unexpected = {"lm_head"} def __init__(self, config: UnlimitedOcrConfig): super().__init__(config) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index f091ae9d0723..c504192687b4 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -617,6 +617,8 @@ def forward( class UnlimitedOcrModel(DeepseekOcr2Model): + _keys_to_ignore_on_load_unexpected = {"lm_head"} + def __init__(self, config: UnlimitedOcrConfig): super().__init__(config) self.multi_modal_projector = nn.Linear( From 6cf169b78e104ba3666e660974bdaa42ecdfc7d6 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 16:17:02 +0000 Subject: [PATCH 016/195] update conversion mapping --- src/transformers/conversion_mapping.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 92d4461c9a3d..d4a005b97847 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -1380,15 +1380,17 @@ def _build_checkpoint_conversion_mapping(): ), WeightRenaming(r"decode_head\.conv_matting\.", r"decode_head.predictor."), ], - "UnlimitedOcrModel": [ - WeightRenaming(r"model\.sam_model", "vision_tower.sam_encoder"), - WeightRenaming(r"model\.vision_model", "vision_tower.vision_encoder"), - WeightRenaming(r"model\.view_seperator", "view_separator"), - WeightRenaming(r"model\.image_newline", "image_newline"), - WeightRenaming(r"model\.projector\.layers\.(weight|bias)", r"multi_modal_projector.\1"), - WeightRenaming(r"model\.embed_tokens\.", r"language_model.embed_tokens."), - WeightRenaming(r"model\.layers\.", r"language_model.layers."), - WeightRenaming(r"model\.norm\.", r"language_model.norm."), + # TODO: unlimited-ocr instead of unlimited_ocr because hub config has model_type="unlimited-ocr" + "unlimited-ocr": [ + WeightRenaming(r"(^|model\.)sam_model", r"\1vision_tower.sam_encoder"), + WeightRenaming(r"(^|model\.)vision_model", r"\1vision_tower.vision_encoder"), + WeightRenaming(r"(^|model\.)view_seperator", r"\1view_separator"), + WeightRenaming(r"(^|model\.)image_newline", r"\1image_newline"), + WeightRenaming(r"(^|model\.)projector\.layers\.weight", r"\1multi_modal_projector.weight"), + WeightRenaming(r"(^|model\.)projector\.layers\.bias", r"\1multi_modal_projector.bias"), + WeightRenaming(r"(^|model\.)embed_tokens\.", r"\1language_model.embed_tokens."), + WeightRenaming(r"(^|model\.)layers\.", r"\1language_model.layers."), + WeightRenaming(r"(^|model\.)norm\.", r"\1language_model.norm."), ], "UnlimitedOcrSamVisionEncoder": [ WeightRenaming(r"blocks\.(\d+)\.norm1\.", r"layers.\1.layer_norm1."), From 73281a717b351c12294323ee641e07ab55d0d623 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 23 Jun 2026 18:17:45 +0200 Subject: [PATCH 017/195] draft docs --- docs/source/en/model_doc/unlimited_ocr.md | 109 +++++++++++++++++----- 1 file changed, 88 insertions(+), 21 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 6d762b4a9b8e..75adc1d11678 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -16,30 +16,90 @@ limitations under the License. ⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer. --> -*This model was contributed to Hugging Face Transformers on 2026-06-23.* +*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-06-23.* # UnlimitedOcr ## Overview -The UnlimitedOcr model was proposed in []() by . - +The UnlimitedOcr model was proposed in [Unlimited OCR Works](https://huggingface.co/papers/2606.23050) by Youyang Yin, Huanhuan Liu, Qunyi Xie, Chaorun Liu, Shiqi Yang, Shaohua Wang, Zhanlong Liu, Hao Zou, Jinyue Chen, Shu Wei, Jingjing Wu, Mingxin Huang, Zhen Wu, Guibin Wang, Tengyu Du, and Lei Jia from Baidu Inc. -The abstract from the paper is the following: - - +Unlimited-OCR is an OCR-specialized vision-language model designed for one-shot long-horizon parsing of single images and multi-page documents. It extends [DeepSeek-OCR-2](deepseek_ocr2) with a two-stage vision pipeline: a SAM ViT-B vision encoder extracts spatial features, which are then fed into a CLIP ViT encoder; the concatenated outputs are projected through an MLP into a DeepSeek-V2 Mixture-of-Experts language model. The 3B-parameter model supports up to 32,768 context tokens, making it suited for parsing long or multi-page documents in a single forward pass. Tips: - +- Unlimited-OCR supports two inference configurations: the default "gundam" mode uses 640Γ—640 tiles with dynamic cropping for high-resolution documents, and "base" mode uses a single 1024Γ—1024 global view for standard-resolution inputs. Gundam mode is enabled by default via the image processor (`crop_to_patches=True`, `tile_size=640`). +- For multi-page documents, pass all page images together with one `` token per page in the text prompt. The model processes all pages jointly within a single context window. +- The sliding-window attention applies only to generated tokens. All image and prompt tokens from the prefill remain fully visible throughout decoding, so long documents do not lose context from earlier pages. +- Use `\nFree OCR.` for plain text extraction and `\nDocument parsing.` for richer structured output. -This model was contributed by [INSERT YOUR HF USERNAME HERE](https://huggingface.co/). -The original code can be found [here](). +This model was contributed by [guarin](https://huggingface.co/guarin). +The original code can be found [here](https://github.com/baidu/Unlimited-OCR). ## Usage examples - +### Single-page OCR + +```python +from transformers import AutoProcessor, AutoModelForImageTextToText + +model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", device_map="auto") +processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") + +image = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" +inputs = processor(images=image, text="\nFree OCR.", return_tensors="pt").to(model.device) + +output = model.generate(**inputs, max_new_tokens=4096) +processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) +# "R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM\nName/Phone Ext. : (...)" +``` + +### Document parsing + +For richer structured output such as markdown-formatted documents, use the `Document parsing.` prompt: + +```python +from transformers import AutoProcessor, AutoModelForImageTextToText + +model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", device_map="auto") +processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") + +image = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" +inputs = processor( + images=image, + text="\nDocument parsing.", + return_tensors="pt", +).to(model.device) + +output = model.generate(**inputs, max_new_tokens=4096) +processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) +``` + +### Multi-page document OCR + +Multi-page documents can be parsed jointly in a single forward pass by passing all page images together. Include one `` token per page in the text prompt so the model processes all pages as a continuous document: + +```python +from PIL import Image +from transformers import AutoProcessor, AutoModelForImageTextToText + +model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", device_map="auto") +processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") + +page1 = Image.open("page1.png") +page2 = Image.open("page2.png") +num_pages = 2 + +inputs = processor( + images=[page1, page2], + text="" * num_pages + "\nMulti page document parsing.", + return_tensors="pt", +).to(model.device) + +output = model.generate(**inputs, max_new_tokens=32768) +processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) +``` ## UnlimitedOcrConfig @@ -61,31 +121,38 @@ The original code can be found [here](). [[autodoc]] UnlimitedOcrSamVisionConfig -## UnlimitedOcrForConditionalGeneration +## UnlimitedOcrImageProcessor -[[autodoc]] UnlimitedOcrForConditionalGeneration +[[autodoc]] UnlimitedOcrImageProcessor -## UnlimitedOcrModel +## UnlimitedOcrProcessor -[[autodoc]] UnlimitedOcrModel - - forward +[[autodoc]] UnlimitedOcrProcessor ## UnlimitedOcrPreTrainedModel [[autodoc]] UnlimitedOcrPreTrainedModel - - forward + +## UnlimitedOcrTextPreTrainedModel + +[[autodoc]] UnlimitedOcrTextPreTrainedModel ## UnlimitedOcrTextModel [[autodoc]] UnlimitedOcrTextModel - forward -## UnlimitedOcrTextPreTrainedModel +## UnlimitedOcrVisionModel -[[autodoc]] UnlimitedOcrTextPreTrainedModel +[[autodoc]] UnlimitedOcrVisionModel - forward -## UnlimitedOcrVisionModel +## UnlimitedOcrModel -[[autodoc]] UnlimitedOcrVisionModel - - forward \ No newline at end of file +[[autodoc]] UnlimitedOcrModel + - forward + +## UnlimitedOcrForConditionalGeneration + +[[autodoc]] UnlimitedOcrForConditionalGeneration + - forward From 8ecbf2c9e5f30e6dc97a72bc4e9cedab39e9be8a Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 08:09:51 +0200 Subject: [PATCH 018/195] add default config values --- .../configuration_unlimited_ocr.py | 63 ++++++++++++------- .../unlimited_ocr/modular_unlimited_ocr.py | 55 ++++++++++++++-- 2 files changed, 91 insertions(+), 27 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index 5dd0486f2956..acbd3914e932 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -40,7 +40,7 @@ class UnlimitedOcrSamVisionConfig(PreTrainedConfig): The indexes of the global attention layers. mlp_dim (`int`, *optional*, defaults to 3072): The dimensionality of the MLP layer in the Transformer encoder. - downsample_channels (`list[int]`, *optional*, defaults to `(512, 896)`): + downsample_channels (`list[int]`, *optional*, defaults to `(512, 1024)`): The channel dimensions for the multi-scale downsampling neck layers. """ @@ -69,7 +69,7 @@ class UnlimitedOcrSamVisionConfig(PreTrainedConfig): def __post_init__(self, **kwargs): if self.downsample_channels is None: - self.downsample_channels = [512, 896] + self.downsample_channels = [512, 1024] return PretrainedConfig.__post_init__(self, **kwargs) @@ -88,15 +88,14 @@ class UnlimitedOcrVisionEncoderConfig(PreTrainedConfig): model_type = "unlimited_ocr_vision_encoder" base_config_key = "encoder_config" - - hidden_size: int = 768 - intermediate_size: int = 3072 + hidden_size: int = 1024 + intermediate_size: int = 4096 projection_dim: int = 512 - num_hidden_layers = 24 - num_attention_heads: int = 12 + num_hidden_layers: int = 24 + num_attention_heads: int = 16 num_channels: int = 3 image_size: int | list[int] | tuple[int, int] | None = 224 - patch_size: int | list[int] | tuple[int, int] | None = 32 + patch_size: int | list[int] | tuple[int, int] | None = 14 hidden_act: str = "quick_gelu" layer_norm_eps: float = 1e-5 attention_dropout: int | float | None = 0.0 @@ -133,6 +132,15 @@ class UnlimitedOcrVisionConfig(PreTrainedConfig): encoder_config: dict | PreTrainedConfig | None = None def __post_init__(self, **kwargs): + if self.sam_config is None: + self.sam_config = self.sub_configs["sam_config"]() + elif isinstance(self.sam_config, dict): + self.sam_config = self.sub_configs["sam_config"](**self.sam_config) + + if self.encoder_config is None: + self.encoder_config = self.sub_configs["encoder_config"]() + elif isinstance(self.encoder_config, dict): + self.encoder_config = self.sub_configs["encoder_config"](**self.encoder_config) if self.sam_config is None: self.sam_config = UnlimitedOcrSamVisionConfig() elif isinstance(self.sam_config, dict): @@ -190,21 +198,20 @@ class UnlimitedOcrTextConfig(PreTrainedConfig): "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } - - vocab_size: int = 32000 - hidden_size: int = 4096 - intermediate_size: int = 11008 + vocab_size: int = 129280 + hidden_size: int = 1280 + intermediate_size: int = 6848 num_hidden_layers: int = 12 - num_attention_heads: int = 32 - num_key_value_heads: int | None = None + num_attention_heads: int = 10 + num_key_value_heads: int | None = 10 hidden_act: str = "silu" - max_position_embeddings: int = 2048 + max_position_embeddings: int = 32768 initializer_range: float = 0.02 rms_norm_eps: float = 1e-6 use_cache: bool = True pad_token_id: int | None = None - bos_token_id: int | None = 1 - eos_token_id: int | list[int] | None = 2 + bos_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 1 pretraining_tp: int | None = 1 tie_word_embeddings: bool = False rope_parameters: RopeParameters | dict | None = None @@ -212,14 +219,14 @@ class UnlimitedOcrTextConfig(PreTrainedConfig): attention_dropout: float | None = 0.0 mlp_bias: bool = False head_dim: int | None = None - n_group: int | None = None + n_group: int | None = 1 n_routed_experts: int = 64 n_shared_experts: int = 2 routed_scaling_factor: float = 1.0 - topk_group: int | None = None + topk_group: int | None = 1 topk_method: str | None = "greedy" - num_experts_per_tok: int | None = None - moe_intermediate_size: int = 1407 + num_experts_per_tok: int | None = 6 + moe_intermediate_size: int = 896 base_config_key = "text_config" mlp_layer_types: list[str] | None = None layer_types: list[str] | None = None @@ -231,8 +238,11 @@ def __post_init__(self, **kwargs): # mask over generated tokens, while the image/prompt prefill is always retained. self.layer_types = ["full_attention"] * self.num_hidden_layers if self.mlp_layer_types is None: + # Some configs may use `first_k_dense_replace` instead of `layer_types`/`mlp_layer_types` + first_k_dense_replace = kwargs.pop("first_k_dense_replace", 1) self.mlp_layer_types = [ - "sparse" if layer_idx >= 1 else "dense" for layer_idx in range(self.num_hidden_layers) + "sparse" if layer_idx >= first_k_dense_replace else "dense" + for layer_idx in range(self.num_hidden_layers) ] self.head_dim = self.hidden_size // self.num_attention_heads if self.num_key_value_heads is None: @@ -268,6 +278,15 @@ class UnlimitedOcrConfig(PreTrainedConfig): tie_word_embeddings: bool = False def __post_init__(self, **kwargs): + if self.vision_config is None: + self.vision_config = self.sub_configs["vision_config"]() + elif isinstance(self.vision_config, dict): + self.vision_config = self.sub_configs["vision_config"](**self.vision_config) + + if self.text_config is None: + self.text_config = self.sub_configs["text_config"]() + elif isinstance(self.text_config, dict): + self.text_config = self.sub_configs["text_config"](**self.text_config) if self.vision_config is None: self.vision_config = UnlimitedOcrVisionConfig() elif isinstance(self.vision_config, dict): diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index c504192687b4..5fef255c7c20 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -221,7 +221,7 @@ class UnlimitedOcrSamVisionConfig(GotOcr2VisionConfig): The indexes of the global attention layers. mlp_dim (`int`, *optional*, defaults to 3072): The dimensionality of the MLP layer in the Transformer encoder. - downsample_channels (`list[int]`, *optional*, defaults to `(512, 896)`): + downsample_channels (`list[int]`, *optional*, defaults to `(512, 1024)`): The channel dimensions for the multi-scale downsampling neck layers. """ @@ -232,7 +232,7 @@ class UnlimitedOcrSamVisionConfig(GotOcr2VisionConfig): def __post_init__(self, **kwargs): if self.downsample_channels is None: - self.downsample_channels = [512, 896] + self.downsample_channels = [512, 1024] return PretrainedConfig.__post_init__(self, **kwargs) @@ -251,7 +251,11 @@ class UnlimitedOcrVisionEncoderConfig(CLIPVisionConfig): model_type = "unlimited_ocr_vision_encoder" base_config_key = "encoder_config" - num_hidden_layers = 24 + hidden_size: int = 1024 + intermediate_size: int = 4096 + num_hidden_layers: int = 24 + num_attention_heads: int = 16 + patch_size: int | list[int] | tuple[int, int] | None = 14 @auto_docstring(checkpoint="baidu/Unlimited-OCR") @@ -264,6 +268,19 @@ class UnlimitedOcrVisionConfig(DeepseekOcr2VisionConfig): "encoder_config": UnlimitedOcrVisionEncoderConfig, } + def __post_init__(self, **kwargs): + if self.sam_config is None: + self.sam_config = self.sub_configs["sam_config"]() + elif isinstance(self.sam_config, dict): + self.sam_config = self.sub_configs["sam_config"](**self.sam_config) + + if self.encoder_config is None: + self.encoder_config = self.sub_configs["encoder_config"]() + elif isinstance(self.encoder_config, dict): + self.encoder_config = self.sub_configs["encoder_config"](**self.encoder_config) + + super().__post_init__(**kwargs) + @auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict @@ -287,8 +304,20 @@ class UnlimitedOcrTextConfig(DeepseekOcr2TextConfig): model_type = "unlimited_ocr_text" base_config_key = "text_config" - layer_types: list[str] | None = None + vocab_size: int = 129280 + hidden_size: int = 1280 + intermediate_size: int = 6848 num_hidden_layers: int = 12 + num_attention_heads: int = 10 + num_key_value_heads: int | None = 10 + max_position_embeddings: int = 32768 + bos_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 1 + moe_intermediate_size: int = 896 + n_group: int | None = 1 + topk_group: int | None = 1 + num_experts_per_tok: int | None = 6 + layer_types: list[str] | None = None sliding_window: int | None = 128 def __post_init__(self, **kwargs): @@ -297,8 +326,11 @@ def __post_init__(self, **kwargs): # mask over generated tokens, while the image/prompt prefill is always retained. self.layer_types = ["full_attention"] * self.num_hidden_layers if self.mlp_layer_types is None: + # Some configs may use `first_k_dense_replace` instead of `layer_types`/`mlp_layer_types` + first_k_dense_replace = kwargs.pop("first_k_dense_replace", 1) self.mlp_layer_types = [ - "sparse" if layer_idx >= 1 else "dense" for layer_idx in range(self.num_hidden_layers) + "sparse" if layer_idx >= first_k_dense_replace else "dense" + for layer_idx in range(self.num_hidden_layers) ] super().__post_init__(**kwargs) @@ -312,6 +344,19 @@ class UnlimitedOcrConfig(DeepseekOcr2Config): "text_config": UnlimitedOcrTextConfig, } + def __post_init__(self, **kwargs): + if self.vision_config is None: + self.vision_config = self.sub_configs["vision_config"]() + elif isinstance(self.vision_config, dict): + self.vision_config = self.sub_configs["vision_config"](**self.vision_config) + + if self.text_config is None: + self.text_config = self.sub_configs["text_config"]() + elif isinstance(self.text_config, dict): + self.text_config = self.sub_configs["text_config"](**self.text_config) + + super().__post_init__(**kwargs) + class UnlimitedOcrModelOutputWithPooling(DeepseekOcr2ModelOutputWithPooling): pass From 96a794a3543282606357e5077d1fe5f4d3268f02 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 08:17:46 +0200 Subject: [PATCH 019/195] update text config handling --- .../models/unlimited_ocr/configuration_unlimited_ocr.py | 5 ++++- .../models/unlimited_ocr/modular_unlimited_ocr.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index acbd3914e932..c02c4fb508f9 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -284,7 +284,10 @@ def __post_init__(self, **kwargs): self.vision_config = self.sub_configs["vision_config"](**self.vision_config) if self.text_config is None: - self.text_config = self.sub_configs["text_config"]() + text_cls = self.sub_configs["text_config"] + text_keys = text_cls().to_dict().keys() + text_kwargs = {key: kwargs.pop(key) for key in text_keys if key in kwargs} + self.text_config = text_cls(**text_kwargs) elif isinstance(self.text_config, dict): self.text_config = self.sub_configs["text_config"](**self.text_config) if self.vision_config is None: diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 5fef255c7c20..88dbd52747ca 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -351,7 +351,10 @@ def __post_init__(self, **kwargs): self.vision_config = self.sub_configs["vision_config"](**self.vision_config) if self.text_config is None: - self.text_config = self.sub_configs["text_config"]() + text_cls = self.sub_configs["text_config"] + text_keys = text_cls().to_dict().keys() + text_kwargs = {key: kwargs.pop(key) for key in text_keys if key in kwargs} + self.text_config = text_cls(**text_kwargs) elif isinstance(self.text_config, dict): self.text_config = self.sub_configs["text_config"](**self.text_config) From d4aa87892679a6a43767778a33d435fc7e9ed80b Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 08:41:38 +0200 Subject: [PATCH 020/195] update post init --- .../configuration_unlimited_ocr.py | 24 +++---------------- .../unlimited_ocr/modular_unlimited_ocr.py | 6 ++--- 2 files changed, 6 insertions(+), 24 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index c02c4fb508f9..e24fef3c0966 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -141,17 +141,8 @@ def __post_init__(self, **kwargs): self.encoder_config = self.sub_configs["encoder_config"]() elif isinstance(self.encoder_config, dict): self.encoder_config = self.sub_configs["encoder_config"](**self.encoder_config) - if self.sam_config is None: - self.sam_config = UnlimitedOcrSamVisionConfig() - elif isinstance(self.sam_config, dict): - self.sam_config = UnlimitedOcrSamVisionConfig(**self.sam_config) - if self.encoder_config is None: - self.encoder_config = UnlimitedOcrVisionEncoderConfig() - elif isinstance(self.encoder_config, dict): - self.encoder_config = UnlimitedOcrVisionEncoderConfig(**self.encoder_config) - - super().__post_init__(**kwargs) + PretrainedConfig.__post_init__(self, **kwargs) @auto_docstring(checkpoint="baidu/Unlimited-OCR") @@ -287,20 +278,11 @@ def __post_init__(self, **kwargs): text_cls = self.sub_configs["text_config"] text_keys = text_cls().to_dict().keys() text_kwargs = {key: kwargs.pop(key) for key in text_keys if key in kwargs} + text_kwargs = {} self.text_config = text_cls(**text_kwargs) elif isinstance(self.text_config, dict): self.text_config = self.sub_configs["text_config"](**self.text_config) - if self.vision_config is None: - self.vision_config = UnlimitedOcrVisionConfig() - elif isinstance(self.vision_config, dict): - self.vision_config = UnlimitedOcrVisionConfig(**self.vision_config) - - if self.text_config is None: - self.text_config = UnlimitedOcrTextConfig() - elif isinstance(self.text_config, dict): - self.text_config = UnlimitedOcrTextConfig(**self.text_config) - - super().__post_init__(**kwargs) + PretrainedConfig.__post_init__(self, **kwargs) __all__ = [ diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 88dbd52747ca..7fe89365fea2 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -279,7 +279,7 @@ def __post_init__(self, **kwargs): elif isinstance(self.encoder_config, dict): self.encoder_config = self.sub_configs["encoder_config"](**self.encoder_config) - super().__post_init__(**kwargs) + PretrainedConfig.__post_init__(self, **kwargs) @auto_docstring(checkpoint="baidu/Unlimited-OCR") @@ -354,11 +354,11 @@ def __post_init__(self, **kwargs): text_cls = self.sub_configs["text_config"] text_keys = text_cls().to_dict().keys() text_kwargs = {key: kwargs.pop(key) for key in text_keys if key in kwargs} + text_kwargs = {} self.text_config = text_cls(**text_kwargs) elif isinstance(self.text_config, dict): self.text_config = self.sub_configs["text_config"](**self.text_config) - - super().__post_init__(**kwargs) + PretrainedConfig.__post_init__(self, **kwargs) class UnlimitedOcrModelOutputWithPooling(DeepseekOcr2ModelOutputWithPooling): From bffafd0a55cc2233e445e3bd8221ff494b842d5f Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 08:45:53 +0200 Subject: [PATCH 021/195] update --- .../models/unlimited_ocr/configuration_unlimited_ocr.py | 9 ++++----- .../models/unlimited_ocr/modular_unlimited_ocr.py | 9 ++++----- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index e24fef3c0966..e095b296ae01 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -275,11 +275,10 @@ def __post_init__(self, **kwargs): self.vision_config = self.sub_configs["vision_config"](**self.vision_config) if self.text_config is None: - text_cls = self.sub_configs["text_config"] - text_keys = text_cls().to_dict().keys() - text_kwargs = {key: kwargs.pop(key) for key in text_keys if key in kwargs} - text_kwargs = {} - self.text_config = text_cls(**text_kwargs) + text_config_cls = self.sub_configs["text_config"] + text_keys = text_config_cls().to_dict().keys() + text_kwargs = {key: kwargs[key] for key in text_keys if key in kwargs} + self.text_config = text_config_cls(**text_kwargs) elif isinstance(self.text_config, dict): self.text_config = self.sub_configs["text_config"](**self.text_config) PretrainedConfig.__post_init__(self, **kwargs) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 7fe89365fea2..6b830ab6c3a6 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -351,11 +351,10 @@ def __post_init__(self, **kwargs): self.vision_config = self.sub_configs["vision_config"](**self.vision_config) if self.text_config is None: - text_cls = self.sub_configs["text_config"] - text_keys = text_cls().to_dict().keys() - text_kwargs = {key: kwargs.pop(key) for key in text_keys if key in kwargs} - text_kwargs = {} - self.text_config = text_cls(**text_kwargs) + text_config_cls = self.sub_configs["text_config"] + text_keys = text_config_cls().to_dict().keys() + text_kwargs = {key: kwargs[key] for key in text_keys if key in kwargs} + self.text_config = text_config_cls(**text_kwargs) elif isinstance(self.text_config, dict): self.text_config = self.sub_configs["text_config"](**self.text_config) PretrainedConfig.__post_init__(self, **kwargs) From 16dcdbba3fd6a9f07e0a81ac5f418927f72ea81d Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 06:54:15 +0000 Subject: [PATCH 022/195] fix to dtype --- tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index d72300ca015b..9b4f5c50cc7c 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -173,9 +173,7 @@ def test_small_model_integration_test_document_parsing(self): "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" ) ) - inputs = self.processor(images=image, text="\ndocument parsing.", return_tensors="pt").to( - model.device, dtype=torch.bfloat16 - ) + inputs = self.processor(images=image, text="\ndocument parsing.", return_tensors="pt").to(model.device) generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True) EXPECTED_DECODED_TEXT = Expectations( @@ -199,7 +197,7 @@ def test_small_model_integration_test_document_parsing_grounding(self): images=image, text="\ndocument parsing.", return_tensors="pt", - ).to(model.device, dtype=torch.bfloat16) + ).to(model.device) generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=False) EXPECTED_DECODED_TEXT = Expectations( @@ -229,7 +227,7 @@ def test_small_model_integration_test_batched(self): text=["\ndocument parsing.", "\ndocument parsing."], return_tensors="pt", padding=True, - ).to(model.device, dtype=torch.bfloat16) + ).to(model.device) generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) decoded = self.processor.batch_decode( generate_ids[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True From 3bba66bd935884e45280eeb5ab9df5da147a79d7 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 08:54:28 +0200 Subject: [PATCH 023/195] update docs --- docs/source/en/model_doc/unlimited_ocr.md | 43 ++++++++++++++++------- 1 file changed, 30 insertions(+), 13 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 75adc1d11678..2fd8fa120258 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -25,21 +25,27 @@ limitations under the License. The UnlimitedOcr model was proposed in [Unlimited OCR Works](https://huggingface.co/papers/2606.23050) by Youyang Yin, Huanhuan Liu, Qunyi Xie, Chaorun Liu, Shiqi Yang, Shaohua Wang, Zhanlong Liu, Hao Zou, Jinyue Chen, Shu Wei, Jingjing Wu, Mingxin Huang, Zhen Wu, Guibin Wang, Tengyu Du, and Lei Jia from Baidu Inc. -Unlimited-OCR is an OCR-specialized vision-language model designed for one-shot long-horizon parsing of single images and multi-page documents. It extends [DeepSeek-OCR-2](deepseek_ocr2) with a two-stage vision pipeline: a SAM ViT-B vision encoder extracts spatial features, which are then fed into a CLIP ViT encoder; the concatenated outputs are projected through an MLP into a DeepSeek-V2 Mixture-of-Experts language model. The 3B-parameter model supports up to 32,768 context tokens, making it suited for parsing long or multi-page documents in a single forward pass. +The abstract from the paper is the following: -Tips: - -- Unlimited-OCR supports two inference configurations: the default "gundam" mode uses 640Γ—640 tiles with dynamic cropping for high-resolution documents, and "base" mode uses a single 1024Γ—1024 global view for standard-resolution inputs. Gundam mode is enabled by default via the image processor (`crop_to_patches=True`, `tile_size=640`). -- For multi-page documents, pass all page images together with one `` token per page in the text prompt. The model processes all pages jointly within a single context window. -- The sliding-window attention applies only to generated tokens. All image and prompt tokens from the prefill remain fully visible throughout decoding, so long documents do not lose context from earlier pages. -- Use `\nFree OCR.` for plain text extraction and `\nDocument parsing.` for richer structured output. +*Recently, end-to-end OCR models, exemplified by DeepSeek OCR, have once again thrust OCR into the spotlight. A widely held view is that employing a large language model (LLM) as the decoder allows the model to leverage the prior distribution of language, leading to improved OCR performance. However, the downside is equally evident: as the output sequence lengthens, the accumulated KV cache drives up memory consumption and progressively slows down generation. This stands in stark contrast to humans, who exhibit no such decline in efficiency during long-horizon copying tasks. In this technical report, we propose Unlimited OCR, a model designed to emulate human parsing working memory. Taking DeepSeek OCR as the baseline, we replace all attention layers in the decoder with our proposed Reference Sliding Window Attention (R-SWA), which reduces attention computation costs while maintaining a constant KV cache throughout the entire decoding process. By combining the high compression rate of DeepSeek OCR's encoder with our constant KV cache design, Unlimited OCR can transcribe dozens of pages of documents in a single forward pass under a standard maximum length of 32K. More importantly, R-SWA is a general-purpose parsing attention mechanism β€” beyond OCR, it is equally applicable to tasks such as ASR, translation, etc. Codes and model weights are publicly available at http://github.com/baidu/Unlimited-OCR* This model was contributed by [guarin](https://huggingface.co/guarin). The original code can be found [here](https://github.com/baidu/Unlimited-OCR). -## Usage examples +> [!TIP] +> Unlimited-OCR supports two inference configurations: the default "gundam" mode uses 640Γ—640 tiles with dynamic cropping for high-resolution documents, and "base" mode uses a single 1024Γ—1024 global view for standard-resolution inputs. Gundam mode is enabled by default via the image processor (`crop_to_patches=True`, `tile_size=640`). + +> [!TIP] +> For multi-page documents, pass all page images together with one `` token per page in the text prompt. The model processes all pages jointly within a single context window. + +> [!TIP] +> The Reference Sliding Window Attention (R-SWA) applies only to generated tokens. All image and prompt tokens from the prefill remain fully visible throughout decoding, so long documents do not lose context from earlier pages. + +> [!TIP] +> Use `\nFree OCR.` for plain text extraction and `\nDocument parsing.` for richer structured output. -### Single-page OCR + + ```python from transformers import AutoProcessor, AutoModelForImageTextToText @@ -55,9 +61,10 @@ processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=T # "R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM\nName/Phone Ext. : (...)" ``` -### Document parsing + + -For richer structured output such as markdown-formatted documents, use the `Document parsing.` prompt: +Use the `Document parsing.` prompt for richer structured output such as markdown-formatted documents. ```python from transformers import AutoProcessor, AutoModelForImageTextToText @@ -76,9 +83,10 @@ output = model.generate(**inputs, max_new_tokens=4096) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) ``` -### Multi-page document OCR + + -Multi-page documents can be parsed jointly in a single forward pass by passing all page images together. Include one `` token per page in the text prompt so the model processes all pages as a continuous document: +Multi-page documents can be parsed jointly in a single forward pass by passing all page images together. Include one `` token per page in the text prompt so the model processes all pages as a continuous document. ```python from PIL import Image @@ -101,6 +109,15 @@ output = model.generate(**inputs, max_new_tokens=32768) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) ``` + + + +## Notes + +- [`UnlimitedOcrForConditionalGeneration`] extends [`DeepseekOcr2ForConditionalGeneration`](deepseek_ocr2) with a two-stage vision pipeline: a SAM ViT-B encoder feeds into a CLIP ViT encoder, and their concatenated features are projected to the language model hidden size. The 3B-parameter model supports up to 32,768 context tokens. +- The Reference Sliding Window Attention (R-SWA) in the decoder keeps the KV cache constant throughout decoding. The prefill (image tokens and prompt) is retained in full; the sliding window applies only across generated tokens. +- Image inputs are only forwarded during the first generation step. Subsequent decode steps skip `pixel_values` to avoid reprocessing the image. + ## UnlimitedOcrConfig [[autodoc]] UnlimitedOcrConfig From ab76c3a8d98ef197fd042a3dcce1ce9901792aba Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 07:30:21 +0000 Subject: [PATCH 024/195] update MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS --- src/transformers/models/auto/tokenization_auto.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/models/auto/tokenization_auto.py b/src/transformers/models/auto/tokenization_auto.py index 77b34d7123d5..cb4c07029e28 100644 --- a/src/transformers/models/auto/tokenization_auto.py +++ b/src/transformers/models/auto/tokenization_auto.py @@ -369,6 +369,7 @@ "deepseek_vl_v2", "deepseek_ocr", "deepseek_ocr2", + "unlimited_ocr", "ernie4_5", "ernie4_5_moe", "fuyu", From f608f9760dc717559f2284ff4cef56d95b29ca1d Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 07:34:02 +0000 Subject: [PATCH 025/195] update conversion mapping to unlimted_ocr --- src/transformers/conversion_mapping.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index d4a005b97847..f6d19b37c0c7 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -1380,8 +1380,7 @@ def _build_checkpoint_conversion_mapping(): ), WeightRenaming(r"decode_head\.conv_matting\.", r"decode_head.predictor."), ], - # TODO: unlimited-ocr instead of unlimited_ocr because hub config has model_type="unlimited-ocr" - "unlimited-ocr": [ + "unlimited_ocr": [ WeightRenaming(r"(^|model\.)sam_model", r"\1vision_tower.sam_encoder"), WeightRenaming(r"(^|model\.)vision_model", r"\1vision_tower.vision_encoder"), WeightRenaming(r"(^|model\.)view_seperator", r"\1view_separator"), From f3463814faf833aad9a8b4478d8db2287df4d2bb Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 09:39:21 +0200 Subject: [PATCH 026/195] add num_local_patches --- .../models/unlimited_ocr/image_processing_unlimited_ocr.py | 3 +++ .../models/unlimited_ocr/modular_unlimited_ocr.py | 4 ++++ 2 files changed, 7 insertions(+) diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index 36f4f83048e9..c0f92725526a 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -251,6 +251,9 @@ def _preprocess( num_columns, num_rows = 1, 1 image_spatial_crop.append([num_columns, num_rows]) batch_feature["image_spatial_crop"] = torch.tensor(image_spatial_crop, dtype=torch.long) + batch_feature["num_local_patches"] = ( + batch_feature["image_spatial_crop"][:, 0] * batch_feature["image_spatial_crop"][:, 1] + ) return batch_feature def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None) -> int: diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 6b830ab6c3a6..f45b7a3c3f18 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -129,7 +129,11 @@ def _preprocess( else: num_columns, num_rows = 1, 1 image_spatial_crop.append([num_columns, num_rows]) + #Β TODO: Rename and cleanup batch_feature["image_spatial_crop"] = torch.tensor(image_spatial_crop, dtype=torch.long) + batch_feature["num_local_patches"] = ( + batch_feature["image_spatial_crop"][:, 0] * batch_feature["image_spatial_crop"][:, 1] + ) return batch_feature From b91da12297e72ef0e2e96de1e3b9b1bbed904a45 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 10:36:59 +0200 Subject: [PATCH 027/195] update num_local_patches handling --- .../unlimited_ocr/image_processing_unlimited_ocr.py | 11 ++++++++--- .../models/unlimited_ocr/modular_unlimited_ocr.py | 12 ++++++++---- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index c0f92725526a..4e3d4a95709a 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -241,19 +241,24 @@ def _preprocess( ) image_spatial_crop = [] + num_local_patches = [] for image in images: height, width = image.shape[-2:] if crop_to_patches and max(height, width) > tile_size: num_columns, num_rows = get_optimal_tiled_canvas( (height, width), (tile_size, tile_size), min_patches, max_patches ) + num_local_patches.append(num_columns * num_rows) else: num_columns, num_rows = 1, 1 + num_local_patches.append(0) image_spatial_crop.append([num_columns, num_rows]) - batch_feature["image_spatial_crop"] = torch.tensor(image_spatial_crop, dtype=torch.long) - batch_feature["num_local_patches"] = ( - batch_feature["image_spatial_crop"][:, 0] * batch_feature["image_spatial_crop"][:, 1] + + # TODO: rename + batch_feature["image_spatial_crop"] = torch.tensor( + image_spatial_crop, dtype=torch.long, device=images[0].device ) + batch_feature["num_local_patches"] = torch.tensor(num_local_patches, dtype=torch.long, device=images[0].device) return batch_feature def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None) -> int: diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index f45b7a3c3f18..66ce58fdb062 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -120,20 +120,24 @@ def _preprocess( ) image_spatial_crop = [] + num_local_patches = [] for image in images: height, width = image.shape[-2:] if crop_to_patches and max(height, width) > tile_size: num_columns, num_rows = get_optimal_tiled_canvas( (height, width), (tile_size, tile_size), min_patches, max_patches ) + num_local_patches.append(num_columns * num_rows) else: num_columns, num_rows = 1, 1 + num_local_patches.append(0) image_spatial_crop.append([num_columns, num_rows]) - #Β TODO: Rename and cleanup - batch_feature["image_spatial_crop"] = torch.tensor(image_spatial_crop, dtype=torch.long) - batch_feature["num_local_patches"] = ( - batch_feature["image_spatial_crop"][:, 0] * batch_feature["image_spatial_crop"][:, 1] + + # TODO: rename + batch_feature["image_spatial_crop"] = torch.tensor( + image_spatial_crop, dtype=torch.long, device=images[0].device ) + batch_feature["num_local_patches"] = torch.tensor(num_local_patches, dtype=torch.long, device=images[0].device) return batch_feature From 4c3dc45faa6587381731331332ae3d213653e49e Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 11:27:44 +0200 Subject: [PATCH 028/195] inline _preprocess --- .../image_processing_unlimited_ocr.py | 87 +++++++++++++----- .../unlimited_ocr/modular_unlimited_ocr.py | 89 ++++++++++++++----- 2 files changed, 129 insertions(+), 47 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index 4e3d4a95709a..296a7bd29a5a 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -25,9 +25,10 @@ from ...feature_extraction_utils import BatchFeature from ...image_processing_backends import TorchvisionBackend +from ...image_transforms import group_images_by_shape, reorder_images from ...image_utils import IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD, PILImageResampling, SizeDict from ...processing_utils import ImagesKwargs, Unpack -from ...utils import auto_docstring +from ...utils import TensorType, auto_docstring class UnlimitedOcrImageProcessorKwargs(ImagesKwargs, total=False): @@ -207,39 +208,79 @@ def crop_image_to_patches( def _preprocess( self, images: list["torch.Tensor"], - size, + size: SizeDict, crop_to_patches: bool, min_patches: int, max_patches: int, tile_size: int, - resample, + resample: "PILImageResampling | None", do_rescale: bool, rescale_factor: float, do_normalize: bool, - image_mean, - image_std, + image_mean: "float | list[float] | None", + image_std: "float | list[float] | None", disable_grouping: bool | None, - return_tensors, + return_tensors: "str | TensorType | None", **kwargs, ) -> BatchFeature: - batch_feature = super()._preprocess( - images, - size=size, - crop_to_patches=crop_to_patches, - min_patches=min_patches, - max_patches=max_patches, - tile_size=tile_size, - resample=resample, - do_rescale=do_rescale, - rescale_factor=rescale_factor, - do_normalize=do_normalize, - image_mean=image_mean, - image_std=image_std, - disable_grouping=disable_grouping, - return_tensors=return_tensors, - **kwargs, - ) + # --- Local patches (batched by shape group) --- + num_local_patches_grouped = {} + local_patches_grouped = {} + + if crop_to_patches: + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + + for shape, stacked_images in grouped_images.items(): + h, w = shape[-2:] + if max(h, w) > tile_size: + stacked_patches, n_patches = self.crop_image_to_patches( + stacked_images, + min_patches=min_patches, + max_patches=max_patches, + tile_size=tile_size, + resample=resample, + ) + flat_patches = stacked_patches.reshape(-1, *stacked_patches.shape[2:]) + flat_patches = self.rescale_and_normalize( + flat_patches, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + local_patches_grouped[shape] = flat_patches.reshape(stacked_patches.shape) + num_local_patches_grouped[shape] = [n_patches] * stacked_images.shape[0] + else: + local_patches_grouped[shape] = [None] * stacked_images.shape[0] + num_local_patches_grouped[shape] = [0] * stacked_images.shape[0] + + ordered_local = reorder_images(local_patches_grouped, grouped_images_index) + else: + ordered_local = [] + + flat_local_list = [patch for item in ordered_local if item is not None for patch in item] + + # --- Global view (batched by shape group) --- + global_target_size = size.height if crop_to_patches else tile_size + + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + processed_global_grouped = {} + for shape, stacked in grouped_images.items(): + h, w = shape[-2:] + scale = global_target_size / max(h, w) + new_h = round(h * scale) + new_w = round(w * scale) + stacked = self.resize(stacked, SizeDict(height=new_h, width=new_w), resample=resample) + stacked = self.pad_to_square(stacked, background_color=self.background_color) + stacked = self.rescale_and_normalize( + stacked, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + processed_global_grouped[shape] = stacked + all_pixel_values_global = reorder_images(processed_global_grouped, grouped_images_index) + + data = {"pixel_values": all_pixel_values_global} + if flat_local_list: + data["pixel_values_local"] = flat_local_list + + batch_feature = BatchFeature(data=data, tensor_type=return_tensors) + # Compute per-image spatial crop grid and local-patch counts. image_spatial_crop = [] num_local_patches = [] for image in images: diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 66ce58fdb062..4d81afdebf94 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -21,7 +21,8 @@ from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PretrainedConfig from ...feature_extraction_utils import BatchFeature -from ...image_utils import ImageInput +from ...image_transforms import group_images_by_shape, reorder_images +from ...image_utils import ImageInput, SizeDict from ...masking_utils import ( and_masks, causal_mask_function, @@ -31,7 +32,7 @@ from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast from ...processing_utils import Unpack from ...tokenization_utils_base import PreTokenizedInput, TextInput -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_int +from ...utils import TensorType, TransformersKwargs, auto_docstring, can_return_tuple, torch_int from ...utils.generic import merge_with_config_defaults from ...utils.output_capturing import capture_outputs from ..clip.configuration_clip import CLIPVisionConfig @@ -86,39 +87,79 @@ class UnlimitedOcrImageProcessor(DeepseekOcr2ImageProcessor): def _preprocess( self, images: list["torch.Tensor"], - size, + size: SizeDict, crop_to_patches: bool, min_patches: int, max_patches: int, tile_size: int, - resample, + resample: "PILImageResampling | None", do_rescale: bool, rescale_factor: float, do_normalize: bool, - image_mean, - image_std, + image_mean: "float | list[float] | None", + image_std: "float | list[float] | None", disable_grouping: bool | None, - return_tensors, + return_tensors: "str | TensorType | None", **kwargs, ) -> BatchFeature: - batch_feature = super()._preprocess( - images, - size=size, - crop_to_patches=crop_to_patches, - min_patches=min_patches, - max_patches=max_patches, - tile_size=tile_size, - resample=resample, - do_rescale=do_rescale, - rescale_factor=rescale_factor, - do_normalize=do_normalize, - image_mean=image_mean, - image_std=image_std, - disable_grouping=disable_grouping, - return_tensors=return_tensors, - **kwargs, - ) + # --- Local patches (batched by shape group) --- + num_local_patches_grouped = {} + local_patches_grouped = {} + + if crop_to_patches: + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + + for shape, stacked_images in grouped_images.items(): + h, w = shape[-2:] + if max(h, w) > tile_size: + stacked_patches, n_patches = self.crop_image_to_patches( + stacked_images, + min_patches=min_patches, + max_patches=max_patches, + tile_size=tile_size, + resample=resample, + ) + flat_patches = stacked_patches.reshape(-1, *stacked_patches.shape[2:]) + flat_patches = self.rescale_and_normalize( + flat_patches, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + local_patches_grouped[shape] = flat_patches.reshape(stacked_patches.shape) + num_local_patches_grouped[shape] = [n_patches] * stacked_images.shape[0] + else: + local_patches_grouped[shape] = [None] * stacked_images.shape[0] + num_local_patches_grouped[shape] = [0] * stacked_images.shape[0] + + ordered_local = reorder_images(local_patches_grouped, grouped_images_index) + else: + ordered_local = [] + + flat_local_list = [patch for item in ordered_local if item is not None for patch in item] + + # --- Global view (batched by shape group) --- + global_target_size = size.height if crop_to_patches else tile_size + + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + processed_global_grouped = {} + for shape, stacked in grouped_images.items(): + h, w = shape[-2:] + scale = global_target_size / max(h, w) + new_h = round(h * scale) + new_w = round(w * scale) + stacked = self.resize(stacked, SizeDict(height=new_h, width=new_w), resample=resample) + stacked = self.pad_to_square(stacked, background_color=self.background_color) + stacked = self.rescale_and_normalize( + stacked, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + processed_global_grouped[shape] = stacked + all_pixel_values_global = reorder_images(processed_global_grouped, grouped_images_index) + + data = {"pixel_values": all_pixel_values_global} + if flat_local_list: + data["pixel_values_local"] = flat_local_list + + batch_feature = BatchFeature(data=data, tensor_type=return_tensors) + # Compute per-image spatial crop grid and local-patch counts. image_spatial_crop = [] num_local_patches = [] for image in images: From 0270a2f85283da88aa093722c2d9d43030695a4e Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 09:36:15 +0000 Subject: [PATCH 029/195] update tests --- .../test_modeling_unlimited_ocr.py | 35 +++++++++++-------- 1 file changed, 21 insertions(+), 14 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 9b4f5c50cc7c..aac0661e1860 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -21,7 +21,14 @@ is_torch_available, is_vision_available, ) -from transformers.testing_utils import Expectations, cleanup, require_torch, slow, torch_device +from transformers.testing_utils import ( + Expectations, + cleanup, + require_torch, + require_torch_accelerator, + slow, + torch_device, +) from ...test_processing_common import url_to_local_path from ...vlm_tester import VLMModelTest, VLMModelTester @@ -164,17 +171,17 @@ def tearDown(self): cleanup(torch_device, gc_collect=True) @slow + @require_torch_accelerator def test_small_model_integration_test_document_parsing(self): - model = UnlimitedOcrForConditionalGeneration.from_pretrained( - self.model_id, torch_dtype=torch.bfloat16, device_map=torch_device - ) + model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, device_map=torch_device).eval() image = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" ) ) inputs = self.processor(images=image, text="\ndocument parsing.", return_tensors="pt").to(model.device) - generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) + with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): + generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True) EXPECTED_DECODED_TEXT = Expectations( { @@ -184,10 +191,9 @@ def test_small_model_integration_test_document_parsing(self): self.assertEqual(decoded, EXPECTED_DECODED_TEXT) @slow + @require_torch_accelerator def test_small_model_integration_test_document_parsing_grounding(self): - model = UnlimitedOcrForConditionalGeneration.from_pretrained( - self.model_id, torch_dtype=torch.bfloat16, device_map=torch_device - ) + model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, device_map=torch_device).eval() image = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -198,7 +204,8 @@ def test_small_model_integration_test_document_parsing_grounding(self): text="\ndocument parsing.", return_tensors="pt", ).to(model.device) - generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) + with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): + generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=False) EXPECTED_DECODED_TEXT = Expectations( { @@ -208,10 +215,9 @@ def test_small_model_integration_test_document_parsing_grounding(self): self.assertEqual(decoded, EXPECTED_DECODED_TEXT) @slow + @require_torch_accelerator def test_small_model_integration_test_batched(self): - model = UnlimitedOcrForConditionalGeneration.from_pretrained( - self.model_id, torch_dtype=torch.bfloat16, device_map=torch_device - ) + model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, device_map=torch_device).eval() image1 = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -227,8 +233,9 @@ def test_small_model_integration_test_batched(self): text=["\ndocument parsing.", "\ndocument parsing."], return_tensors="pt", padding=True, - ).to(model.device) - generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) + ).to(model.device, dtype=torch.bfloat16) + with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): + generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) decoded = self.processor.batch_decode( generate_ids[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True ) From 91e599658b29efb087355ed6ea62e8119c747e1d Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 11:36:25 +0200 Subject: [PATCH 030/195] update docs --- docs/source/en/model_doc/unlimited_ocr.md | 32 +++-------------------- 1 file changed, 3 insertions(+), 29 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 2fd8fa120258..e32d2f988c2c 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -33,7 +33,7 @@ This model was contributed by [guarin](https://huggingface.co/guarin). The original code can be found [here](https://github.com/baidu/Unlimited-OCR). > [!TIP] -> Unlimited-OCR supports two inference configurations: the default "gundam" mode uses 640Γ—640 tiles with dynamic cropping for high-resolution documents, and "base" mode uses a single 1024Γ—1024 global view for standard-resolution inputs. Gundam mode is enabled by default via the image processor (`crop_to_patches=True`, `tile_size=640`). +> Unlimited-OCR supports two inference configurations: the default "gundam" mode uses 640x640 tiles with dynamic cropping for high-resolution documents, and "base" mode uses a single 1024x1024 global view for standard-resolution inputs. Gundam mode is enabled by default via the image processor (`crop_to_patches=True`, `tile_size=640`). > [!TIP] > For multi-page documents, pass all page images together with one `` token per page in the text prompt. The model processes all pages jointly within a single context window. @@ -41,9 +41,6 @@ The original code can be found [here](https://github.com/baidu/Unlimited-OCR). > [!TIP] > The Reference Sliding Window Attention (R-SWA) applies only to generated tokens. All image and prompt tokens from the prefill remain fully visible throughout decoding, so long documents do not lose context from earlier pages. -> [!TIP] -> Use `\nFree OCR.` for plain text extraction and `\nDocument parsing.` for richer structured output. - @@ -54,36 +51,13 @@ model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", devic processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") image = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" -inputs = processor(images=image, text="\nFree OCR.", return_tensors="pt").to(model.device) +inputs = processor(images=image, text="\ndocument parsing.", return_tensors="pt").to(model.device) output = model.generate(**inputs, max_new_tokens=4096) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) # "R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM\nName/Phone Ext. : (...)" ``` - - - -Use the `Document parsing.` prompt for richer structured output such as markdown-formatted documents. - -```python -from transformers import AutoProcessor, AutoModelForImageTextToText - -model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", device_map="auto") -processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") - -image = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" -inputs = processor( - images=image, - text="\nDocument parsing.", - return_tensors="pt", -).to(model.device) - -output = model.generate(**inputs, max_new_tokens=4096) -processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) -``` - - Multi-page documents can be parsed jointly in a single forward pass by passing all page images together. Include one `` token per page in the text prompt so the model processes all pages as a continuous document. @@ -101,7 +75,7 @@ num_pages = 2 inputs = processor( images=[page1, page2], - text="" * num_pages + "\nMulti page document parsing.", + text="" * num_pages + "\nMulti page parsing.", return_tensors="pt", ).to(model.device) From 15150c856a5d37b0e591a4a46a7277769446d63f Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 12:07:58 +0200 Subject: [PATCH 031/195] add multi page test --- .../unlimited_ocr/modular_unlimited_ocr.py | 2 +- .../test_modeling_unlimited_ocr.py | 32 +++++++++++++++++-- 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 4d81afdebf94..9003566eb31f 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -22,7 +22,7 @@ from ...configuration_utils import PretrainedConfig from ...feature_extraction_utils import BatchFeature from ...image_transforms import group_images_by_shape, reorder_images -from ...image_utils import ImageInput, SizeDict +from ...image_utils import ImageInput, PILImageResampling, SizeDict from ...masking_utils import ( and_masks, causal_mask_function, diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index aac0661e1860..e63810408828 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -216,7 +216,7 @@ def test_small_model_integration_test_document_parsing_grounding(self): @slow @require_torch_accelerator - def test_small_model_integration_test_batched(self): + def test_small_model_integration_test_document_parsing_batched(self): model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, device_map=torch_device).eval() image1 = load_image( url_to_local_path( @@ -234,11 +234,13 @@ def test_small_model_integration_test_batched(self): return_tensors="pt", padding=True, ).to(model.device, dtype=torch.bfloat16) + with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) decoded = self.processor.batch_decode( generate_ids[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True ) + EXPECTED_DECODED_TEXT = Expectations( { ("cuda", None): [ @@ -246,5 +248,31 @@ def test_small_model_integration_test_batched(self): "header [53, 23, 365, 41]Advanced Template and Styl", ], } - ).get_expectation() # fmt: skip + ).get_expectation() + self.assertEqual(decoded, EXPECTED_DECODED_TEXT) + + @slow + @require_torch_accelerator + def test_small_model_integration_test_multi_page_document_parsing(self): + model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, device_map=torch_device).eval() + image = load_image( + url_to_local_path( + "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" + ) + ) + inputs = self.processor(images=image, text="\nMulti page parsing.", return_tensors="pt").to( + model.device + ) + + with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): + generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) + decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True) + + EXPECTED_DECODED_TEXT = Expectations( + { + ("cuda", None): [ + "image [382, 87, 489, 174]\n", + ], + } + ).get_expectation() self.assertEqual(decoded, EXPECTED_DECODED_TEXT) From 76307e919ee4104ebd5e1ce70a878cc0a2316dc8 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 12:29:49 +0200 Subject: [PATCH 032/195] update pad_to_square --- .../image_processing_unlimited_ocr.py | 4 +- .../unlimited_ocr/modular_unlimited_ocr.py | 43 +++++++++++++++++++ 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index 296a7bd29a5a..9ffceb8a6e7b 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -354,8 +354,8 @@ def pad_to_square( ) max_dim = max(height, width) - paste_x_left = (max_dim - width) // 2 - paste_y_left = (max_dim - height) // 2 + paste_x_left = round((max_dim - width) / 2) + paste_y_left = round((max_dim - height) / 2) paste_x_right = max_dim - width - paste_x_left paste_y_right = max_dim - height - paste_y_left padded_images = tvF.pad( diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 9003566eb31f..8f41bc730590 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -16,6 +16,7 @@ import torch from huggingface_hub.dataclasses import strict from torch import nn +from torchvision.transforms.v2 import functional as tvF from ... import initialization as init from ...cache_utils import Cache, DynamicCache @@ -84,6 +85,48 @@ class UnlimitedOcrImageProcessor(DeepseekOcr2ImageProcessor): max_patches = 32 model_input_names = ["pixel_values", "num_local_patches", "image_spatial_crop"] + def pad_to_square( + self, + images: "torch.Tensor", + background_color: int | tuple[int, int, int] = 0, + ) -> "torch.Tensor": + """ + Pads an image to a square based on the longest edge. + + Args: + images (`torch.Tensor`): + The images to pad. Shape: (batch_size, num_channels, height, width) or (num_channels, height, width). + background_color (`int` or `tuple[int, int, int]`, *optional*, defaults to 0): + The color to use for the padding. Can be an integer for single channel or a + tuple of integers representing for multi-channel images. If passed as integer + in multi-channel mode, it will default to `0` in subsequent channels. + Returns: + `torch.Tensor`: The padded images. + """ + height, width = images.shape[-2:] + + if height == width: + return images + + num_channels = images.shape[1] if len(images.shape) == 4 else images.shape[0] + if isinstance(background_color, int): + background_color = [background_color] + [0] * (num_channels - 1) + elif len(background_color) != num_channels: + raise ValueError( + f"background_color must have no more than {num_channels} elements to match the number of channels" + ) + + max_dim = max(height, width) + paste_x_left = round((max_dim - width) / 2) + paste_y_left = round((max_dim - height) / 2) + paste_x_right = max_dim - width - paste_x_left + paste_y_right = max_dim - height - paste_y_left + padded_images = tvF.pad( + images, padding=[paste_x_left, paste_y_left, paste_x_right, paste_y_right], fill=background_color + ) + + return padded_images + def _preprocess( self, images: list["torch.Tensor"], From 07782c901695928467b05edb2fc109239791af83 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 12:05:45 +0000 Subject: [PATCH 033/195] update tile handling --- .../models/unlimited_ocr/image_processing_unlimited_ocr.py | 2 +- src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index 9ffceb8a6e7b..4f47fb052425 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -257,7 +257,7 @@ def _preprocess( flat_local_list = [patch for item in ordered_local if item is not None for patch in item] # --- Global view (batched by shape group) --- - global_target_size = size.height if crop_to_patches else tile_size + global_target_size = size.height grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) processed_global_grouped = {} diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 8f41bc730590..a2030b3273b9 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -179,7 +179,7 @@ def _preprocess( flat_local_list = [patch for item in ordered_local if item is not None for patch in item] # --- Global view (batched by shape group) --- - global_target_size = size.height if crop_to_patches else tile_size + global_target_size = size.height grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) processed_global_grouped = {} From 64ab0cd1a028f8ed1013e5c017a6d03b21381683 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 12:05:45 +0000 Subject: [PATCH 034/195] update tile handling --- .../image_processing_unlimited_ocr.py | 14 +++++++------- .../models/unlimited_ocr/modular_unlimited_ocr.py | 14 +++++++------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index 9ffceb8a6e7b..1a2c90d6f932 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -257,7 +257,7 @@ def _preprocess( flat_local_list = [patch for item in ordered_local if item is not None for patch in item] # --- Global view (batched by shape group) --- - global_target_size = size.height if crop_to_patches else tile_size + global_target_size = size.height grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) processed_global_grouped = {} @@ -278,8 +278,6 @@ def _preprocess( if flat_local_list: data["pixel_values_local"] = flat_local_list - batch_feature = BatchFeature(data=data, tensor_type=return_tensors) - # Compute per-image spatial crop grid and local-patch counts. image_spatial_crop = [] num_local_patches = [] @@ -296,11 +294,13 @@ def _preprocess( image_spatial_crop.append([num_columns, num_rows]) # TODO: rename - batch_feature["image_spatial_crop"] = torch.tensor( - image_spatial_crop, dtype=torch.long, device=images[0].device + data["image_spatial_crop"] = image_spatial_crop + data["num_local_patches"] = num_local_patches + + return BatchFeature( + data=data, + tensor_type=return_tensors, ) - batch_feature["num_local_patches"] = torch.tensor(num_local_patches, dtype=torch.long, device=images[0].device) - return batch_feature def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None) -> int: """ diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 8f41bc730590..e5afe98eb140 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -179,7 +179,7 @@ def _preprocess( flat_local_list = [patch for item in ordered_local if item is not None for patch in item] # --- Global view (batched by shape group) --- - global_target_size = size.height if crop_to_patches else tile_size + global_target_size = size.height grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) processed_global_grouped = {} @@ -200,8 +200,6 @@ def _preprocess( if flat_local_list: data["pixel_values_local"] = flat_local_list - batch_feature = BatchFeature(data=data, tensor_type=return_tensors) - # Compute per-image spatial crop grid and local-patch counts. image_spatial_crop = [] num_local_patches = [] @@ -218,11 +216,13 @@ def _preprocess( image_spatial_crop.append([num_columns, num_rows]) # TODO: rename - batch_feature["image_spatial_crop"] = torch.tensor( - image_spatial_crop, dtype=torch.long, device=images[0].device + data["image_spatial_crop"] = image_spatial_crop + data["num_local_patches"] = num_local_patches + + return BatchFeature( + data=data, + tensor_type=return_tensors, ) - batch_feature["num_local_patches"] = torch.tensor(num_local_patches, dtype=torch.long, device=images[0].device) - return batch_feature class UnlimitedOcrProcessorKwargs(DeepseekOcr2ProcessorKwargs, total=False): From 0ec1eb0ff12c5e443075f6283586dbcb254c2cfb Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 12:10:19 +0000 Subject: [PATCH 035/195] update multi page test --- .../models/unlimited_ocr/test_modeling_unlimited_ocr.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index e63810408828..cad1a51eee3c 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -255,12 +255,17 @@ def test_small_model_integration_test_document_parsing_batched(self): @require_torch_accelerator def test_small_model_integration_test_multi_page_document_parsing(self): model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, device_map=torch_device).eval() - image = load_image( + image1 = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" ) ) - inputs = self.processor(images=image, text="\nMulti page parsing.", return_tensors="pt").to( + image2 = load_image( + url_to_local_path( + "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png" + ) + ) + inputs = self.processor(images=[image1, image2], text="\nMulti page parsing.", crop_to_patches=False, return_tensors="pt").to( model.device ) From 5b40b2a2823f97da55e85b4af7a41eb32de02ee4 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 14:12:15 +0200 Subject: [PATCH 036/195] update docs --- docs/source/en/model_doc/unlimited_ocr.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index e32d2f988c2c..df0866d4e2be 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -33,7 +33,7 @@ This model was contributed by [guarin](https://huggingface.co/guarin). The original code can be found [here](https://github.com/baidu/Unlimited-OCR). > [!TIP] -> Unlimited-OCR supports two inference configurations: the default "gundam" mode uses 640x640 tiles with dynamic cropping for high-resolution documents, and "base" mode uses a single 1024x1024 global view for standard-resolution inputs. Gundam mode is enabled by default via the image processor (`crop_to_patches=True`, `tile_size=640`). +> Unlimited-OCR supports two inference configurations: the default "gundam" mode uses 640x640 tiles with dynamic cropping for high-resolution documents, and "base" mode uses a single 1024x1024 global view for standard-resolution inputs. Gundam mode is enabled by default via the image processor (`crop_to_patches=True`). > [!TIP] > For multi-page documents, pass all page images together with one `` token per page in the text prompt. The model processes all pages jointly within a single context window. From 975c5b2125b1a9e299bdff4f586793d07f689c26 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 24 Jun 2026 15:51:13 +0000 Subject: [PATCH 037/195] add cache --- .../configuration_unlimited_ocr.py | 4 +- .../unlimited_ocr/modeling_unlimited_ocr.py | 165 +++++++++++++++-- .../unlimited_ocr/modular_unlimited_ocr.py | 175 +++++++++++++++--- 3 files changed, 307 insertions(+), 37 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index e095b296ae01..44c6e74ae479 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -225,8 +225,8 @@ class UnlimitedOcrTextConfig(PreTrainedConfig): def __post_init__(self, **kwargs): if self.layer_types is None: - # Full attention on every layer keeps the KV cache complete: the sliding window is realized as a - # mask over generated tokens, while the image/prompt prefill is always retained. + # Full attention on every layer keeps the KV cache complete; the ring-buffer logic in + # UnlimitedOcrDynamicCache limits generated tokens to `sliding_window` entries at runtime. self.layer_types = ["full_attention"] * self.num_hidden_layers if self.mlp_layer_types is None: # Some configs may use `first_k_dense_replace` instead of `layer_types`/`mlp_layer_types` diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 0a94901cdc67..e96d727637f3 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -31,13 +31,14 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin +from ...generation.utils import GenerationMode from ...integrations import ( use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func, ) -from ...masking_utils import and_masks, causal_mask_function, create_causal_mask, create_sliding_window_causal_mask +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 @@ -67,6 +68,109 @@ logger = logging.get_logger(__name__) +class UnlimitedOcrDynamicCache(DynamicCache): + """DynamicCache with a fixed-size ring buffer for generated tokens. + + All prefill (image + prompt) tokens are kept intact; generated tokens are stored in a + ``sliding_window``-sized ring buffer that is overwritten in place. This mirrors the reference + model's sliding-window attention (each generated token attends to all prefill tokens plus the + last ``sliding_window`` generated tokens) and keeps the cached tensors at a constant + ``prefill + sliding_window`` length from the very first decode step onwards. + + The constant length is what makes generation fast: the cuDNN SDPA backend re-plans its kernel + on every distinct key/value sequence length, so a cache that grows by one each step (the plain + `DynamicCache`) stalls ~400ms per step on CPU re-planning. With a constant length cuDNN plans + once and reuses it. The not-yet-filled ring slots during the first ``sliding_window`` steps are + hidden by the attention mask built in `UnlimitedOcrTextModel._create_attention_mask`. + + `get_seq_length()` returns the *logical* cumulative token count so `generate()` computes correct + position ids; the physically stored length is bounded at ``prefill + sliding_window``. + """ + + def __init__(self, config=None): + super().__init__(config=config) + self._sliding_window = getattr(config, "sliding_window", None) if config is not None else None + self._ring_prefill: dict[int, int] = {} + self._ring_pos: dict[int, int] = {} + self._n_gen: dict[int, int] = {} + self._cum_len: dict[int, int] = {} + + def update(self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, *args, **kwargs): + if self._sliding_window is None: + return super().update(key_states, value_states, layer_idx, *args, **kwargs) + + layer = self.layers[layer_idx] + if not layer.is_initialized: + layer.lazy_initialization(key_states, value_states) + + seq_len = key_states.shape[-2] + self._cum_len[layer_idx] = self._cum_len.get(layer_idx, 0) + seq_len + + if layer_idx not in self._ring_prefill: + if seq_len > 1: + # Prefill: accumulate normally (single forward, so no re-planning concern). + layer.keys = torch.cat([layer.keys, key_states], dim=-2) + layer.values = torch.cat([layer.values, value_states], dim=-2) + return layer.keys, layer.values + # First decode step: record the prefill boundary and pre-allocate the whole ring up front + # so the cached length is constant (prefill + sliding_window) for every subsequent step. + self._ring_prefill[layer_idx] = layer.keys.shape[-2] + self._ring_pos[layer_idx] = 0 + self._n_gen[layer_idx] = 0 + ring = torch.zeros( + *layer.keys.shape[:-2], + self._sliding_window, + layer.keys.shape[-1], + dtype=layer.keys.dtype, + device=layer.keys.device, + ) + layer.keys = torch.cat([layer.keys, ring], dim=-2) + layer.values = torch.cat([layer.values, torch.zeros_like(ring)], dim=-2) + + pfl = self._ring_prefill[layer_idx] + rpos = self._ring_pos[layer_idx] + for t in range(seq_len): + slot = pfl + rpos + layer.keys[..., slot : slot + 1, :] = key_states[..., t : t + 1, :] + layer.values[..., slot : slot + 1, :] = value_states[..., t : t + 1, :] + rpos = (rpos + 1) % self._sliding_window + self._n_gen[layer_idx] = min(self._n_gen[layer_idx] + 1, self._sliding_window) + self._ring_pos[layer_idx] = rpos + + # Always return the full fixed-size buffer; unfilled ring slots are masked out by the model. + return layer.keys, layer.values + + def decode_kv_layout(self, layer_idx: int = 0) -> tuple[int, int] | None: + """`(kv_total, kv_valid)` for the upcoming decode step, or `None` if the ring is inactive. + + `kv_total` is the constant physical length (`prefill + sliding_window`) and `kv_valid` is the + number of leading entries that hold real keys/values (`prefill + filled ring slots`) once this + step's token has been written. Slots in `[kv_valid, kv_total)` are unfilled and must be masked. + """ + if self._sliding_window is None: + return None + if layer_idx in self._ring_prefill: + prefill_len = self._ring_prefill[layer_idx] + n_gen_after = min(self._n_gen[layer_idx] + 1, self._sliding_window) + elif layer_idx < len(self.layers) and self.layers[layer_idx].is_initialized: + # First decode step: the ring has not been allocated yet, so the current physical length is + # exactly the prefill length and this step writes the first generated token. + prefill_len = self.layers[layer_idx].keys.shape[-2] + n_gen_after = 1 + else: + return None + return prefill_len + self._sliding_window, prefill_len + n_gen_after + + def get_seq_length(self, layer_idx: int = 0) -> int: + return self._cum_len.get(layer_idx, 0) + + def get_mask_sizes(self, query_length: int, layer_idx: int) -> tuple[int, int]: + cum = self._cum_len.get(layer_idx, 0) + if layer_idx not in self._ring_prefill: + return cum + query_length, 0 + return self._ring_prefill[layer_idx] + self._sliding_window, 0 + + @dataclass class UnlimitedOcrModelOutputWithPooling(BaseModelOutputWithPooling): """ @@ -1574,7 +1678,7 @@ def forward( inputs_embeds = self.embed_tokens(input_ids) if use_cache and past_key_values is None: - past_key_values = DynamicCache(config=self.config) + past_key_values = UnlimitedOcrDynamicCache(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 @@ -1614,19 +1718,26 @@ def _create_attention_mask(self, inputs_embeds, attention_mask, past_key_values, if self.config.sliding_window is None: return create_causal_mask(**mask_kwargs) - prefill_length = getattr(past_key_values, "prefill_length", None) - if prefill_length is None: - prefill_length = torch.tensor(inputs_embeds.shape[1], device=inputs_embeds.device) - if past_key_values is not None: - past_key_values.prefill_length = prefill_length - - def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): - return kv_idx < prefill_length + # Ring-buffer decode (single query token): the cache returns a constant-length buffer of + # [prefill tokens] + [sliding_window ring slots]. Attention is order-agnostic over keys (RoPE + # is baked into the cached keys), so the query may attend to every *filled* slot β€” that is all + # prefill tokens plus the generated tokens currently in the ring, i.e. the model's sliding + # window. We only need to hide the not-yet-filled ring slots during the first window of steps. + # Keeping the mask shape constant (even when all slots are valid) lets the cuDNN SDPA backend + # plan its kernel once instead of re-planning every step. + layout = past_key_values.decode_kv_layout() if isinstance(past_key_values, UnlimitedOcrDynamicCache) else None + if inputs_embeds.shape[1] == 1 and layout is not None: + kv_total, kv_valid = layout + min_value = torch.finfo(inputs_embeds.dtype).min + mask = torch.zeros( + inputs_embeds.shape[0], 1, 1, kv_total, dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + if kv_valid < kv_total: + mask[..., kv_valid:] = min_value + return mask - return create_sliding_window_causal_mask( - **mask_kwargs, - or_mask_function=and_masks(prefill_overlay, causal_mask_function), - ) + # Prefill: standard full causal mask so image tokens can attend to each other correctly. + return create_causal_mask(**mask_kwargs) @auto_docstring( @@ -1952,6 +2063,32 @@ def prepare_inputs_for_generation( return model_inputs + def _prepare_cache_for_generation( + self, generation_config, model_kwargs, generation_mode, batch_size, max_cache_length + ): + # `generate()` would otherwise build a plain `DynamicCache` whose KV grows unbounded, making decode + # O(nΒ²). For the default dynamic-cache case we instead use `UnlimitedOcrDynamicCache`, which keeps all + # prefill tokens but bounds generated tokens to `sliding_window` entries via a ring buffer (O(n) total). + # Caches that need rollback (assisted/contrastive) or an explicit `cache_implementation` defer to super. + uses_default_dynamic_cache = ( + model_kwargs.get("past_key_values") is None + and generation_config.use_cache + and generation_config.cache_implementation is None + and generation_mode not in (GenerationMode.ASSISTED_GENERATION, GenerationMode.CONTRASTIVE_SEARCH) + and self._supports_default_dynamic_cache() + ) + if uses_default_dynamic_cache: + model_kwargs["past_key_values"] = UnlimitedOcrDynamicCache( + config=self.config.get_text_config(decoder=True) + ) + return + # Defer every other case (static/quantized/offloaded caches, encoder-decoder, no cache, ...) + # to the base implementation. Called on `GenerationMixin` directly rather than via `super()` + # so the modular converter doesn't try to inline a method that isn't in the modeling lineage. + return super()._prepare_cache_for_generation( + generation_config, model_kwargs, generation_mode, batch_size, max_cache_length + ) + __all__ = [ "UnlimitedOcrForConditionalGeneration", diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index e5afe98eb140..4650fbcbff5d 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -22,14 +22,11 @@ from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PretrainedConfig from ...feature_extraction_utils import BatchFeature +from ...generation import GenerationMixin +from ...generation.utils import GenerationMode from ...image_transforms import group_images_by_shape, reorder_images from ...image_utils import ImageInput, PILImageResampling, SizeDict -from ...masking_utils import ( - and_masks, - causal_mask_function, - create_causal_mask, - create_sliding_window_causal_mask, -) +from ...masking_utils import create_causal_mask from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast from ...processing_utils import Unpack from ...tokenization_utils_base import PreTokenizedInput, TextInput @@ -350,6 +347,109 @@ class UnlimitedOcrVisionEncoderConfig(CLIPVisionConfig): patch_size: int | list[int] | tuple[int, int] | None = 14 +class UnlimitedOcrDynamicCache(DynamicCache): + """DynamicCache with a fixed-size ring buffer for generated tokens. + + All prefill (image + prompt) tokens are kept intact; generated tokens are stored in a + ``sliding_window``-sized ring buffer that is overwritten in place. This mirrors the reference + model's sliding-window attention (each generated token attends to all prefill tokens plus the + last ``sliding_window`` generated tokens) and keeps the cached tensors at a constant + ``prefill + sliding_window`` length from the very first decode step onwards. + + The constant length is what makes generation fast: the cuDNN SDPA backend re-plans its kernel + on every distinct key/value sequence length, so a cache that grows by one each step (the plain + `DynamicCache`) stalls ~400ms per step on CPU re-planning. With a constant length cuDNN plans + once and reuses it. The not-yet-filled ring slots during the first ``sliding_window`` steps are + hidden by the attention mask built in `UnlimitedOcrTextModel._create_attention_mask`. + + `get_seq_length()` returns the *logical* cumulative token count so `generate()` computes correct + position ids; the physically stored length is bounded at ``prefill + sliding_window``. + """ + + def __init__(self, config=None): + super().__init__(config=config) + self._sliding_window = getattr(config, "sliding_window", None) if config is not None else None + self._ring_prefill: dict[int, int] = {} + self._ring_pos: dict[int, int] = {} + self._n_gen: dict[int, int] = {} + self._cum_len: dict[int, int] = {} + + def update(self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, *args, **kwargs): + if self._sliding_window is None: + return super().update(key_states, value_states, layer_idx, *args, **kwargs) + + layer = self.layers[layer_idx] + if not layer.is_initialized: + layer.lazy_initialization(key_states, value_states) + + seq_len = key_states.shape[-2] + self._cum_len[layer_idx] = self._cum_len.get(layer_idx, 0) + seq_len + + if layer_idx not in self._ring_prefill: + if seq_len > 1: + # Prefill: accumulate normally (single forward, so no re-planning concern). + layer.keys = torch.cat([layer.keys, key_states], dim=-2) + layer.values = torch.cat([layer.values, value_states], dim=-2) + return layer.keys, layer.values + # First decode step: record the prefill boundary and pre-allocate the whole ring up front + # so the cached length is constant (prefill + sliding_window) for every subsequent step. + self._ring_prefill[layer_idx] = layer.keys.shape[-2] + self._ring_pos[layer_idx] = 0 + self._n_gen[layer_idx] = 0 + ring = torch.zeros( + *layer.keys.shape[:-2], + self._sliding_window, + layer.keys.shape[-1], + dtype=layer.keys.dtype, + device=layer.keys.device, + ) + layer.keys = torch.cat([layer.keys, ring], dim=-2) + layer.values = torch.cat([layer.values, torch.zeros_like(ring)], dim=-2) + + pfl = self._ring_prefill[layer_idx] + rpos = self._ring_pos[layer_idx] + for t in range(seq_len): + slot = pfl + rpos + layer.keys[..., slot : slot + 1, :] = key_states[..., t : t + 1, :] + layer.values[..., slot : slot + 1, :] = value_states[..., t : t + 1, :] + rpos = (rpos + 1) % self._sliding_window + self._n_gen[layer_idx] = min(self._n_gen[layer_idx] + 1, self._sliding_window) + self._ring_pos[layer_idx] = rpos + + # Always return the full fixed-size buffer; unfilled ring slots are masked out by the model. + return layer.keys, layer.values + + def decode_kv_layout(self, layer_idx: int = 0) -> tuple[int, int] | None: + """`(kv_total, kv_valid)` for the upcoming decode step, or `None` if the ring is inactive. + + `kv_total` is the constant physical length (`prefill + sliding_window`) and `kv_valid` is the + number of leading entries that hold real keys/values (`prefill + filled ring slots`) once this + step's token has been written. Slots in `[kv_valid, kv_total)` are unfilled and must be masked. + """ + if self._sliding_window is None: + return None + if layer_idx in self._ring_prefill: + prefill_len = self._ring_prefill[layer_idx] + n_gen_after = min(self._n_gen[layer_idx] + 1, self._sliding_window) + elif layer_idx < len(self.layers) and self.layers[layer_idx].is_initialized: + # First decode step: the ring has not been allocated yet, so the current physical length is + # exactly the prefill length and this step writes the first generated token. + prefill_len = self.layers[layer_idx].keys.shape[-2] + n_gen_after = 1 + else: + return None + return prefill_len + self._sliding_window, prefill_len + n_gen_after + + def get_seq_length(self, layer_idx: int = 0) -> int: + return self._cum_len.get(layer_idx, 0) + + def get_mask_sizes(self, query_length: int, layer_idx: int) -> tuple[int, int]: + cum = self._cum_len.get(layer_idx, 0) + if layer_idx not in self._ring_prefill: + return cum + query_length, 0 + return self._ring_prefill[layer_idx] + self._sliding_window, 0 + + @auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict class UnlimitedOcrVisionConfig(DeepseekOcr2VisionConfig): @@ -414,8 +514,8 @@ class UnlimitedOcrTextConfig(DeepseekOcr2TextConfig): def __post_init__(self, **kwargs): if self.layer_types is None: - # Full attention on every layer keeps the KV cache complete: the sliding window is realized as a - # mask over generated tokens, while the image/prompt prefill is always retained. + # Full attention on every layer keeps the KV cache complete; the ring-buffer logic in + # UnlimitedOcrDynamicCache limits generated tokens to `sliding_window` entries at runtime. self.layer_types = ["full_attention"] * self.num_hidden_layers if self.mlp_layer_types is None: # Some configs may use `first_k_dense_replace` instead of `layer_types`/`mlp_layer_types` @@ -691,19 +791,26 @@ def _create_attention_mask(self, inputs_embeds, attention_mask, past_key_values, if self.config.sliding_window is None: return create_causal_mask(**mask_kwargs) - prefill_length = getattr(past_key_values, "prefill_length", None) - if prefill_length is None: - prefill_length = torch.tensor(inputs_embeds.shape[1], device=inputs_embeds.device) - if past_key_values is not None: - past_key_values.prefill_length = prefill_length - - def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): - return kv_idx < prefill_length + # Ring-buffer decode (single query token): the cache returns a constant-length buffer of + # [prefill tokens] + [sliding_window ring slots]. Attention is order-agnostic over keys (RoPE + # is baked into the cached keys), so the query may attend to every *filled* slot β€” that is all + # prefill tokens plus the generated tokens currently in the ring, i.e. the model's sliding + # window. We only need to hide the not-yet-filled ring slots during the first window of steps. + # Keeping the mask shape constant (even when all slots are valid) lets the cuDNN SDPA backend + # plan its kernel once instead of re-planning every step. + layout = past_key_values.decode_kv_layout() if isinstance(past_key_values, UnlimitedOcrDynamicCache) else None + if inputs_embeds.shape[1] == 1 and layout is not None: + kv_total, kv_valid = layout + min_value = torch.finfo(inputs_embeds.dtype).min + mask = torch.zeros( + inputs_embeds.shape[0], 1, 1, kv_total, dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + if kv_valid < kv_total: + mask[..., kv_valid:] = min_value + return mask - return create_sliding_window_causal_mask( - **mask_kwargs, - or_mask_function=and_masks(prefill_overlay, causal_mask_function), - ) + # Prefill: standard full causal mask so image tokens can attend to each other correctly. + return create_causal_mask(**mask_kwargs) @merge_with_config_defaults @capture_outputs @@ -725,7 +832,7 @@ def forward( inputs_embeds = self.embed_tokens(input_ids) if use_cache and past_key_values is None: - past_key_values = DynamicCache(config=self.config) + past_key_values = UnlimitedOcrDynamicCache(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 @@ -998,6 +1105,32 @@ def prepare_inputs_for_generation( return model_inputs + def _prepare_cache_for_generation( + self, generation_config, model_kwargs, generation_mode, batch_size, max_cache_length + ): + # `generate()` would otherwise build a plain `DynamicCache` whose KV grows unbounded, making decode + # O(nΒ²). For the default dynamic-cache case we instead use `UnlimitedOcrDynamicCache`, which keeps all + # prefill tokens but bounds generated tokens to `sliding_window` entries via a ring buffer (O(n) total). + # Caches that need rollback (assisted/contrastive) or an explicit `cache_implementation` defer to super. + uses_default_dynamic_cache = ( + model_kwargs.get("past_key_values") is None + and generation_config.use_cache + and generation_config.cache_implementation is None + and generation_mode not in (GenerationMode.ASSISTED_GENERATION, GenerationMode.CONTRASTIVE_SEARCH) + and self._supports_default_dynamic_cache() + ) + if uses_default_dynamic_cache: + model_kwargs["past_key_values"] = UnlimitedOcrDynamicCache( + config=self.config.get_text_config(decoder=True) + ) + return + # Defer every other case (static/quantized/offloaded caches, encoder-decoder, no cache, ...) + # to the base implementation. Called on `GenerationMixin` directly rather than via `super()` + # so the modular converter doesn't try to inline a method that isn't in the modeling lineage. + return GenerationMixin._prepare_cache_for_generation( + self, generation_config, model_kwargs, generation_mode, batch_size, max_cache_length + ) + __all__ = [ "UnlimitedOcrConfig", From 2d4d952ed8f9e26fd74a0bc8f323a7da2fed350a Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 07:56:59 +0000 Subject: [PATCH 038/195] update test --- tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index cad1a51eee3c..ba563862b7bf 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -275,9 +275,7 @@ def test_small_model_integration_test_multi_page_document_parsing(self): EXPECTED_DECODED_TEXT = Expectations( { - ("cuda", None): [ - "image [382, 87, 489, 174]\n", - ], + ("cuda", None): "image [382, 87, 489, 174]\n", } ).get_expectation() self.assertEqual(decoded, EXPECTED_DECODED_TEXT) From 927be912f39fad8c2b168d93d365be3241e0ef23 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 07:57:08 +0000 Subject: [PATCH 039/195] add better error message --- .../models/unlimited_ocr/modular_unlimited_ocr.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 4650fbcbff5d..0148c742d6fa 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -237,6 +237,15 @@ def _expand_image_tokens( image_spatial_crop: torch.Tensor, num_local_patches: list[int] | torch.Tensor, ) -> list[str]: + num_images = len(image_spatial_crop) + total_image_tokens = sum(t.count(self.image_token) for t in text) + if total_image_tokens != num_images: + raise ValueError( + f"Number of `{self.image_token}` tokens in the text ({total_image_tokens}) does not match the " + f"number of images passed ({num_images}). Use one `{self.image_token}` placeholder per image, " + f"e.g. `'' * len(images) + '\\nMulti page parsing.'`" + ) + size = self.image_processor.size["height"] tile_size = self.image_processor.tile_size From 50f717a0f08b9d7d6895a3e7341b35f507cdfe42 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 10:14:26 +0200 Subject: [PATCH 040/195] add reference_sliding_attention --- docs/source/en/model_doc/unlimited_ocr.md | 2 +- src/transformers/configuration_utils.py | 1 + .../configuration_unlimited_ocr.py | 6 +- .../unlimited_ocr/modeling_unlimited_ocr.py | 256 ++++++++--------- .../unlimited_ocr/modular_unlimited_ocr.py | 261 +++++++++--------- .../test_modeling_unlimited_ocr.py | 19 +- 6 files changed, 265 insertions(+), 280 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index df0866d4e2be..ecbc06dc44b9 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -16,7 +16,7 @@ limitations under the License. ⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer. --> -*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-06-23.* +*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-06-25.* # UnlimitedOcr diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py index 2b6fdce06e28..acd0f5b2c2db 100755 --- a/src/transformers/configuration_utils.py +++ b/src/transformers/configuration_utils.py @@ -75,6 +75,7 @@ "hybrid", # for layers that have both mamba and attention in zamba and zamba2 "moe", # for nemotron_h, which uses either attention, mamba or moe "deepseek_sparse_attention", # for models with DSA indexer (GLM MoE DSA, DeepSeek V32) + "reference_sliding_attention", # UnlimitedOcr: full attention over prefill + ring-buffered sliding window over generated tokens ) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index 44c6e74ae479..f060974056d3 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -225,9 +225,9 @@ class UnlimitedOcrTextConfig(PreTrainedConfig): def __post_init__(self, **kwargs): if self.layer_types is None: - # Full attention on every layer keeps the KV cache complete; the ring-buffer logic in - # UnlimitedOcrDynamicCache limits generated tokens to `sliding_window` entries at runtime. - self.layer_types = ["full_attention"] * self.num_hidden_layers + self.layer_types = [ + "full_attention" if self.sliding_window is None else "reference_sliding_attention" + ] * self.num_hidden_layers if self.mlp_layer_types is None: # Some configs may use `first_k_dense_replace` instead of `layer_types`/`mlp_layer_types` first_k_dense_replace = kwargs.pop("first_k_dense_replace", 1) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index e96d727637f3..b2efc7a96c7f 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -29,7 +29,8 @@ from ... import initialization as init from ...activations import ACT2FN -from ...cache_utils import Cache, DynamicCache +from ...cache_utils import Cache, DynamicCache, DynamicSlidingWindowLayer +from ...configuration_utils import PretrainedConfig from ...generation import GenerationMixin from ...generation.utils import GenerationMode from ...integrations import ( @@ -68,109 +69,6 @@ logger = logging.get_logger(__name__) -class UnlimitedOcrDynamicCache(DynamicCache): - """DynamicCache with a fixed-size ring buffer for generated tokens. - - All prefill (image + prompt) tokens are kept intact; generated tokens are stored in a - ``sliding_window``-sized ring buffer that is overwritten in place. This mirrors the reference - model's sliding-window attention (each generated token attends to all prefill tokens plus the - last ``sliding_window`` generated tokens) and keeps the cached tensors at a constant - ``prefill + sliding_window`` length from the very first decode step onwards. - - The constant length is what makes generation fast: the cuDNN SDPA backend re-plans its kernel - on every distinct key/value sequence length, so a cache that grows by one each step (the plain - `DynamicCache`) stalls ~400ms per step on CPU re-planning. With a constant length cuDNN plans - once and reuses it. The not-yet-filled ring slots during the first ``sliding_window`` steps are - hidden by the attention mask built in `UnlimitedOcrTextModel._create_attention_mask`. - - `get_seq_length()` returns the *logical* cumulative token count so `generate()` computes correct - position ids; the physically stored length is bounded at ``prefill + sliding_window``. - """ - - def __init__(self, config=None): - super().__init__(config=config) - self._sliding_window = getattr(config, "sliding_window", None) if config is not None else None - self._ring_prefill: dict[int, int] = {} - self._ring_pos: dict[int, int] = {} - self._n_gen: dict[int, int] = {} - self._cum_len: dict[int, int] = {} - - def update(self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, *args, **kwargs): - if self._sliding_window is None: - return super().update(key_states, value_states, layer_idx, *args, **kwargs) - - layer = self.layers[layer_idx] - if not layer.is_initialized: - layer.lazy_initialization(key_states, value_states) - - seq_len = key_states.shape[-2] - self._cum_len[layer_idx] = self._cum_len.get(layer_idx, 0) + seq_len - - if layer_idx not in self._ring_prefill: - if seq_len > 1: - # Prefill: accumulate normally (single forward, so no re-planning concern). - layer.keys = torch.cat([layer.keys, key_states], dim=-2) - layer.values = torch.cat([layer.values, value_states], dim=-2) - return layer.keys, layer.values - # First decode step: record the prefill boundary and pre-allocate the whole ring up front - # so the cached length is constant (prefill + sliding_window) for every subsequent step. - self._ring_prefill[layer_idx] = layer.keys.shape[-2] - self._ring_pos[layer_idx] = 0 - self._n_gen[layer_idx] = 0 - ring = torch.zeros( - *layer.keys.shape[:-2], - self._sliding_window, - layer.keys.shape[-1], - dtype=layer.keys.dtype, - device=layer.keys.device, - ) - layer.keys = torch.cat([layer.keys, ring], dim=-2) - layer.values = torch.cat([layer.values, torch.zeros_like(ring)], dim=-2) - - pfl = self._ring_prefill[layer_idx] - rpos = self._ring_pos[layer_idx] - for t in range(seq_len): - slot = pfl + rpos - layer.keys[..., slot : slot + 1, :] = key_states[..., t : t + 1, :] - layer.values[..., slot : slot + 1, :] = value_states[..., t : t + 1, :] - rpos = (rpos + 1) % self._sliding_window - self._n_gen[layer_idx] = min(self._n_gen[layer_idx] + 1, self._sliding_window) - self._ring_pos[layer_idx] = rpos - - # Always return the full fixed-size buffer; unfilled ring slots are masked out by the model. - return layer.keys, layer.values - - def decode_kv_layout(self, layer_idx: int = 0) -> tuple[int, int] | None: - """`(kv_total, kv_valid)` for the upcoming decode step, or `None` if the ring is inactive. - - `kv_total` is the constant physical length (`prefill + sliding_window`) and `kv_valid` is the - number of leading entries that hold real keys/values (`prefill + filled ring slots`) once this - step's token has been written. Slots in `[kv_valid, kv_total)` are unfilled and must be masked. - """ - if self._sliding_window is None: - return None - if layer_idx in self._ring_prefill: - prefill_len = self._ring_prefill[layer_idx] - n_gen_after = min(self._n_gen[layer_idx] + 1, self._sliding_window) - elif layer_idx < len(self.layers) and self.layers[layer_idx].is_initialized: - # First decode step: the ring has not been allocated yet, so the current physical length is - # exactly the prefill length and this step writes the first generated token. - prefill_len = self.layers[layer_idx].keys.shape[-2] - n_gen_after = 1 - else: - return None - return prefill_len + self._sliding_window, prefill_len + n_gen_after - - def get_seq_length(self, layer_idx: int = 0) -> int: - return self._cum_len.get(layer_idx, 0) - - def get_mask_sizes(self, query_length: int, layer_idx: int) -> tuple[int, int]: - cum = self._cum_len.get(layer_idx, 0) - if layer_idx not in self._ring_prefill: - return cum + query_length, 0 - return self._ring_prefill[layer_idx] + self._sliding_window, 0 - - @dataclass class UnlimitedOcrModelOutputWithPooling(BaseModelOutputWithPooling): """ @@ -1309,6 +1207,100 @@ def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwarg ) +class DynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): + """Sliding-window cache layer that keeps all prefill tokens and windows only the generated ones. + + The stock `DynamicSlidingWindowLayer` evicts the oldest tokens once the window fills, which would + discard the image/prompt prefill. Here the prefill (every token cached before the first single-token + decode step) is kept intact and never evicted. Generated tokens are appended right after the prefill + until ``sliding_window`` of them have accumulated; from then on the newest generated token overwrites + the oldest one in place (a ring buffer). This mirrors the reference model's attention: each generated + token attends to all prefill tokens plus the last ``sliding_window`` generated tokens. + + While the window is filling, the cached length grows by one each step (exactly like the stock sliding + layer), so within the first ``sliding_window`` generated tokens the layer is indistinguishable from + full attention. Once the ring is full the cached tensors stay at a constant length + (``prefill + sliding_window``), which lets the cuDNN SDPA backend reuse a single kernel plan across the + (typically long) steady-state decode instead of re-planning on every distinct sequence length. + + `get_seq_length()` (inherited) returns the *logical* cumulative token count so positions are computed + correctly; the physically stored length is bounded at ``prefill + sliding_window``. Because keys/values + are overwritten in place the layer cannot roll back, so cache-rollback generation modes (assisted, + contrastive) use a plain full `DynamicCache` instead (see + `UnlimitedOcrForConditionalGeneration._prepare_cache_for_generation`). + """ + + layer_type = "reference_sliding_attention" + + def __init__(self, config: PretrainedConfig | None = None, sliding_window: int | None = None): + super().__init__(config=config, sliding_window=sliding_window) + self.prefill_length: int | None = None + self.ring_position = 0 + + def update( + self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs + ) -> tuple[torch.Tensor, torch.Tensor]: + if not self.is_initialized: + self.lazy_initialization(key_states, value_states) + + seq_len = key_states.shape[-2] + self.cumulative_length += seq_len + + # Prefill: the first update carries the image/prompt context (more than one token). Keep it intact. + if self.prefill_length is None and seq_len > 1: + self.keys = torch.cat([self.keys, key_states], dim=-2) + self.values = torch.cat([self.values, value_states], dim=-2) + return self.keys, self.values + + # First decode step: freeze the prefill boundary. `self.keys` still holds the (possibly empty) + # prefill here; an empty buffer means generation started without any prompt (BOS only). + if self.prefill_length is None: + self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 + + # TODO: can we remove the for loop? + for token_idx in range(seq_len): + key = key_states[..., token_idx : token_idx + 1, :] + value = value_states[..., token_idx : token_idx + 1, :] + generated_length = self.keys.shape[-2] - self.prefill_length if self.keys.dim() > 1 else 0 + if generated_length < self.sliding_window: + # Window not full yet: append (the buffer grows by one, like the stock sliding layer). + self.keys = torch.cat([self.keys, key], dim=-2) + self.values = torch.cat([self.values, value], dim=-2) + else: + # Window full: overwrite the oldest generated slot in place (ring buffer, constant length). + slot = self.prefill_length + self.ring_position + self.keys[..., slot : slot + 1, :] = key + self.values[..., slot : slot + 1, :] = value + self.ring_position = (self.ring_position + 1) % self.sliding_window + + return self.keys, self.values + + def get_mask_sizes(self, query_length: int) -> tuple[int, int]: + # Full visibility, no sliding offset: every query attends to all cached keys (all prefill plus the + # generated tokens currently held in the ring), so the prefill is never masked out. + if self.prefill_length is None: + return self.cumulative_length + query_length, 0 + return self.decode_kv_length(query_length), 0 + + def decode_kv_length(self, query_length: int = 1) -> int | None: + """Physical length of the cached buffer after the upcoming `update` of `query_length` tokens. + + Returns `None` before the layer is initialized. Used to size the all-visible decode mask, whose + width must match the key/value tensors returned by `update`. + """ + if not self.is_initialized: + return None + if self.prefill_length is not None: + prefill_length = self.prefill_length + generated_before = self.keys.shape[-2] - prefill_length if self.keys.dim() > 1 else 0 + else: + # Before the first decode step the whole buffer is prefill. + prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 + generated_before = 0 + generated_after = min(generated_before + query_length, self.sliding_window) + return prefill_length + generated_after + + class UnlimitedOcrTextRotaryEmbedding(nn.Module): inv_freq: torch.Tensor # fix linting for `register_buffer` @@ -1678,7 +1670,7 @@ def forward( inputs_embeds = self.embed_tokens(input_ids) if use_cache and past_key_values is None: - past_key_values = UnlimitedOcrDynamicCache(config=self.config) + 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 @@ -1718,23 +1710,20 @@ def _create_attention_mask(self, inputs_embeds, attention_mask, past_key_values, if self.config.sliding_window is None: return create_causal_mask(**mask_kwargs) - # Ring-buffer decode (single query token): the cache returns a constant-length buffer of - # [prefill tokens] + [sliding_window ring slots]. Attention is order-agnostic over keys (RoPE - # is baked into the cached keys), so the query may attend to every *filled* slot β€” that is all - # prefill tokens plus the generated tokens currently in the ring, i.e. the model's sliding - # window. We only need to hide the not-yet-filled ring slots during the first window of steps. - # Keeping the mask shape constant (even when all slots are valid) lets the cuDNN SDPA backend - # plan its kernel once instead of re-planning every step. - layout = past_key_values.decode_kv_layout() if isinstance(past_key_values, UnlimitedOcrDynamicCache) else None - if inputs_embeds.shape[1] == 1 and layout is not None: - kv_total, kv_valid = layout - min_value = torch.finfo(inputs_embeds.dtype).min - mask = torch.zeros( - inputs_embeds.shape[0], 1, 1, kv_total, dtype=inputs_embeds.dtype, device=inputs_embeds.device + # Ring-buffer decode (single query token): the cache holds [prefill tokens] + [the generated + # tokens currently in the ring], i.e. all prefill plus the last `sliding_window` generated tokens. + # Attention is order-agnostic over keys (RoPE is baked into the cached keys), so the query attends + # to every cached key β€” a plain causal/sliding mask would wrongly hide the prefill. We therefore + # build an all-visible mask sized to the cache's physical length. + kv_length = None + if past_key_values is not None and past_key_values.layers: + first_layer = past_key_values.layers[0] + if isinstance(first_layer, DynamicReferenceSlidingWindowLayer): + kv_length = first_layer.decode_kv_length(inputs_embeds.shape[1]) + if inputs_embeds.shape[1] == 1 and kv_length is not None: + return torch.zeros( + inputs_embeds.shape[0], 1, 1, kv_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device ) - if kv_valid < kv_total: - mask[..., kv_valid:] = min_value - return mask # Prefill: standard full causal mask so image tokens can attend to each other correctly. return create_causal_mask(**mask_kwargs) @@ -2066,24 +2055,21 @@ def prepare_inputs_for_generation( def _prepare_cache_for_generation( self, generation_config, model_kwargs, generation_mode, batch_size, max_cache_length ): - # `generate()` would otherwise build a plain `DynamicCache` whose KV grows unbounded, making decode - # O(nΒ²). For the default dynamic-cache case we instead use `UnlimitedOcrDynamicCache`, which keeps all - # prefill tokens but bounds generated tokens to `sliding_window` entries via a ring buffer (O(n) total). - # Caches that need rollback (assisted/contrastive) or an explicit `cache_implementation` defer to super. - uses_default_dynamic_cache = ( + # The default path builds a `DynamicCache(config)`, which dispatches `DynamicReferenceSlidingWindowLayer` + # per `config.layer_types` (ring buffer: all prefill kept, generated tokens bounded to `sliding_window`). + # The ring layer returns a fixed-size KV buffer and cannot support the cache rollback / multi-token + # verification that assisted and contrastive generation need, so for those modes we force a plain + # full-attention `DynamicCache` (rollback-capable, full causal attention) instead. + if ( model_kwargs.get("past_key_values") is None and generation_config.use_cache - and generation_config.cache_implementation is None - and generation_mode not in (GenerationMode.ASSISTED_GENERATION, GenerationMode.CONTRASTIVE_SEARCH) + and generation_mode in (GenerationMode.ASSISTED_GENERATION, GenerationMode.CONTRASTIVE_SEARCH) and self._supports_default_dynamic_cache() - ) - if uses_default_dynamic_cache: - model_kwargs["past_key_values"] = UnlimitedOcrDynamicCache( - config=self.config.get_text_config(decoder=True) - ) + ): + model_kwargs["past_key_values"] = DynamicCache() return - # Defer every other case (static/quantized/offloaded caches, encoder-decoder, no cache, ...) - # to the base implementation. Called on `GenerationMixin` directly rather than via `super()` + # Every other case (default dynamic ring cache, static/quantized/offloaded, encoder-decoder, no cache, + # ...) defers to the base implementation. Called on `GenerationMixin` directly rather than via `super()` # so the modular converter doesn't try to inline a method that isn't in the modeling lineage. return super()._prepare_cache_for_generation( generation_config, model_kwargs, generation_mode, batch_size, max_cache_length diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 4650fbcbff5d..4dea1deed50f 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -19,7 +19,7 @@ from torchvision.transforms.v2 import functional as tvF from ... import initialization as init -from ...cache_utils import Cache, DynamicCache +from ...cache_utils import Cache, DynamicCache, DynamicSlidingWindowLayer from ...configuration_utils import PretrainedConfig from ...feature_extraction_utils import BatchFeature from ...generation import GenerationMixin @@ -347,109 +347,6 @@ class UnlimitedOcrVisionEncoderConfig(CLIPVisionConfig): patch_size: int | list[int] | tuple[int, int] | None = 14 -class UnlimitedOcrDynamicCache(DynamicCache): - """DynamicCache with a fixed-size ring buffer for generated tokens. - - All prefill (image + prompt) tokens are kept intact; generated tokens are stored in a - ``sliding_window``-sized ring buffer that is overwritten in place. This mirrors the reference - model's sliding-window attention (each generated token attends to all prefill tokens plus the - last ``sliding_window`` generated tokens) and keeps the cached tensors at a constant - ``prefill + sliding_window`` length from the very first decode step onwards. - - The constant length is what makes generation fast: the cuDNN SDPA backend re-plans its kernel - on every distinct key/value sequence length, so a cache that grows by one each step (the plain - `DynamicCache`) stalls ~400ms per step on CPU re-planning. With a constant length cuDNN plans - once and reuses it. The not-yet-filled ring slots during the first ``sliding_window`` steps are - hidden by the attention mask built in `UnlimitedOcrTextModel._create_attention_mask`. - - `get_seq_length()` returns the *logical* cumulative token count so `generate()` computes correct - position ids; the physically stored length is bounded at ``prefill + sliding_window``. - """ - - def __init__(self, config=None): - super().__init__(config=config) - self._sliding_window = getattr(config, "sliding_window", None) if config is not None else None - self._ring_prefill: dict[int, int] = {} - self._ring_pos: dict[int, int] = {} - self._n_gen: dict[int, int] = {} - self._cum_len: dict[int, int] = {} - - def update(self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, *args, **kwargs): - if self._sliding_window is None: - return super().update(key_states, value_states, layer_idx, *args, **kwargs) - - layer = self.layers[layer_idx] - if not layer.is_initialized: - layer.lazy_initialization(key_states, value_states) - - seq_len = key_states.shape[-2] - self._cum_len[layer_idx] = self._cum_len.get(layer_idx, 0) + seq_len - - if layer_idx not in self._ring_prefill: - if seq_len > 1: - # Prefill: accumulate normally (single forward, so no re-planning concern). - layer.keys = torch.cat([layer.keys, key_states], dim=-2) - layer.values = torch.cat([layer.values, value_states], dim=-2) - return layer.keys, layer.values - # First decode step: record the prefill boundary and pre-allocate the whole ring up front - # so the cached length is constant (prefill + sliding_window) for every subsequent step. - self._ring_prefill[layer_idx] = layer.keys.shape[-2] - self._ring_pos[layer_idx] = 0 - self._n_gen[layer_idx] = 0 - ring = torch.zeros( - *layer.keys.shape[:-2], - self._sliding_window, - layer.keys.shape[-1], - dtype=layer.keys.dtype, - device=layer.keys.device, - ) - layer.keys = torch.cat([layer.keys, ring], dim=-2) - layer.values = torch.cat([layer.values, torch.zeros_like(ring)], dim=-2) - - pfl = self._ring_prefill[layer_idx] - rpos = self._ring_pos[layer_idx] - for t in range(seq_len): - slot = pfl + rpos - layer.keys[..., slot : slot + 1, :] = key_states[..., t : t + 1, :] - layer.values[..., slot : slot + 1, :] = value_states[..., t : t + 1, :] - rpos = (rpos + 1) % self._sliding_window - self._n_gen[layer_idx] = min(self._n_gen[layer_idx] + 1, self._sliding_window) - self._ring_pos[layer_idx] = rpos - - # Always return the full fixed-size buffer; unfilled ring slots are masked out by the model. - return layer.keys, layer.values - - def decode_kv_layout(self, layer_idx: int = 0) -> tuple[int, int] | None: - """`(kv_total, kv_valid)` for the upcoming decode step, or `None` if the ring is inactive. - - `kv_total` is the constant physical length (`prefill + sliding_window`) and `kv_valid` is the - number of leading entries that hold real keys/values (`prefill + filled ring slots`) once this - step's token has been written. Slots in `[kv_valid, kv_total)` are unfilled and must be masked. - """ - if self._sliding_window is None: - return None - if layer_idx in self._ring_prefill: - prefill_len = self._ring_prefill[layer_idx] - n_gen_after = min(self._n_gen[layer_idx] + 1, self._sliding_window) - elif layer_idx < len(self.layers) and self.layers[layer_idx].is_initialized: - # First decode step: the ring has not been allocated yet, so the current physical length is - # exactly the prefill length and this step writes the first generated token. - prefill_len = self.layers[layer_idx].keys.shape[-2] - n_gen_after = 1 - else: - return None - return prefill_len + self._sliding_window, prefill_len + n_gen_after - - def get_seq_length(self, layer_idx: int = 0) -> int: - return self._cum_len.get(layer_idx, 0) - - def get_mask_sizes(self, query_length: int, layer_idx: int) -> tuple[int, int]: - cum = self._cum_len.get(layer_idx, 0) - if layer_idx not in self._ring_prefill: - return cum + query_length, 0 - return self._ring_prefill[layer_idx] + self._sliding_window, 0 - - @auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict class UnlimitedOcrVisionConfig(DeepseekOcr2VisionConfig): @@ -514,9 +411,9 @@ class UnlimitedOcrTextConfig(DeepseekOcr2TextConfig): def __post_init__(self, **kwargs): if self.layer_types is None: - # Full attention on every layer keeps the KV cache complete; the ring-buffer logic in - # UnlimitedOcrDynamicCache limits generated tokens to `sliding_window` entries at runtime. - self.layer_types = ["full_attention"] * self.num_hidden_layers + self.layer_types = [ + "full_attention" if self.sliding_window is None else "reference_sliding_attention" + ] * self.num_hidden_layers if self.mlp_layer_types is None: # Some configs may use `first_k_dense_replace` instead of `layer_types`/`mlp_layer_types` first_k_dense_replace = kwargs.pop("first_k_dense_replace", 1) @@ -747,6 +644,100 @@ def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwarg ) +class DynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): + """Sliding-window cache layer that keeps all prefill tokens and windows only the generated ones. + + The stock `DynamicSlidingWindowLayer` evicts the oldest tokens once the window fills, which would + discard the image/prompt prefill. Here the prefill (every token cached before the first single-token + decode step) is kept intact and never evicted. Generated tokens are appended right after the prefill + until ``sliding_window`` of them have accumulated; from then on the newest generated token overwrites + the oldest one in place (a ring buffer). This mirrors the reference model's attention: each generated + token attends to all prefill tokens plus the last ``sliding_window`` generated tokens. + + While the window is filling, the cached length grows by one each step (exactly like the stock sliding + layer), so within the first ``sliding_window`` generated tokens the layer is indistinguishable from + full attention. Once the ring is full the cached tensors stay at a constant length + (``prefill + sliding_window``), which lets the cuDNN SDPA backend reuse a single kernel plan across the + (typically long) steady-state decode instead of re-planning on every distinct sequence length. + + `get_seq_length()` (inherited) returns the *logical* cumulative token count so positions are computed + correctly; the physically stored length is bounded at ``prefill + sliding_window``. Because keys/values + are overwritten in place the layer cannot roll back, so cache-rollback generation modes (assisted, + contrastive) use a plain full `DynamicCache` instead (see + `UnlimitedOcrForConditionalGeneration._prepare_cache_for_generation`). + """ + + layer_type = "reference_sliding_attention" + + def __init__(self, config: PretrainedConfig | None = None, sliding_window: int | None = None): + super().__init__(config=config, sliding_window=sliding_window) + self.prefill_length: int | None = None + self.ring_position = 0 + + def update( + self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs + ) -> tuple[torch.Tensor, torch.Tensor]: + if not self.is_initialized: + self.lazy_initialization(key_states, value_states) + + seq_len = key_states.shape[-2] + self.cumulative_length += seq_len + + # Prefill: the first update carries the image/prompt context (more than one token). Keep it intact. + if self.prefill_length is None and seq_len > 1: + self.keys = torch.cat([self.keys, key_states], dim=-2) + self.values = torch.cat([self.values, value_states], dim=-2) + return self.keys, self.values + + # First decode step: freeze the prefill boundary. `self.keys` still holds the (possibly empty) + # prefill here; an empty buffer means generation started without any prompt (BOS only). + if self.prefill_length is None: + self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 + + # TODO: can we remove the for loop? + for token_idx in range(seq_len): + key = key_states[..., token_idx : token_idx + 1, :] + value = value_states[..., token_idx : token_idx + 1, :] + generated_length = self.keys.shape[-2] - self.prefill_length if self.keys.dim() > 1 else 0 + if generated_length < self.sliding_window: + # Window not full yet: append (the buffer grows by one, like the stock sliding layer). + self.keys = torch.cat([self.keys, key], dim=-2) + self.values = torch.cat([self.values, value], dim=-2) + else: + # Window full: overwrite the oldest generated slot in place (ring buffer, constant length). + slot = self.prefill_length + self.ring_position + self.keys[..., slot : slot + 1, :] = key + self.values[..., slot : slot + 1, :] = value + self.ring_position = (self.ring_position + 1) % self.sliding_window + + return self.keys, self.values + + def get_mask_sizes(self, query_length: int) -> tuple[int, int]: + # Full visibility, no sliding offset: every query attends to all cached keys (all prefill plus the + # generated tokens currently held in the ring), so the prefill is never masked out. + if self.prefill_length is None: + return self.cumulative_length + query_length, 0 + return self.decode_kv_length(query_length), 0 + + def decode_kv_length(self, query_length: int = 1) -> int | None: + """Physical length of the cached buffer after the upcoming `update` of `query_length` tokens. + + Returns `None` before the layer is initialized. Used to size the all-visible decode mask, whose + width must match the key/value tensors returned by `update`. + """ + if not self.is_initialized: + return None + if self.prefill_length is not None: + prefill_length = self.prefill_length + generated_before = self.keys.shape[-2] - prefill_length if self.keys.dim() > 1 else 0 + else: + # Before the first decode step the whole buffer is prefill. + prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 + generated_before = 0 + generated_after = min(generated_before + query_length, self.sliding_window) + return prefill_length + generated_after + + class UnlimitedOcrTextRotaryEmbedding(DeepseekOcr2TextRotaryEmbedding): pass @@ -791,23 +782,20 @@ def _create_attention_mask(self, inputs_embeds, attention_mask, past_key_values, if self.config.sliding_window is None: return create_causal_mask(**mask_kwargs) - # Ring-buffer decode (single query token): the cache returns a constant-length buffer of - # [prefill tokens] + [sliding_window ring slots]. Attention is order-agnostic over keys (RoPE - # is baked into the cached keys), so the query may attend to every *filled* slot β€” that is all - # prefill tokens plus the generated tokens currently in the ring, i.e. the model's sliding - # window. We only need to hide the not-yet-filled ring slots during the first window of steps. - # Keeping the mask shape constant (even when all slots are valid) lets the cuDNN SDPA backend - # plan its kernel once instead of re-planning every step. - layout = past_key_values.decode_kv_layout() if isinstance(past_key_values, UnlimitedOcrDynamicCache) else None - if inputs_embeds.shape[1] == 1 and layout is not None: - kv_total, kv_valid = layout - min_value = torch.finfo(inputs_embeds.dtype).min - mask = torch.zeros( - inputs_embeds.shape[0], 1, 1, kv_total, dtype=inputs_embeds.dtype, device=inputs_embeds.device + # Ring-buffer decode (single query token): the cache holds [prefill tokens] + [the generated + # tokens currently in the ring], i.e. all prefill plus the last `sliding_window` generated tokens. + # Attention is order-agnostic over keys (RoPE is baked into the cached keys), so the query attends + # to every cached key β€” a plain causal/sliding mask would wrongly hide the prefill. We therefore + # build an all-visible mask sized to the cache's physical length. + kv_length = None + if past_key_values is not None and past_key_values.layers: + first_layer = past_key_values.layers[0] + if isinstance(first_layer, DynamicReferenceSlidingWindowLayer): + kv_length = first_layer.decode_kv_length(inputs_embeds.shape[1]) + if inputs_embeds.shape[1] == 1 and kv_length is not None: + return torch.zeros( + inputs_embeds.shape[0], 1, 1, kv_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device ) - if kv_valid < kv_total: - mask[..., kv_valid:] = min_value - return mask # Prefill: standard full causal mask so image tokens can attend to each other correctly. return create_causal_mask(**mask_kwargs) @@ -832,7 +820,7 @@ def forward( inputs_embeds = self.embed_tokens(input_ids) if use_cache and past_key_values is None: - past_key_values = UnlimitedOcrDynamicCache(config=self.config) + 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 @@ -1108,24 +1096,21 @@ def prepare_inputs_for_generation( def _prepare_cache_for_generation( self, generation_config, model_kwargs, generation_mode, batch_size, max_cache_length ): - # `generate()` would otherwise build a plain `DynamicCache` whose KV grows unbounded, making decode - # O(nΒ²). For the default dynamic-cache case we instead use `UnlimitedOcrDynamicCache`, which keeps all - # prefill tokens but bounds generated tokens to `sliding_window` entries via a ring buffer (O(n) total). - # Caches that need rollback (assisted/contrastive) or an explicit `cache_implementation` defer to super. - uses_default_dynamic_cache = ( + # The default path builds a `DynamicCache(config)`, which dispatches `DynamicReferenceSlidingWindowLayer` + # per `config.layer_types` (ring buffer: all prefill kept, generated tokens bounded to `sliding_window`). + # The ring layer returns a fixed-size KV buffer and cannot support the cache rollback / multi-token + # verification that assisted and contrastive generation need, so for those modes we force a plain + # full-attention `DynamicCache` (rollback-capable, full causal attention) instead. + if ( model_kwargs.get("past_key_values") is None and generation_config.use_cache - and generation_config.cache_implementation is None - and generation_mode not in (GenerationMode.ASSISTED_GENERATION, GenerationMode.CONTRASTIVE_SEARCH) + and generation_mode in (GenerationMode.ASSISTED_GENERATION, GenerationMode.CONTRASTIVE_SEARCH) and self._supports_default_dynamic_cache() - ) - if uses_default_dynamic_cache: - model_kwargs["past_key_values"] = UnlimitedOcrDynamicCache( - config=self.config.get_text_config(decoder=True) - ) + ): + model_kwargs["past_key_values"] = DynamicCache() return - # Defer every other case (static/quantized/offloaded caches, encoder-decoder, no cache, ...) - # to the base implementation. Called on `GenerationMixin` directly rather than via `super()` + # Every other case (default dynamic ring cache, static/quantized/offloaded, encoder-decoder, no cache, + # ...) defers to the base implementation. Called on `GenerationMixin` directly rather than via `super()` # so the modular converter doesn't try to inline a method that isn't in the modeling lineage. return GenerationMixin._prepare_cache_for_generation( self, generation_config, model_kwargs, generation_mode, batch_size, max_cache_length diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index cad1a51eee3c..996401ec2bbe 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -125,6 +125,16 @@ class UnlimitedOcrModelTest(VLMModelTest, unittest.TestCase): test_all_params_have_gradient = False test_torch_exportable = False + def prepare_config_and_inputs_for_generate(self, batch_size=2): + config, inputs_dict = super().prepare_config_and_inputs_for_generate(batch_size=batch_size) + # `is_moe_model` (used by generation tests to pick a looser tolerance for MoE routing noise) checks + # `config._experts_implementation`, which is only populated when a model is instantiated. Some tests + # (e.g. `test_generate_with_static_cache`) check it on the bare config, so mark the model's default + # experts implementation here to get the MoE tolerance. The ring-buffer vs. static sliding cache + # differ by ~1e-3 once amplified through the sparse expert routing, which is well within that bound. + config._experts_implementation = "grouped_mm" + return config, inputs_dict + @unittest.skip( reason="UnlimitedOcrVisionModel builds a hybrid bidirectional+causal mask internally, so SDPA is always called with a non-null `attn_mask`." ) @@ -265,9 +275,12 @@ def test_small_model_integration_test_multi_page_document_parsing(self): "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png" ) ) - inputs = self.processor(images=[image1, image2], text="\nMulti page parsing.", crop_to_patches=False, return_tensors="pt").to( - model.device - ) + inputs = self.processor( + images=[image1, image2], + text="\nMulti page parsing.", + crop_to_patches=False, + return_tensors="pt", + ).to(model.device) with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) From a9eb8ae9e4c76b8b6195d17081d6344459c824e9 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 10:44:38 +0200 Subject: [PATCH 041/195] update --- .../unlimited_ocr/modeling_unlimited_ocr.py | 24 ---------------- .../unlimited_ocr/modular_unlimited_ocr.py | 28 +------------------ .../unlimited_ocr/processing_unlimited_ocr.py | 9 ++++++ 3 files changed, 10 insertions(+), 51 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index b2efc7a96c7f..684972676438 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -32,7 +32,6 @@ from ...cache_utils import Cache, DynamicCache, DynamicSlidingWindowLayer from ...configuration_utils import PretrainedConfig from ...generation import GenerationMixin -from ...generation.utils import GenerationMode from ...integrations import ( use_experts_implementation, use_kernel_forward_from_hub, @@ -2052,29 +2051,6 @@ def prepare_inputs_for_generation( return model_inputs - def _prepare_cache_for_generation( - self, generation_config, model_kwargs, generation_mode, batch_size, max_cache_length - ): - # The default path builds a `DynamicCache(config)`, which dispatches `DynamicReferenceSlidingWindowLayer` - # per `config.layer_types` (ring buffer: all prefill kept, generated tokens bounded to `sliding_window`). - # The ring layer returns a fixed-size KV buffer and cannot support the cache rollback / multi-token - # verification that assisted and contrastive generation need, so for those modes we force a plain - # full-attention `DynamicCache` (rollback-capable, full causal attention) instead. - if ( - model_kwargs.get("past_key_values") is None - and generation_config.use_cache - and generation_mode in (GenerationMode.ASSISTED_GENERATION, GenerationMode.CONTRASTIVE_SEARCH) - and self._supports_default_dynamic_cache() - ): - model_kwargs["past_key_values"] = DynamicCache() - return - # Every other case (default dynamic ring cache, static/quantized/offloaded, encoder-decoder, no cache, - # ...) defers to the base implementation. Called on `GenerationMixin` directly rather than via `super()` - # so the modular converter doesn't try to inline a method that isn't in the modeling lineage. - return super()._prepare_cache_for_generation( - generation_config, model_kwargs, generation_mode, batch_size, max_cache_length - ) - __all__ = [ "UnlimitedOcrForConditionalGeneration", diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 672a756e1a96..62f100427425 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -22,8 +22,6 @@ from ...cache_utils import Cache, DynamicCache, DynamicSlidingWindowLayer from ...configuration_utils import PretrainedConfig from ...feature_extraction_utils import BatchFeature -from ...generation import GenerationMixin -from ...generation.utils import GenerationMode from ...image_transforms import group_images_by_shape, reorder_images from ...image_utils import ImageInput, PILImageResampling, SizeDict from ...masking_utils import create_causal_mask @@ -1092,8 +1090,7 @@ def prepare_inputs_for_generation( **kwargs, ) - # Image inputs are only needed during prefill (or when the cache is disabled); once the image tokens - # have been embedded they must be dropped so later decode steps don't reprocess the pixel values. + # Image inputs are only needed during prefill or when the cache is disabled 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 @@ -1102,29 +1099,6 @@ def prepare_inputs_for_generation( return model_inputs - def _prepare_cache_for_generation( - self, generation_config, model_kwargs, generation_mode, batch_size, max_cache_length - ): - # The default path builds a `DynamicCache(config)`, which dispatches `DynamicReferenceSlidingWindowLayer` - # per `config.layer_types` (ring buffer: all prefill kept, generated tokens bounded to `sliding_window`). - # The ring layer returns a fixed-size KV buffer and cannot support the cache rollback / multi-token - # verification that assisted and contrastive generation need, so for those modes we force a plain - # full-attention `DynamicCache` (rollback-capable, full causal attention) instead. - if ( - model_kwargs.get("past_key_values") is None - and generation_config.use_cache - and generation_mode in (GenerationMode.ASSISTED_GENERATION, GenerationMode.CONTRASTIVE_SEARCH) - and self._supports_default_dynamic_cache() - ): - model_kwargs["past_key_values"] = DynamicCache() - return - # Every other case (default dynamic ring cache, static/quantized/offloaded, encoder-decoder, no cache, - # ...) defers to the base implementation. Called on `GenerationMixin` directly rather than via `super()` - # so the modular converter doesn't try to inline a method that isn't in the modeling lineage. - return GenerationMixin._prepare_cache_for_generation( - self, generation_config, model_kwargs, generation_mode, batch_size, max_cache_length - ) - __all__ = [ "UnlimitedOcrConfig", diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index 8c0b571f6040..5bdb54f8dcea 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -78,6 +78,15 @@ def _expand_image_tokens( Returns: `list[str]`: Text with expanded image token placeholders. """ + num_images = len(image_spatial_crop) + total_image_tokens = sum(t.count(self.image_token) for t in text) + if total_image_tokens != num_images: + raise ValueError( + f"Number of `{self.image_token}` tokens in the text ({total_image_tokens}) does not match the " + f"number of images passed ({num_images}). Use one `{self.image_token}` placeholder per image, " + f"e.g. `'' * len(images) + '\\nMulti page parsing.'`" + ) + size = self.image_processor.size["height"] tile_size = self.image_processor.tile_size From c20454372d1040a1cc8554781b473e79717882bb Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 16:05:46 +0200 Subject: [PATCH 042/195] refactor --- src/transformers/masking_utils.py | 134 ++++++++++++++++++ .../unlimited_ocr/modeling_unlimited_ocr.py | 79 ++++------- .../unlimited_ocr/modular_unlimited_ocr.py | 92 ++++++------ 3 files changed, 204 insertions(+), 101 deletions(-) diff --git a/src/transformers/masking_utils.py b/src/transformers/masking_utils.py index bf97836cb865..4f87217a6227 100644 --- a/src/transformers/masking_utils.py +++ b/src/transformers/masking_utils.py @@ -1243,6 +1243,140 @@ def create_sliding_window_causal_mask( return causal_mask +# NOTE: The only difference with create_causal_mask is that we select here sliding instead of full layers. +# TODO: Check if we really want to keep this. The only reason we need it is in case we have a hybrid cache +# with different layer types. In the default case where config.layer_types contains a single type of layer +# create_causal_mask is sufficient. +def create_reference_sliding_window_causal_mask( + config: PreTrainedConfig, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + past_key_values: Cache | None, + position_ids: torch.Tensor | None = None, + or_mask_function: Callable | None = None, + and_mask_function: Callable | None = None, + block_sequence_ids: torch.Tensor | None = None, +) -> torch.Tensor | BlockMask | None: + """ + Create the mask for the reference-sliding layers. This mirrors `create_sliding_window_causal_mask` (it reads the + sizes from the sliding cache layer), but uses a plain causal mask function instead of a sliding window one: the + reference-sliding cache already holds the window physically as a ring buffer (prefill + the last `sliding_window` + generated tokens) and reports its physical length with offset 0. A sliding window mask function would wrongly + evict the prefill; a plain causal mask instead stays correctly causal during prefill and degenerates to full + visibility over the physical buffer during single-token decode (every cached key index <= the large query + offset) β€” which is exactly what the order-agnostic ring buffer needs (RoPE is baked into the cached keys). + + Args: + config (`PreTrainedConfig`): + The model config. + inputs_embeds (`torch.Tensor`): + The input embeddings of shape (batch_size, query_length, hidden_dim). This is used only to infer the + batch size, query length and dtype. + attention_mask (`torch.Tensor`, optional): + The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length). + It can also be an already prepared 4D mask, in which case it is returned as-is. + cache_position (`torch.Tensor`): + Deprecated and unused. + past_key_values (`Cache`, optional): + The past key values, if we use a cache. + position_ids (`torch.Tensor`, optional) + A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences. + or_mask_function (`Callable`, optional): + An optional mask function to combine with the causal mask function (by doing the union of both). This is + useful to easily overlay another mask on top of the causal one, for example for image tokens handling. + and_mask_function (`Callable`, optional): + An optional mask function to combine with the causal mask function (by doing the intersection of both). This is + useful to easily overlay another mask on top of the causal one, for example for image tokens handling. + block_sequence_ids (`torch.Tensor`, *optional*): + A tensor of same shape as input IDs indicating to which block or group each token belongs to. Tokens from + the same block will keep a bidirectional mask within the block, attending causally to the past. Index `-1` + can be used for blocks that have to keep complete causality within itself. + """ + # Power feature: if `is_causal` is False, then fallback to bi-directional mask for bi-directional attention. + # It allows to use decoder-only models with bi-directional attention as well + if not getattr(config, "is_causal", True): + return create_bidirectional_mask( + config, + inputs_embeds, + attention_mask, + past_key_values=past_key_values, + or_mask_function=or_mask_function, + and_mask_function=and_mask_function, + ) + + # NOTE: The only difference with create_causal_mask is that we select here sliding instead of full layers + # If we have an hybrid cache structure, here we want to create the mask for the sliding layers + if hasattr(past_key_values, "is_sliding") and True in past_key_values.is_sliding: + layer_idx = past_key_values.is_sliding.index(True) + else: + layer_idx = 0 + + early_exit, attention_mask, packed_sequence_mask, q_length, kv_length, q_offset, kv_offset = ( + _preprocess_mask_arguments(config, inputs_embeds, attention_mask, past_key_values, position_ids, layer_idx) + ) + if early_exit: + return attention_mask + + batch_size, dtype, device = inputs_embeds.shape[0], inputs_embeds.dtype, inputs_embeds.device + mask_factory_function = causal_mask_function + mask_interface = ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation] + + # Defaulting to using non-vmap based mask creations except when detecting + # users passing custom mask functions (as we cannot guarantee that they + # are properly index-based as required by our implementation). + use_vmap = False + + # Do not allow skip if we are compiling (this is to match BC) + # TODO: cyril -> probably revisit and remove this, but a lot of tests rely on it + if _is_torch_xpu_available: + # Do not allow skip if we are compiling for decoding, but for prefill, we still allow skip to optimization the perf of 1st token generation + allow_is_causal_skip = not (getattr(past_key_values, "is_compileable", False) and q_length == 1) + else: + allow_is_causal_skip = not getattr(past_key_values, "is_compileable", False) + + # Allow slight deviations from causal mask + # Note that it is very important to apply this before any other deviations of the mask (such as packed sequence mask, + # padding mask, etc) as the resulting mask may otherwise not be correct! + if or_mask_function is not None: + if not _is_torch_greater_or_equal_than_2_6: + raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6") + mask_factory_function = or_masks(mask_factory_function, or_mask_function) + allow_is_causal_skip = False + use_vmap = True + if and_mask_function is not None: + if not _is_torch_greater_or_equal_than_2_6: + raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6") + mask_factory_function = and_masks(mask_factory_function, and_mask_function) + allow_is_causal_skip = False + use_vmap = True + + # If we detected packing format or blockwise overlay + if packed_sequence_mask is not None: + mask_factory_function = and_masks(mask_factory_function, packed_sequence_mask_function(packed_sequence_mask)) + allow_is_causal_skip = False + if block_sequence_ids is not None: + block_sequence_ids = maybe_pad_block_sequence_ids(block_sequence_ids, attention_mask, kv_length, kv_offset) + mask_factory_function = or_masks(mask_factory_function, blockwise_overlay(block_sequence_ids)) + allow_is_causal_skip = False + + # We now create the mask + causal_mask = mask_interface( + batch_size=batch_size, + q_length=q_length, + kv_length=kv_length, + q_offset=q_offset, + kv_offset=kv_offset, + mask_function=mask_factory_function, + attention_mask=attention_mask, + allow_is_causal_skip=allow_is_causal_skip, # additional kwarg for sdpa + dtype=dtype, # Additional kwarg for eager + config=config, # Pass the config as well, in case someone wants to easily have their own mask_interface + use_vmap=use_vmap, # Short-circuit to non-vmap expansions for the mask + device=device, + ) + return causal_mask + + def create_bidirectional_sliding_window_mask( config: PreTrainedConfig, inputs_embeds: torch.Tensor, diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 684972676438..fee142f20659 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -38,7 +38,7 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...masking_utils import create_causal_mask +from ...masking_utils import create_causal_mask, create_reference_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling @@ -1207,26 +1207,20 @@ def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwarg class DynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): - """Sliding-window cache layer that keeps all prefill tokens and windows only the generated ones. + """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens and windows + only the generated ones. The stock `DynamicSlidingWindowLayer` evicts the oldest tokens once the window fills, which would discard the image/prompt prefill. Here the prefill (every token cached before the first single-token decode step) is kept intact and never evicted. Generated tokens are appended right after the prefill until ``sliding_window`` of them have accumulated; from then on the newest generated token overwrites - the oldest one in place (a ring buffer). This mirrors the reference model's attention: each generated - token attends to all prefill tokens plus the last ``sliding_window`` generated tokens. + the oldest one. While the window is filling, the cached length grows by one each step (exactly like the stock sliding layer), so within the first ``sliding_window`` generated tokens the layer is indistinguishable from full attention. Once the ring is full the cached tensors stay at a constant length (``prefill + sliding_window``), which lets the cuDNN SDPA backend reuse a single kernel plan across the (typically long) steady-state decode instead of re-planning on every distinct sequence length. - - `get_seq_length()` (inherited) returns the *logical* cumulative token count so positions are computed - correctly; the physically stored length is bounded at ``prefill + sliding_window``. Because keys/values - are overwritten in place the layer cannot roll back, so cache-rollback generation modes (assisted, - contrastive) use a plain full `DynamicCache` instead (see - `UnlimitedOcrForConditionalGeneration._prepare_cache_for_generation`). """ layer_type = "reference_sliding_attention" @@ -1239,25 +1233,26 @@ def __init__(self, config: PretrainedConfig | None = None, sliding_window: int | def update( self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs ) -> tuple[torch.Tensor, torch.Tensor]: + # Lazy initialization if not self.is_initialized: self.lazy_initialization(key_states, value_states) - seq_len = key_states.shape[-2] - self.cumulative_length += seq_len + sequence_length = key_states.shape[-2] + self.cumulative_length += sequence_length - # Prefill: the first update carries the image/prompt context (more than one token). Keep it intact. - if self.prefill_length is None and seq_len > 1: + # Prefill with prompt context + if self.prefill_length is None and sequence_length > 1: self.keys = torch.cat([self.keys, key_states], dim=-2) self.values = torch.cat([self.values, value_states], dim=-2) return self.keys, self.values - # First decode step: freeze the prefill boundary. `self.keys` still holds the (possibly empty) - # prefill here; an empty buffer means generation started without any prompt (BOS only). + # First decode step + # Handle generation with empty prompt if self.prefill_length is None: self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 # TODO: can we remove the for loop? - for token_idx in range(seq_len): + for token_idx in range(sequence_length): key = key_states[..., token_idx : token_idx + 1, :] value = value_states[..., token_idx : token_idx + 1, :] generated_length = self.keys.shape[-2] - self.prefill_length if self.keys.dim() > 1 else 0 @@ -1671,20 +1666,32 @@ def forward( if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) + # It may already have been prepared by, e.g., `generate` 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 = self._create_attention_mask(inputs_embeds, attention_mask, past_key_values, position_ids) + if not isinstance(causal_mask_mapping := attention_mask, dict): + mask_kwargs = { + "config": self.config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + causal_mask_mapping = { + "full_attention": create_causal_mask(**mask_kwargs), + "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), + } 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]: + for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): hidden_states = decoder_layer( hidden_states, - attention_mask=causal_mask, + attention_mask=causal_mask_mapping[self.config.layer_types[layer_idx]], position_embeddings=position_embeddings, position_ids=position_ids, past_key_values=past_key_values, @@ -1698,35 +1705,6 @@ def forward( past_key_values=past_key_values, ) - def _create_attention_mask(self, inputs_embeds, attention_mask, past_key_values, position_ids): - mask_kwargs = { - "config": self.config, - "inputs_embeds": inputs_embeds, - "attention_mask": attention_mask, - "past_key_values": past_key_values, - "position_ids": position_ids, - } - if self.config.sliding_window is None: - return create_causal_mask(**mask_kwargs) - - # Ring-buffer decode (single query token): the cache holds [prefill tokens] + [the generated - # tokens currently in the ring], i.e. all prefill plus the last `sliding_window` generated tokens. - # Attention is order-agnostic over keys (RoPE is baked into the cached keys), so the query attends - # to every cached key β€” a plain causal/sliding mask would wrongly hide the prefill. We therefore - # build an all-visible mask sized to the cache's physical length. - kv_length = None - if past_key_values is not None and past_key_values.layers: - first_layer = past_key_values.layers[0] - if isinstance(first_layer, DynamicReferenceSlidingWindowLayer): - kv_length = first_layer.decode_kv_length(inputs_embeds.shape[1]) - if inputs_embeds.shape[1] == 1 and kv_length is not None: - return torch.zeros( - inputs_embeds.shape[0], 1, 1, kv_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device - ) - - # Prefill: standard full causal mask so image tokens can attend to each other correctly. - return create_causal_mask(**mask_kwargs) - @auto_docstring( custom_intro=""" @@ -2041,8 +2019,7 @@ def prepare_inputs_for_generation( **kwargs, ) - # Image inputs are only needed during prefill (or when the cache is disabled); once the image tokens - # have been embedded they must be dropped so later decode steps don't reprocess the pixel values. + # Image inputs are only needed during prefill or when the cache is disabled 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 diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 62f100427425..bb044a6913a5 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -24,11 +24,21 @@ from ...feature_extraction_utils import BatchFeature from ...image_transforms import group_images_by_shape, reorder_images from ...image_utils import ImageInput, PILImageResampling, SizeDict -from ...masking_utils import create_causal_mask +from ...masking_utils import ( + create_causal_mask, + create_reference_sliding_window_causal_mask, +) from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast from ...processing_utils import Unpack from ...tokenization_utils_base import PreTokenizedInput, TextInput -from ...utils import TensorType, TransformersKwargs, auto_docstring, can_return_tuple, torch_int +from ...utils import ( + TensorType, + TransformersKwargs, + auto_docstring, + can_return_tuple, + is_torch_flex_attn_available, + torch_int, +) from ...utils.generic import merge_with_config_defaults from ...utils.output_capturing import capture_outputs from ..clip.configuration_clip import CLIPVisionConfig @@ -75,6 +85,10 @@ from ..got_ocr2.configuration_got_ocr2 import GotOcr2VisionConfig +if is_torch_flex_attn_available(): + pass + + class UnlimitedOcrImageProcessor(DeepseekOcr2ImageProcessor): tile_size = 640 max_patches = 32 @@ -652,26 +666,20 @@ def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwarg class DynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): - """Sliding-window cache layer that keeps all prefill tokens and windows only the generated ones. + """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens and windows + only the generated ones. The stock `DynamicSlidingWindowLayer` evicts the oldest tokens once the window fills, which would discard the image/prompt prefill. Here the prefill (every token cached before the first single-token decode step) is kept intact and never evicted. Generated tokens are appended right after the prefill until ``sliding_window`` of them have accumulated; from then on the newest generated token overwrites - the oldest one in place (a ring buffer). This mirrors the reference model's attention: each generated - token attends to all prefill tokens plus the last ``sliding_window`` generated tokens. + the oldest one. While the window is filling, the cached length grows by one each step (exactly like the stock sliding layer), so within the first ``sliding_window`` generated tokens the layer is indistinguishable from full attention. Once the ring is full the cached tensors stay at a constant length (``prefill + sliding_window``), which lets the cuDNN SDPA backend reuse a single kernel plan across the (typically long) steady-state decode instead of re-planning on every distinct sequence length. - - `get_seq_length()` (inherited) returns the *logical* cumulative token count so positions are computed - correctly; the physically stored length is bounded at ``prefill + sliding_window``. Because keys/values - are overwritten in place the layer cannot roll back, so cache-rollback generation modes (assisted, - contrastive) use a plain full `DynamicCache` instead (see - `UnlimitedOcrForConditionalGeneration._prepare_cache_for_generation`). """ layer_type = "reference_sliding_attention" @@ -684,25 +692,26 @@ def __init__(self, config: PretrainedConfig | None = None, sliding_window: int | def update( self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs ) -> tuple[torch.Tensor, torch.Tensor]: + # Lazy initialization if not self.is_initialized: self.lazy_initialization(key_states, value_states) - seq_len = key_states.shape[-2] - self.cumulative_length += seq_len + sequence_length = key_states.shape[-2] + self.cumulative_length += sequence_length - # Prefill: the first update carries the image/prompt context (more than one token). Keep it intact. - if self.prefill_length is None and seq_len > 1: + # Prefill with prompt context + if self.prefill_length is None and sequence_length > 1: self.keys = torch.cat([self.keys, key_states], dim=-2) self.values = torch.cat([self.values, value_states], dim=-2) return self.keys, self.values - # First decode step: freeze the prefill boundary. `self.keys` still holds the (possibly empty) - # prefill here; an empty buffer means generation started without any prompt (BOS only). + # First decode step + # Handle generation with empty prompt if self.prefill_length is None: self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 # TODO: can we remove the for loop? - for token_idx in range(seq_len): + for token_idx in range(sequence_length): key = key_states[..., token_idx : token_idx + 1, :] value = value_states[..., token_idx : token_idx + 1, :] generated_length = self.keys.shape[-2] - self.prefill_length if self.keys.dim() > 1 else 0 @@ -778,35 +787,6 @@ class UnlimitedOcrTextPreTrainedModel(DeepseekOcr2TextPreTrainedModel): class UnlimitedOcrTextModel(DeepseekOcr2TextModel): - def _create_attention_mask(self, inputs_embeds, attention_mask, past_key_values, position_ids): - mask_kwargs = { - "config": self.config, - "inputs_embeds": inputs_embeds, - "attention_mask": attention_mask, - "past_key_values": past_key_values, - "position_ids": position_ids, - } - if self.config.sliding_window is None: - return create_causal_mask(**mask_kwargs) - - # Ring-buffer decode (single query token): the cache holds [prefill tokens] + [the generated - # tokens currently in the ring], i.e. all prefill plus the last `sliding_window` generated tokens. - # Attention is order-agnostic over keys (RoPE is baked into the cached keys), so the query attends - # to every cached key β€” a plain causal/sliding mask would wrongly hide the prefill. We therefore - # build an all-visible mask sized to the cache's physical length. - kv_length = None - if past_key_values is not None and past_key_values.layers: - first_layer = past_key_values.layers[0] - if isinstance(first_layer, DynamicReferenceSlidingWindowLayer): - kv_length = first_layer.decode_kv_length(inputs_embeds.shape[1]) - if inputs_embeds.shape[1] == 1 and kv_length is not None: - return torch.zeros( - inputs_embeds.shape[0], 1, 1, kv_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device - ) - - # Prefill: standard full causal mask so image tokens can attend to each other correctly. - return create_causal_mask(**mask_kwargs) - @merge_with_config_defaults @capture_outputs @auto_docstring @@ -829,20 +809,32 @@ def forward( if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) + # It may already have been prepared by, e.g., `generate` 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 = self._create_attention_mask(inputs_embeds, attention_mask, past_key_values, position_ids) + if not isinstance(causal_mask_mapping := attention_mask, dict): + mask_kwargs = { + "config": self.config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + causal_mask_mapping = { + "full_attention": create_causal_mask(**mask_kwargs), + "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), + } 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]: + for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): hidden_states = decoder_layer( hidden_states, - attention_mask=causal_mask, + attention_mask=causal_mask_mapping[self.config.layer_types[layer_idx]], position_embeddings=position_embeddings, position_ids=position_ids, past_key_values=past_key_values, From a67620e594a4bd79064c69e121fa399e7476e38d Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 16:20:08 +0200 Subject: [PATCH 043/195] optimize --- .../unlimited_ocr/modeling_unlimited_ocr.py | 33 ++++++++++--------- .../unlimited_ocr/modular_unlimited_ocr.py | 33 ++++++++++--------- 2 files changed, 36 insertions(+), 30 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index fee142f20659..5a06b40da675 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1251,21 +1251,24 @@ def update( if self.prefill_length is None: self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 - # TODO: can we remove the for loop? - for token_idx in range(sequence_length): - key = key_states[..., token_idx : token_idx + 1, :] - value = value_states[..., token_idx : token_idx + 1, :] - generated_length = self.keys.shape[-2] - self.prefill_length if self.keys.dim() > 1 else 0 - if generated_length < self.sliding_window: - # Window not full yet: append (the buffer grows by one, like the stock sliding layer). - self.keys = torch.cat([self.keys, key], dim=-2) - self.values = torch.cat([self.values, value], dim=-2) - else: - # Window full: overwrite the oldest generated slot in place (ring buffer, constant length). - slot = self.prefill_length + self.ring_position - self.keys[..., slot : slot + 1, :] = key - self.values[..., slot : slot + 1, :] = value - self.ring_position = (self.ring_position + 1) % self.sliding_window + # Append while window still grows + generated_length = self.keys.shape[-2] - self.prefill_length if self.keys.dim() > 1 else 0 + append_length = min(sequence_length, max(0, self.sliding_window - generated_length)) + if append_length > 0: + self.keys = torch.cat([self.keys, key_states[..., :append_length, :]], dim=-2) + self.values = torch.cat([self.values, value_states[..., :append_length, :]], dim=-2) + + # Overwrite once window size is reached + overwrite_length = sequence_length - append_length + if overwrite_length > 0: + # Only the most recent `sliding_window` overwrites survive + write_length = min(overwrite_length, self.sliding_window) + start = self.ring_position + overwrite_length - write_length + offsets = torch.arange(write_length, device=key_states.device) + slots = self.prefill_length + (start + offsets) % self.sliding_window + self.keys[..., slots, :] = key_states[..., sequence_length - write_length :, :] + self.values[..., slots, :] = value_states[..., sequence_length - write_length :, :] + self.ring_position = (self.ring_position + overwrite_length) % self.sliding_window return self.keys, self.values diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index bb044a6913a5..32a0187d385b 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -710,21 +710,24 @@ def update( if self.prefill_length is None: self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 - # TODO: can we remove the for loop? - for token_idx in range(sequence_length): - key = key_states[..., token_idx : token_idx + 1, :] - value = value_states[..., token_idx : token_idx + 1, :] - generated_length = self.keys.shape[-2] - self.prefill_length if self.keys.dim() > 1 else 0 - if generated_length < self.sliding_window: - # Window not full yet: append (the buffer grows by one, like the stock sliding layer). - self.keys = torch.cat([self.keys, key], dim=-2) - self.values = torch.cat([self.values, value], dim=-2) - else: - # Window full: overwrite the oldest generated slot in place (ring buffer, constant length). - slot = self.prefill_length + self.ring_position - self.keys[..., slot : slot + 1, :] = key - self.values[..., slot : slot + 1, :] = value - self.ring_position = (self.ring_position + 1) % self.sliding_window + # Append while window grows + generated_length = self.keys.shape[-2] - self.prefill_length if self.keys.dim() > 1 else 0 + append_length = min(sequence_length, max(0, self.sliding_window - generated_length)) + if append_length > 0: + self.keys = torch.cat([self.keys, key_states[..., :append_length, :]], dim=-2) + self.values = torch.cat([self.values, value_states[..., :append_length, :]], dim=-2) + + # Overwrite if window size is reached + overwrite_length = sequence_length - append_length + if overwrite_length > 0: + # Only the most recent `sliding_window` overwrites survive + write_length = min(overwrite_length, self.sliding_window) + start = self.ring_position + overwrite_length - write_length + offsets = torch.arange(write_length, device=key_states.device) + slots = self.prefill_length + (start + offsets) % self.sliding_window + self.keys[..., slots, :] = key_states[..., sequence_length - write_length :, :] + self.values[..., slots, :] = value_states[..., sequence_length - write_length :, :] + self.ring_position = (self.ring_position + overwrite_length) % self.sliding_window return self.keys, self.values From 34c8d1dd0ef4d873323eee717ced33e5a3aee94c Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 16:24:12 +0200 Subject: [PATCH 044/195] fix-repo --- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 5a06b40da675..c400b68ff8b9 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1251,14 +1251,14 @@ def update( if self.prefill_length is None: self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 - # Append while window still grows + # Append while window grows generated_length = self.keys.shape[-2] - self.prefill_length if self.keys.dim() > 1 else 0 append_length = min(sequence_length, max(0, self.sliding_window - generated_length)) if append_length > 0: self.keys = torch.cat([self.keys, key_states[..., :append_length, :]], dim=-2) self.values = torch.cat([self.values, value_states[..., :append_length, :]], dim=-2) - # Overwrite once window size is reached + # Overwrite if window size is reached overwrite_length = sequence_length - append_length if overwrite_length > 0: # Only the most recent `sliding_window` overwrites survive From 3dffb4dbec8b4b2115f1435b3cc576e96576837c Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 16:24:52 +0200 Subject: [PATCH 045/195] update docs --- docs/source/en/model_doc/unlimited_ocr.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index ecbc06dc44b9..34d0eae429a8 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -51,7 +51,7 @@ model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", devic processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") image = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" -inputs = processor(images=image, text="\ndocument parsing.", return_tensors="pt").to(model.device) +inputs = processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) output = model.generate(**inputs, max_new_tokens=4096) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) @@ -75,7 +75,7 @@ num_pages = 2 inputs = processor( images=[page1, page2], - text="" * num_pages + "\nMulti page parsing.", + text="" * num_pages + "Multi page parsing.", return_tensors="pt", ).to(model.device) From 1d9a231060fe151c97e38ba8deb01efb81ff57be Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 16:56:00 +0200 Subject: [PATCH 046/195] update tests --- tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index c9e5944938d8..95605e6972fd 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -189,7 +189,7 @@ def test_small_model_integration_test_document_parsing(self): "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" ) ) - inputs = self.processor(images=image, text="\ndocument parsing.", return_tensors="pt").to(model.device) + inputs = self.processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True) @@ -211,7 +211,7 @@ def test_small_model_integration_test_document_parsing_grounding(self): ) inputs = self.processor( images=image, - text="\ndocument parsing.", + text="document parsing.", return_tensors="pt", ).to(model.device) with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): @@ -240,7 +240,7 @@ def test_small_model_integration_test_document_parsing_batched(self): ) inputs = self.processor( images=[image1, image2], - text=["\ndocument parsing.", "\ndocument parsing."], + text=["document parsing.", "document parsing."], return_tensors="pt", padding=True, ).to(model.device, dtype=torch.bfloat16) @@ -277,7 +277,7 @@ def test_small_model_integration_test_multi_page_document_parsing(self): ) inputs = self.processor( images=[image1, image2], - text="\nMulti page parsing.", + text="Multi page parsing.", crop_to_patches=False, return_tensors="pt", ).to(model.device) From 62883fc7fbce4aa403dbcb225e6ce6ecdf178478 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 17:09:39 +0200 Subject: [PATCH 047/195] rename to local_patches_grid --- .../image_processing_unlimited_ocr.py | 9 ++-- .../unlimited_ocr/modeling_unlimited_ocr.py | 22 ++++----- .../unlimited_ocr/modular_unlimited_ocr.py | 46 ++++++++----------- .../unlimited_ocr/processing_unlimited_ocr.py | 10 ++-- 4 files changed, 40 insertions(+), 47 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index 1a2c90d6f932..6c19007ac9c1 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -146,7 +146,7 @@ class UnlimitedOcrImageProcessor(TorchvisionBackend): max_patches = 32 tile_size = 640 background_color = [127, 127, 127] - model_input_names = ["pixel_values", "num_local_patches", "image_spatial_crop"] + model_input_names = ["pixel_values", "num_local_patches", "local_patches_grid"] def __init__(self, **kwargs: Unpack[UnlimitedOcrImageProcessorKwargs]): super().__init__(**kwargs) @@ -279,7 +279,7 @@ def _preprocess( data["pixel_values_local"] = flat_local_list # Compute per-image spatial crop grid and local-patch counts. - image_spatial_crop = [] + local_patches_grid = [] num_local_patches = [] for image in images: height, width = image.shape[-2:] @@ -291,11 +291,10 @@ def _preprocess( else: num_columns, num_rows = 1, 1 num_local_patches.append(0) - image_spatial_crop.append([num_columns, num_rows]) + local_patches_grid.append([num_columns, num_rows]) - # TODO: rename - data["image_spatial_crop"] = image_spatial_crop data["num_local_patches"] = num_local_patches + data["local_patches_grid"] = local_patches_grid return BatchFeature( data=data, diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index c400b68ff8b9..f82405ed3ea7 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1743,7 +1743,7 @@ def get_image_features( pixel_values: torch.FloatTensor, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - image_spatial_crop: torch.Tensor | None = None, + local_patches_grid: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> "UnlimitedOcrModelOutputWithPooling": r""" @@ -1751,7 +1751,7 @@ def get_image_features( 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]`. - image_spatial_crop (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + local_patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): The local crop grid `(num_columns, num_rows)` per image. """ if isinstance(num_local_patches, torch.Tensor): @@ -1788,7 +1788,7 @@ def get_image_features( local_features = per_image_local[idx] if local_features is not None and local_features.shape[0] > 0: - num_columns, num_rows = int(image_spatial_crop[idx][0]), int(image_spatial_crop[idx][1]) + num_columns, num_rows = int(local_patches_grid[idx][0]), int(local_patches_grid[idx][1]) num_queries_local = int(local_features.shape[1] ** 0.5) local_grid = local_features.reshape( num_rows, num_columns, num_queries_local, num_queries_local, hidden_size @@ -1847,7 +1847,7 @@ def forward( pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - image_spatial_crop: torch.Tensor | None = None, + local_patches_grid: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -1860,7 +1860,7 @@ def forward( 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. - image_spatial_crop (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + local_patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): The local crop grid `(num_columns, num_rows)` per image. """ if inputs_embeds is None: @@ -1869,7 +1869,7 @@ def forward( image_features = None if pixel_values is not None: image_features = self.get_image_features( - pixel_values, pixel_values_local, num_local_patches, image_spatial_crop, return_dict=True + pixel_values, pixel_values_local, num_local_patches, local_patches_grid, return_dict=True ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) @@ -1943,7 +1943,7 @@ def forward( pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - image_spatial_crop: torch.Tensor | None = None, + local_patches_grid: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -1958,7 +1958,7 @@ def forward( 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. - image_spatial_crop (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + local_patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): The local crop grid `(num_columns, num_rows)` per image. """ outputs = self.model( @@ -1966,7 +1966,7 @@ def forward( pixel_values=pixel_values, pixel_values_local=pixel_values_local, num_local_patches=num_local_patches, - image_spatial_crop=image_spatial_crop, + local_patches_grid=local_patches_grid, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, @@ -2006,7 +2006,7 @@ def prepare_inputs_for_generation( pixel_values=None, pixel_values_local=None, num_local_patches=None, - image_spatial_crop=None, + local_patches_grid=None, attention_mask=None, logits_to_keep=None, is_first_iteration=False, @@ -2027,7 +2027,7 @@ def prepare_inputs_for_generation( model_inputs["pixel_values"] = pixel_values model_inputs["pixel_values_local"] = pixel_values_local model_inputs["num_local_patches"] = num_local_patches - model_inputs["image_spatial_crop"] = image_spatial_crop + model_inputs["local_patches_grid"] = local_patches_grid return model_inputs diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 32a0187d385b..b5f2c9c49537 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -36,7 +36,6 @@ TransformersKwargs, auto_docstring, can_return_tuple, - is_torch_flex_attn_available, torch_int, ) from ...utils.generic import merge_with_config_defaults @@ -85,14 +84,10 @@ from ..got_ocr2.configuration_got_ocr2 import GotOcr2VisionConfig -if is_torch_flex_attn_available(): - pass - - class UnlimitedOcrImageProcessor(DeepseekOcr2ImageProcessor): tile_size = 640 max_patches = 32 - model_input_names = ["pixel_values", "num_local_patches", "image_spatial_crop"] + model_input_names = ["pixel_values", "num_local_patches", "local_patches_grid"] def pad_to_square( self, @@ -210,7 +205,7 @@ def _preprocess( data["pixel_values_local"] = flat_local_list # Compute per-image spatial crop grid and local-patch counts. - image_spatial_crop = [] + local_patches_grid = [] num_local_patches = [] for image in images: height, width = image.shape[-2:] @@ -222,11 +217,10 @@ def _preprocess( else: num_columns, num_rows = 1, 1 num_local_patches.append(0) - image_spatial_crop.append([num_columns, num_rows]) + local_patches_grid.append([num_columns, num_rows]) - # TODO: rename - data["image_spatial_crop"] = image_spatial_crop data["num_local_patches"] = num_local_patches + data["local_patches_grid"] = local_patches_grid return BatchFeature( data=data, @@ -246,10 +240,10 @@ class UnlimitedOcrProcessor(DeepseekOcr2Processor): def _expand_image_tokens( self, text: list[TextInput], - image_spatial_crop: torch.Tensor, + local_patches_grid: torch.Tensor, num_local_patches: list[int] | torch.Tensor, ) -> list[str]: - num_images = len(image_spatial_crop) + num_images = len(local_patches_grid) total_image_tokens = sum(t.count(self.image_token) for t in text) if total_image_tokens != num_images: raise ValueError( @@ -267,8 +261,8 @@ def _expand_image_tokens( crop_index = 0 for i in range(len(text)): while self.image_token in text[i]: - num_columns = int(image_spatial_crop[crop_index][0]) - num_rows = int(image_spatial_crop[crop_index][1]) + num_columns = int(local_patches_grid[crop_index][0]) + num_rows = int(local_patches_grid[crop_index][1]) num_tokens = num_queries_global * (num_queries_global + 1) + 1 if int(num_local_patches[crop_index]) > 0: num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) @@ -303,7 +297,7 @@ def __call__( text = text.copy() # below lines change text in-place image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) - text = self._expand_image_tokens(text, image_inputs["image_spatial_crop"], image_inputs["num_local_patches"]) + text = self._expand_image_tokens(text, image_inputs["local_patches_grid"], image_inputs["num_local_patches"]) return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) @@ -868,7 +862,7 @@ def get_image_features( pixel_values: torch.FloatTensor, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - image_spatial_crop: torch.Tensor | None = None, + local_patches_grid: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> "UnlimitedOcrModelOutputWithPooling": r""" @@ -876,7 +870,7 @@ def get_image_features( 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]`. - image_spatial_crop (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + local_patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): The local crop grid `(num_columns, num_rows)` per image. """ if isinstance(num_local_patches, torch.Tensor): @@ -913,7 +907,7 @@ def get_image_features( local_features = per_image_local[idx] if local_features is not None and local_features.shape[0] > 0: - num_columns, num_rows = int(image_spatial_crop[idx][0]), int(image_spatial_crop[idx][1]) + num_columns, num_rows = int(local_patches_grid[idx][0]), int(local_patches_grid[idx][1]) num_queries_local = int(local_features.shape[1] ** 0.5) local_grid = local_features.reshape( num_rows, num_columns, num_queries_local, num_queries_local, hidden_size @@ -949,7 +943,7 @@ def forward( pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - image_spatial_crop: torch.Tensor | None = None, + local_patches_grid: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -962,7 +956,7 @@ def forward( 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. - image_spatial_crop (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + local_patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): The local crop grid `(num_columns, num_rows)` per image. """ if inputs_embeds is None: @@ -971,7 +965,7 @@ def forward( image_features = None if pixel_values is not None: image_features = self.get_image_features( - pixel_values, pixel_values_local, num_local_patches, image_spatial_crop, return_dict=True + pixel_values, pixel_values_local, num_local_patches, local_patches_grid, return_dict=True ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) @@ -1006,7 +1000,7 @@ def forward( pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - image_spatial_crop: torch.Tensor | None = None, + local_patches_grid: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -1021,7 +1015,7 @@ def forward( 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. - image_spatial_crop (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + local_patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): The local crop grid `(num_columns, num_rows)` per image. """ outputs = self.model( @@ -1029,7 +1023,7 @@ def forward( pixel_values=pixel_values, pixel_values_local=pixel_values_local, num_local_patches=num_local_patches, - image_spatial_crop=image_spatial_crop, + local_patches_grid=local_patches_grid, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, @@ -1069,7 +1063,7 @@ def prepare_inputs_for_generation( pixel_values=None, pixel_values_local=None, num_local_patches=None, - image_spatial_crop=None, + local_patches_grid=None, attention_mask=None, logits_to_keep=None, is_first_iteration=False, @@ -1090,7 +1084,7 @@ def prepare_inputs_for_generation( model_inputs["pixel_values"] = pixel_values model_inputs["pixel_values_local"] = pixel_values_local model_inputs["num_local_patches"] = num_local_patches - model_inputs["image_spatial_crop"] = image_spatial_crop + model_inputs["local_patches_grid"] = local_patches_grid return model_inputs diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index 5bdb54f8dcea..8c51f6ede89d 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -62,7 +62,7 @@ def __init__( def _expand_image_tokens( self, text: list[TextInput], - image_spatial_crop: torch.Tensor, + local_patches_grid: torch.Tensor, num_local_patches: list[int] | torch.Tensor, ) -> list[str]: """ @@ -78,7 +78,7 @@ def _expand_image_tokens( Returns: `list[str]`: Text with expanded image token placeholders. """ - num_images = len(image_spatial_crop) + num_images = len(local_patches_grid) total_image_tokens = sum(t.count(self.image_token) for t in text) if total_image_tokens != num_images: raise ValueError( @@ -96,8 +96,8 @@ def _expand_image_tokens( crop_index = 0 for i in range(len(text)): while self.image_token in text[i]: - num_columns = int(image_spatial_crop[crop_index][0]) - num_rows = int(image_spatial_crop[crop_index][1]) + num_columns = int(local_patches_grid[crop_index][0]) + num_rows = int(local_patches_grid[crop_index][1]) num_tokens = num_queries_global * (num_queries_global + 1) + 1 if int(num_local_patches[crop_index]) > 0: num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) @@ -143,7 +143,7 @@ def __call__( text = text.copy() # below lines change text in-place image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) - text = self._expand_image_tokens(text, image_inputs["image_spatial_crop"], image_inputs["num_local_patches"]) + text = self._expand_image_tokens(text, image_inputs["local_patches_grid"], image_inputs["num_local_patches"]) return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) From 37da1841397cf7980f12c528cfcc91e44965000d Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 17:12:02 +0200 Subject: [PATCH 048/195] add comment --- .../models/unlimited_ocr/image_processing_unlimited_ocr.py | 1 + src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 1 + 2 files changed, 2 insertions(+) diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index 6c19007ac9c1..47f22b900351 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -353,6 +353,7 @@ def pad_to_square( ) max_dim = max(height, width) + # Use round(x/2) instead of x//2 to match PIL.ImageOps.pad behavior paste_x_left = round((max_dim - width) / 2) paste_y_left = round((max_dim - height) / 2) paste_x_right = max_dim - width - paste_x_left diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index b5f2c9c49537..3b6fb44a41c6 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -121,6 +121,7 @@ def pad_to_square( ) max_dim = max(height, width) + # Use round(x/2) instead of x//2 to match PIL.ImageOps.pad behavior paste_x_left = round((max_dim - width) / 2) paste_y_left = round((max_dim - height) / 2) paste_x_right = max_dim - width - paste_x_left From 95c7f1eb4c8a0bf3c81734333b07febddd9e399e Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 17:14:10 +0200 Subject: [PATCH 049/195] update comment --- src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 3b6fb44a41c6..1f27469eb92b 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -250,7 +250,7 @@ def _expand_image_tokens( raise ValueError( f"Number of `{self.image_token}` tokens in the text ({total_image_tokens}) does not match the " f"number of images passed ({num_images}). Use one `{self.image_token}` placeholder per image, " - f"e.g. `'' * len(images) + '\\nMulti page parsing.'`" + f"e.g. `'' * len(images) + 'Multi page parsing.'`" ) size = self.image_processor.size["height"] From f37b18e55083f0273d472c639b3f7602e5abf4be Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 17:26:52 +0200 Subject: [PATCH 050/195] rename --- .../image_processing_unlimited_ocr.py | 8 +-- .../unlimited_ocr/modeling_unlimited_ocr.py | 28 ++++----- .../unlimited_ocr/modular_unlimited_ocr.py | 59 +++++++++++-------- .../unlimited_ocr/processing_unlimited_ocr.py | 14 +++-- 4 files changed, 62 insertions(+), 47 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index 47f22b900351..617ed1edd3aa 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -146,7 +146,7 @@ class UnlimitedOcrImageProcessor(TorchvisionBackend): max_patches = 32 tile_size = 640 background_color = [127, 127, 127] - model_input_names = ["pixel_values", "num_local_patches", "local_patches_grid"] + model_input_names = ["pixel_values", "num_local_patches", "patches_grid"] def __init__(self, **kwargs: Unpack[UnlimitedOcrImageProcessorKwargs]): super().__init__(**kwargs) @@ -279,7 +279,7 @@ def _preprocess( data["pixel_values_local"] = flat_local_list # Compute per-image spatial crop grid and local-patch counts. - local_patches_grid = [] + patches_grid = [] num_local_patches = [] for image in images: height, width = image.shape[-2:] @@ -291,10 +291,10 @@ def _preprocess( else: num_columns, num_rows = 1, 1 num_local_patches.append(0) - local_patches_grid.append([num_columns, num_rows]) + patches_grid.append([num_columns, num_rows]) data["num_local_patches"] = num_local_patches - data["local_patches_grid"] = local_patches_grid + data["patches_grid"] = patches_grid return BatchFeature( data=data, diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index f82405ed3ea7..1661cc7f8eac 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1743,7 +1743,7 @@ def get_image_features( pixel_values: torch.FloatTensor, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - local_patches_grid: torch.Tensor | None = None, + patches_grid: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> "UnlimitedOcrModelOutputWithPooling": r""" @@ -1751,8 +1751,8 @@ def get_image_features( 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]`. - local_patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): - The local crop grid `(num_columns, num_rows)` per image. + patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + The patches grid `(num_columns, num_rows)` per image. """ if isinstance(num_local_patches, torch.Tensor): num_local_patches = num_local_patches.tolist() @@ -1788,7 +1788,7 @@ def get_image_features( local_features = per_image_local[idx] if local_features is not None and local_features.shape[0] > 0: - num_columns, num_rows = int(local_patches_grid[idx][0]), int(local_patches_grid[idx][1]) + num_columns, num_rows = int(patches_grid[idx][0]), int(patches_grid[idx][1]) num_queries_local = int(local_features.shape[1] ** 0.5) local_grid = local_features.reshape( num_rows, num_columns, num_queries_local, num_queries_local, hidden_size @@ -1847,7 +1847,7 @@ def forward( pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - local_patches_grid: torch.Tensor | None = None, + patches_grid: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -1860,8 +1860,8 @@ def forward( 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. - local_patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): - The local crop grid `(num_columns, num_rows)` per image. + patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + The patches grid `(num_columns, num_rows)` per image. """ if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) @@ -1869,7 +1869,7 @@ def forward( image_features = None if pixel_values is not None: image_features = self.get_image_features( - pixel_values, pixel_values_local, num_local_patches, local_patches_grid, return_dict=True + pixel_values, pixel_values_local, num_local_patches, patches_grid, return_dict=True ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) @@ -1943,7 +1943,7 @@ def forward( pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - local_patches_grid: torch.Tensor | None = None, + patches_grid: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -1958,15 +1958,15 @@ def forward( 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. - local_patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): - The local crop grid `(num_columns, num_rows)` per image. + patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + The patches grid `(num_columns, num_rows)` per image. """ outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, pixel_values_local=pixel_values_local, num_local_patches=num_local_patches, - local_patches_grid=local_patches_grid, + patches_grid=patches_grid, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, @@ -2006,7 +2006,7 @@ def prepare_inputs_for_generation( pixel_values=None, pixel_values_local=None, num_local_patches=None, - local_patches_grid=None, + patches_grid=None, attention_mask=None, logits_to_keep=None, is_first_iteration=False, @@ -2027,7 +2027,7 @@ def prepare_inputs_for_generation( model_inputs["pixel_values"] = pixel_values model_inputs["pixel_values_local"] = pixel_values_local model_inputs["num_local_patches"] = num_local_patches - model_inputs["local_patches_grid"] = local_patches_grid + model_inputs["patches_grid"] = patches_grid return model_inputs diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 1f27469eb92b..e3cd44e1c4f9 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -87,7 +87,7 @@ class UnlimitedOcrImageProcessor(DeepseekOcr2ImageProcessor): tile_size = 640 max_patches = 32 - model_input_names = ["pixel_values", "num_local_patches", "local_patches_grid"] + model_input_names = ["pixel_values", "num_local_patches", "patches_grid"] def pad_to_square( self, @@ -206,7 +206,7 @@ def _preprocess( data["pixel_values_local"] = flat_local_list # Compute per-image spatial crop grid and local-patch counts. - local_patches_grid = [] + patches_grid = [] num_local_patches = [] for image in images: height, width = image.shape[-2:] @@ -218,10 +218,10 @@ def _preprocess( else: num_columns, num_rows = 1, 1 num_local_patches.append(0) - local_patches_grid.append([num_columns, num_rows]) + patches_grid.append([num_columns, num_rows]) data["num_local_patches"] = num_local_patches - data["local_patches_grid"] = local_patches_grid + data["patches_grid"] = patches_grid return BatchFeature( data=data, @@ -241,10 +241,10 @@ class UnlimitedOcrProcessor(DeepseekOcr2Processor): def _expand_image_tokens( self, text: list[TextInput], - local_patches_grid: torch.Tensor, + patches_grid: torch.Tensor, num_local_patches: list[int] | torch.Tensor, ) -> list[str]: - num_images = len(local_patches_grid) + num_images = len(patches_grid) total_image_tokens = sum(t.count(self.image_token) for t in text) if total_image_tokens != num_images: raise ValueError( @@ -262,8 +262,8 @@ def _expand_image_tokens( crop_index = 0 for i in range(len(text)): while self.image_token in text[i]: - num_columns = int(local_patches_grid[crop_index][0]) - num_rows = int(local_patches_grid[crop_index][1]) + num_columns = int(patches_grid[crop_index][0]) + num_rows = int(patches_grid[crop_index][1]) num_tokens = num_queries_global * (num_queries_global + 1) + 1 if int(num_local_patches[crop_index]) > 0: num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) @@ -279,6 +279,19 @@ def __call__( text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, **kwargs: Unpack[UnlimitedOcrProcessorKwargs], ) -> BatchFeature: + r""" + Returns: + [`BatchFeature`]: A [`BatchFeature`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not + `None`). + - **pixel_values** -- Global view pixel values. Returned when `images` is not `None`. + - **pixel_values_local** -- Local patch pixel values. Returned when `images` is not `None`. + - **num_local_patches** -- Number of local patches per image. Returned when `images` is not `None`. + - **patches_grid** -- Number of patch columns and rows per image. Returned when `images` is not `None`. + """ if images is None: raise ValueError("`images` are expected as arguments to a `UnlimitedOcrProcessor` instance.") if text is None: @@ -298,7 +311,7 @@ def __call__( text = text.copy() # below lines change text in-place image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) - text = self._expand_image_tokens(text, image_inputs["local_patches_grid"], image_inputs["num_local_patches"]) + text = self._expand_image_tokens(text, image_inputs["patches_grid"], image_inputs["num_local_patches"]) return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) @@ -863,7 +876,7 @@ def get_image_features( pixel_values: torch.FloatTensor, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - local_patches_grid: torch.Tensor | None = None, + patches_grid: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> "UnlimitedOcrModelOutputWithPooling": r""" @@ -871,8 +884,8 @@ def get_image_features( 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]`. - local_patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): - The local crop grid `(num_columns, num_rows)` per image. + patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + The patches grid `(num_columns, num_rows)` per image. """ if isinstance(num_local_patches, torch.Tensor): num_local_patches = num_local_patches.tolist() @@ -908,7 +921,7 @@ def get_image_features( local_features = per_image_local[idx] if local_features is not None and local_features.shape[0] > 0: - num_columns, num_rows = int(local_patches_grid[idx][0]), int(local_patches_grid[idx][1]) + num_columns, num_rows = int(patches_grid[idx][0]), int(patches_grid[idx][1]) num_queries_local = int(local_features.shape[1] ** 0.5) local_grid = local_features.reshape( num_rows, num_columns, num_queries_local, num_queries_local, hidden_size @@ -944,7 +957,7 @@ def forward( pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - local_patches_grid: torch.Tensor | None = None, + patches_grid: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -957,8 +970,8 @@ def forward( 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. - local_patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): - The local crop grid `(num_columns, num_rows)` per image. + patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + The patches grid `(num_columns, num_rows)` per image. """ if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) @@ -966,7 +979,7 @@ def forward( image_features = None if pixel_values is not None: image_features = self.get_image_features( - pixel_values, pixel_values_local, num_local_patches, local_patches_grid, return_dict=True + pixel_values, pixel_values_local, num_local_patches, patches_grid, return_dict=True ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) @@ -1001,7 +1014,7 @@ def forward( pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - local_patches_grid: torch.Tensor | None = None, + patches_grid: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -1016,15 +1029,15 @@ def forward( 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. - local_patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): - The local crop grid `(num_columns, num_rows)` per image. + patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + The patches grid `(num_columns, num_rows)` per image. """ outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, pixel_values_local=pixel_values_local, num_local_patches=num_local_patches, - local_patches_grid=local_patches_grid, + patches_grid=patches_grid, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, @@ -1064,7 +1077,7 @@ def prepare_inputs_for_generation( pixel_values=None, pixel_values_local=None, num_local_patches=None, - local_patches_grid=None, + patches_grid=None, attention_mask=None, logits_to_keep=None, is_first_iteration=False, @@ -1085,7 +1098,7 @@ def prepare_inputs_for_generation( model_inputs["pixel_values"] = pixel_values model_inputs["pixel_values_local"] = pixel_values_local model_inputs["num_local_patches"] = num_local_patches - model_inputs["local_patches_grid"] = local_patches_grid + model_inputs["patches_grid"] = patches_grid return model_inputs diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index 8c51f6ede89d..0aa50c69368b 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -62,7 +62,7 @@ def __init__( def _expand_image_tokens( self, text: list[TextInput], - local_patches_grid: torch.Tensor, + patches_grid: torch.Tensor, num_local_patches: list[int] | torch.Tensor, ) -> list[str]: """ @@ -78,13 +78,13 @@ def _expand_image_tokens( Returns: `list[str]`: Text with expanded image token placeholders. """ - num_images = len(local_patches_grid) + num_images = len(patches_grid) total_image_tokens = sum(t.count(self.image_token) for t in text) if total_image_tokens != num_images: raise ValueError( f"Number of `{self.image_token}` tokens in the text ({total_image_tokens}) does not match the " f"number of images passed ({num_images}). Use one `{self.image_token}` placeholder per image, " - f"e.g. `'' * len(images) + '\\nMulti page parsing.'`" + f"e.g. `'' * len(images) + 'Multi page parsing.'`" ) size = self.image_processor.size["height"] @@ -96,8 +96,8 @@ def _expand_image_tokens( crop_index = 0 for i in range(len(text)): while self.image_token in text[i]: - num_columns = int(local_patches_grid[crop_index][0]) - num_rows = int(local_patches_grid[crop_index][1]) + num_columns = int(patches_grid[crop_index][0]) + num_rows = int(patches_grid[crop_index][1]) num_tokens = num_queries_global * (num_queries_global + 1) + 1 if int(num_local_patches[crop_index]) > 0: num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) @@ -123,6 +123,8 @@ def __call__( `None`). - **pixel_values** -- Global view pixel values. Returned when `images` is not `None`. - **pixel_values_local** -- Local patch pixel values. Returned when `images` is not `None`. + - **num_local_patches** -- Number of local patches per image. Returned when `images` is not `None`. + - **patches_grid** -- Number of patch columns and rows per image. Returned when `images` is not `None`. """ if images is None: raise ValueError("`images` are expected as arguments to a `UnlimitedOcrProcessor` instance.") @@ -143,7 +145,7 @@ def __call__( text = text.copy() # below lines change text in-place image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) - text = self._expand_image_tokens(text, image_inputs["local_patches_grid"], image_inputs["num_local_patches"]) + text = self._expand_image_tokens(text, image_inputs["patches_grid"], image_inputs["num_local_patches"]) return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) From 62065bcfc2b754bed981a1cd871c789e0967c5dd Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 18:00:15 +0200 Subject: [PATCH 051/195] cleanup --- .../unlimited_ocr/modeling_unlimited_ocr.py | 992 +++++++++--------- .../unlimited_ocr/modular_unlimited_ocr.py | 121 +-- 2 files changed, 502 insertions(+), 611 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 1661cc7f8eac..24aa78225c59 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -135,50 +135,6 @@ class UnlimitedOcrCausalLMOutputWithPast(ModelOutput): image_hidden_states: torch.FloatTensor | None = None -@auto_docstring -class UnlimitedOcrPreTrainedModel(PreTrainedModel): - config: UnlimitedOcrConfig - base_model_prefix = "model" - input_modalities = ("image", "text") - supports_gradient_checkpointing = True - _no_split_modules = [ - "UnlimitedOcrSamVisionLayer", - "UnlimitedOcrVisionEncoderLayer", - "UnlimitedOcrTextDecoderLayer", - ] - _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, UnlimitedOcrSamVisionAttention): - if module.use_rel_pos: - init.zeros_(module.rel_pos_h) - init.zeros_(module.rel_pos_w) - elif isinstance(module, UnlimitedOcrSamVisionEncoder): - if module.pos_embed is not None: - init.zeros_(module.pos_embed) - elif isinstance(module, UnlimitedOcrModel): - embed_std = 1 / math.sqrt(self.config.text_config.hidden_size) - init.normal_(module.view_separator, mean=0.0, std=embed_std) - if isinstance(module, UnlimitedOcrModel): - embed_std = 1 / math.sqrt(self.config.text_config.hidden_size) - init.normal_(module.image_newline, mean=0.0, std=embed_std) - elif isinstance(module, UnlimitedOcrVisionEmbeddings): - factor = module.config.initializer_factor - init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor) - init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor) - init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor) - init.copy_(module.position_ids, torch.arange(module.num_positions).expand((1, -1))) - - class UnlimitedOcrSamVisionAttention(nn.Module): """Multi-head Attention block with relative position embeddings.""" @@ -312,6 +268,79 @@ def forward(self, hidden_states: torch.Tensor, output_attentions=None) -> tuple[ return attn_output, attn_weights +class UnlimitedOcrSamVisionEncoder(UnlimitedOcrPreTrainedModel): + _can_record_outputs = {"hidden_states": UnlimitedOcrSamVisionLayer, "attentions": UnlimitedOcrSamVisionAttention} + + def __init__(self, config: UnlimitedOcrSamVisionConfig): + super().__init__(config) + self.config = config + self.image_size = config.image_size + self.patch_embed = UnlimitedOcrSamPatchEmbeddings(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 = UnlimitedOcrSamVisionLayer( + config, + window_size=config.window_size if i not in config.global_attn_indexes else 0, + ) + self.layers.append(layer) + + self.neck = UnlimitedOcrSamVisionNeck(config) + + self.gradient_checkpointing = False + self.proj = UnlimitedOcrSamVisionProj(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 UnlimitedOcrSamMLPBlock(nn.Module): def __init__(self, config): super().__init__() @@ -567,79 +596,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states -class UnlimitedOcrSamVisionEncoder(UnlimitedOcrPreTrainedModel): - _can_record_outputs = {"hidden_states": UnlimitedOcrSamVisionLayer, "attentions": UnlimitedOcrSamVisionAttention} - - def __init__(self, config: UnlimitedOcrSamVisionConfig): - super().__init__(config) - self.config = config - self.image_size = config.image_size - self.patch_embed = UnlimitedOcrSamPatchEmbeddings(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 = UnlimitedOcrSamVisionLayer( - config, - window_size=config.window_size if i not in config.global_attn_indexes else 0, - ) - self.layers.append(layer) - - self.neck = UnlimitedOcrSamVisionNeck(config) - - self.gradient_checkpointing = False - self.proj = UnlimitedOcrSamVisionProj(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 UnlimitedOcrVisionMLP(nn.Module): def __init__(self, config): super().__init__() @@ -917,124 +873,471 @@ def forward( return hidden_states -class UnlimitedOcrAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" +class UnlimitedOcrTextRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` - def __init__(self, config: UnlimitedOcrVisionEncoderConfig): + def __init__(self, config: UnlimitedOcrTextConfig, device=None): super().__init__() - self.config = config - self.embed_dim = config.hidden_size - self.num_heads = config.num_attention_heads - self.head_dim = self.embed_dim // self.num_heads - self.scale = self.head_dim**-0.5 - self.dropout = config.attention_dropout - self.is_causal = False + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings - self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) - self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) - self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) - self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) - # The shared `eager_attention_forward` calls `repeat_kv(..., num_key_value_groups)`; CLIP attention is - # plain multi-head attention, so the group count is 1 and `repeat_kv` becomes a no-op. - self.num_key_value_groups = 1 + self.config = config - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - **kwargs: Unpack[TransformersKwargs], - ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Input shape: Batch x Time x Channel""" + 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) - input_shape = hidden_states.shape[:-1] + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) - hidden_shape = (*input_shape, -1, self.head_dim) - queries = self.q_proj(hidden_states) - keys = self.k_proj(hidden_states) - values = self.v_proj(hidden_states) + @staticmethod + def compute_default_rope_parameters( + config: UnlimitedOcrTextConfig | 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 - queries = queries.view(hidden_shape).transpose(1, 2) - keys = keys.view(hidden_shape).transpose(1, 2) - values = values.view(hidden_shape).transpose(1, 2) + attention_factor = 1.0 # Unused in this type of RoPE - attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + # 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 UnlimitedOcrTextAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: UnlimitedOcrTextConfig, 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, - queries, - keys, - values, + query_states, + key_states, + value_states, attention_mask, - scaling=self.scale, - dropout=0.0 if not self.training else self.dropout, + 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.out_proj(attn_output) - + attn_output = self.o_proj(attn_output) return attn_output, attn_weights -class UnlimitedOcrMLP(nn.Module): +class UnlimitedOcrTextMLP(nn.Module): + def __init__(self, config: UnlimitedOcrTextConfig, 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 UnlimitedOcrTextExperts(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 UnlimitedOcrTextMoe(nn.Module): + def __init__(self, config: UnlimitedOcrTextConfig): super().__init__() self.config = config - self.activation_fn = ACT2FN[config.hidden_act] - self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) - self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + self.experts = UnlimitedOcrTextExperts(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 = UnlimitedOcrTextMLP(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: - hidden_states = self.fc1(hidden_states) - hidden_states = self.activation_fn(hidden_states) - hidden_states = self.fc2(hidden_states) + 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 -class UnlimitedOcrEncoderLayer(GradientCheckpointingLayer): - def __init__(self, config: UnlimitedOcrVisionConfig | UnlimitedOcrTextConfig): +@use_kernel_forward_from_hub("RMSNorm") +class UnlimitedOcrTextRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + UnlimitedOcrTextRMSNorm is equivalent to T5LayerNorm + """ super().__init__() - self.embed_dim = config.hidden_size - self.self_attn = UnlimitedOcrAttention(config) - self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) - self.mlp = UnlimitedOcrMLP(config) - self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + 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 UnlimitedOcrTextDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = UnlimitedOcrTextAttention(config=config, layer_idx=layer_idx) + self.mlp = ( + UnlimitedOcrTextMoe(config) + if config.mlp_layer_types[layer_idx] == "sparse" + else UnlimitedOcrTextMLP(config) + ) + + self.input_layernorm = UnlimitedOcrTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = UnlimitedOcrTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, - attention_mask: 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.FloatTensor: + ) -> torch.Tensor: residual = hidden_states - - hidden_states = self.layer_norm1(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.layer_norm2(hidden_states) + hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states - return hidden_states -class UnlimitedOcrVisionEmbeddings(nn.Module): - def __init__(self, config: UnlimitedOcrVisionConfig): - super().__init__() - self.config = config - self.embed_dim = config.hidden_size - self.image_size = config.image_size - self.patch_size = config.patch_size - - self.class_embedding = nn.Parameter(torch.randn(self.embed_dim)) +@auto_docstring +class UnlimitedOcrPreTrainedModel(PreTrainedModel): + config: UnlimitedOcrConfig + base_model_prefix = "model" + input_modalities = ("image", "text") + supports_gradient_checkpointing = True + _no_split_modules = [ + "UnlimitedOcrSamVisionLayer", + "UnlimitedOcrVisionEncoderLayer", + "UnlimitedOcrTextDecoderLayer", + ] + _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, UnlimitedOcrSamVisionAttention): + if module.use_rel_pos: + init.zeros_(module.rel_pos_h) + init.zeros_(module.rel_pos_w) + elif isinstance(module, UnlimitedOcrSamVisionEncoder): + if module.pos_embed is not None: + init.zeros_(module.pos_embed) + elif isinstance(module, UnlimitedOcrModel): + embed_std = 1 / math.sqrt(self.config.text_config.hidden_size) + init.normal_(module.view_separator, mean=0.0, std=embed_std) + if isinstance(module, UnlimitedOcrModel): + embed_std = 1 / math.sqrt(self.config.text_config.hidden_size) + init.normal_(module.image_newline, mean=0.0, std=embed_std) + elif isinstance(module, UnlimitedOcrVisionEmbeddings): + factor = module.config.initializer_factor + init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor) + init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor) + init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor) + init.copy_(module.position_ids, torch.arange(module.num_positions).expand((1, -1))) + + +class UnlimitedOcrAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: UnlimitedOcrVisionEncoderConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + self.scale = self.head_dim**-0.5 + self.dropout = config.attention_dropout + self.is_causal = False + + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) + # Required for repeat_kv(..., num_key_value_groups) + self.num_key_value_groups = 1 + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Input shape: Batch x Time x Channel""" + + input_shape = hidden_states.shape[:-1] + + hidden_shape = (*input_shape, -1, self.head_dim) + queries = self.q_proj(hidden_states) + keys = self.k_proj(hidden_states) + values = self.v_proj(hidden_states) + + queries = queries.view(hidden_shape).transpose(1, 2) + keys = keys.view(hidden_shape).transpose(1, 2) + values = values.view(hidden_shape).transpose(1, 2) + + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + + attn_output, attn_weights = attention_interface( + self, + queries, + keys, + values, + attention_mask, + scaling=self.scale, + dropout=0.0 if not self.training else self.dropout, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights + + +class UnlimitedOcrMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + + +class UnlimitedOcrEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: UnlimitedOcrVisionConfig | UnlimitedOcrTextConfig): + super().__init__() + self.embed_dim = config.hidden_size + self.self_attn = UnlimitedOcrAttention(config) + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.mlp = UnlimitedOcrMLP(config) + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.FloatTensor: + residual = hidden_states + + hidden_states = self.layer_norm1(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + **kwargs, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class UnlimitedOcrVisionEmbeddings(nn.Module): + def __init__(self, config: UnlimitedOcrVisionConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size + + self.class_embedding = nn.Parameter(torch.randn(self.embed_dim)) self.patch_embedding = nn.Conv2d( in_channels=config.num_channels, @@ -1062,6 +1365,7 @@ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: position_embedding = self.position_embedding.weight.unsqueeze(0) num_positions = position_embedding.shape[1] - 1 + # always interpolate when tracing to ensure the exported model works for dynamic input shapes if not torch.jit.is_tracing() and num_patches == num_positions and height == width: return self.position_embedding(self.position_ids) @@ -1077,12 +1381,13 @@ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) + # TODO: check if we can drop dtype cast target_dtype = patch_pos_embed.dtype patch_pos_embed = nn.functional.interpolate( patch_pos_embed.to(torch.float32), size=(new_height, new_width), mode="bicubic", - antialias=True, + antialias=True, # different from CLIP align_corners=False, ).to(target_dtype) @@ -1098,6 +1403,7 @@ def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: """ batch_size, _, grid_height, grid_width = patch_embeds.shape patch_embeds = patch_embeds.flatten(2).transpose(1, 2) + class_embeds = self.class_embedding.expand(batch_size, 1, -1) embeddings = torch.cat([class_embeds, patch_embeds], dim=1) embeddings = embeddings + self.interpolate_pos_encoding( @@ -1170,11 +1476,11 @@ def __init__(self, config: UnlimitedOcrVisionEncoderConfig): def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: r""" patch_embeds (`torch.Tensor` of shape `(batch_size, hidden_size, grid_height, grid_width)`): - The SAM feature map used in place of the CLIP patch embeddings. + Patch embeddings. """ hidden_states = self.embeddings(patch_embeds) hidden_states = self.pre_layrnorm(hidden_states) - encoder_outputs: BaseModelOutput = self.encoder(inputs_embeds=hidden_states, **kwargs) + encoder_outputs = self.encoder(inputs_embeds=hidden_states, **kwargs) return BaseModelOutput(last_hidden_state=encoder_outputs.last_hidden_state) @@ -1298,310 +1604,6 @@ def decode_kv_length(self, query_length: int = 1) -> int | None: return prefill_length + generated_after -class UnlimitedOcrTextRotaryEmbedding(nn.Module): - inv_freq: torch.Tensor # fix linting for `register_buffer` - - def __init__(self, config: UnlimitedOcrTextConfig, 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: UnlimitedOcrTextConfig | 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 UnlimitedOcrTextAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" - - def __init__(self, config: UnlimitedOcrTextConfig, 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 UnlimitedOcrTextMLP(nn.Module): - def __init__(self, config: UnlimitedOcrTextConfig, 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 UnlimitedOcrTextExperts(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 UnlimitedOcrTextMoe(nn.Module): - def __init__(self, config: UnlimitedOcrTextConfig): - super().__init__() - self.config = config - self.experts = UnlimitedOcrTextExperts(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 = UnlimitedOcrTextMLP(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 UnlimitedOcrTextRMSNorm(nn.Module): - def __init__(self, hidden_size, eps: float = 1e-6) -> None: - """ - UnlimitedOcrTextRMSNorm 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 UnlimitedOcrTextDecoderLayer(GradientCheckpointingLayer): - def __init__(self, config, layer_idx: int): - super().__init__() - self.hidden_size = config.hidden_size - self.self_attn = UnlimitedOcrTextAttention(config=config, layer_idx=layer_idx) - self.mlp = ( - UnlimitedOcrTextMoe(config) - if config.mlp_layer_types[layer_idx] == "sparse" - else UnlimitedOcrTextMLP(config) - ) - - self.input_layernorm = UnlimitedOcrTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = UnlimitedOcrTextRMSNorm(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 UnlimitedOcrTextPreTrainedModel(PreTrainedModel): config: UnlimitedOcrTextConfig @@ -1800,9 +1802,6 @@ def get_image_features( [local_grid, newline.expand(num_rows * num_queries_local, 1, hidden_size)], dim=1 ) local_flat = local_grid.reshape(-1, hidden_size) - # NOTE: local-then-global ordering matches the reference `forward` (the feature order the weights - # were trained on), NOT the token order built by the reference `infer`. - # TODO: verify correctness all_features.append(torch.cat([local_flat, global_flat, view_separator], dim=0)) else: all_features.append(torch.cat([global_flat, view_separator], dim=0)) @@ -2022,7 +2021,6 @@ def prepare_inputs_for_generation( **kwargs, ) - # Image inputs are only needed during prefill or when the cache is disabled 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 diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index e3cd44e1c4f9..74d688819030 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -55,30 +55,9 @@ DeepseekOcr2ModelOutputWithPast, DeepseekOcr2ModelOutputWithPooling, DeepseekOcr2PreTrainedModel, - DeepseekOcr2SamLayerNorm, - DeepseekOcr2SamMLPBlock, - DeepseekOcr2SamPatchEmbeddings, - DeepseekOcr2SamVisionAttention, - DeepseekOcr2SamVisionEncoder, - DeepseekOcr2SamVisionLayer, - DeepseekOcr2SamVisionNeck, - DeepseekOcr2SamVisionProj, - DeepseekOcr2SamVisionSdpaAttention, - DeepseekOcr2TextAttention, - DeepseekOcr2TextDecoderLayer, - DeepseekOcr2TextExperts, - DeepseekOcr2TextMLP, DeepseekOcr2TextModel, - DeepseekOcr2TextMoe, DeepseekOcr2TextPreTrainedModel, - DeepseekOcr2TextRMSNorm, - DeepseekOcr2TextRotaryEmbedding, - DeepseekOcr2VisionAttention, - DeepseekOcr2VisionEncoderLayer, - DeepseekOcr2VisionMLP, DeepseekOcr2VisionModel, - DeepseekOcr2VisionRMSNorm, - DeepseekOcr2VisionRotaryEmbedding, ) from ..deepseek_ocr2.processing_deepseek_ocr2 import DeepseekOcr2Processor, DeepseekOcr2ProcessorKwargs from ..got_ocr2.configuration_got_ocr2 import GotOcr2VisionConfig @@ -505,67 +484,10 @@ def _init_weights(self, module): init.copy_(module.position_ids, torch.arange(module.num_positions).expand((1, -1))) -class UnlimitedOcrSamVisionAttention(DeepseekOcr2SamVisionAttention): - pass - - -class UnlimitedOcrSamMLPBlock(DeepseekOcr2SamMLPBlock): - pass - - -class UnlimitedOcrSamVisionSdpaAttention(DeepseekOcr2SamVisionSdpaAttention): - pass - - -class UnlimitedOcrSamVisionLayer(DeepseekOcr2SamVisionLayer): - pass - - -class UnlimitedOcrSamLayerNorm(DeepseekOcr2SamLayerNorm): - pass - - -class UnlimitedOcrSamVisionNeck(DeepseekOcr2SamVisionNeck): - pass - - -class UnlimitedOcrSamPatchEmbeddings(DeepseekOcr2SamPatchEmbeddings): - pass - - -class UnlimitedOcrSamVisionProj(DeepseekOcr2SamVisionProj): - pass - - -class UnlimitedOcrSamVisionEncoder(DeepseekOcr2SamVisionEncoder): - pass - - -class UnlimitedOcrVisionMLP(DeepseekOcr2VisionMLP): - pass - - -class UnlimitedOcrVisionRMSNorm(DeepseekOcr2VisionRMSNorm): - pass - - -class UnlimitedOcrVisionRotaryEmbedding(DeepseekOcr2VisionRotaryEmbedding): - pass - - -class UnlimitedOcrVisionAttention(DeepseekOcr2VisionAttention): - pass - - -class UnlimitedOcrVisionEncoderLayer(DeepseekOcr2VisionEncoderLayer): - pass - - class UnlimitedOcrAttention(CLIPAttention): def __init__(self, config: UnlimitedOcrVisionEncoderConfig): super().__init__(config) - # The shared `eager_attention_forward` calls `repeat_kv(..., num_key_value_groups)`; CLIP attention is - # plain multi-head attention, so the group count is 1 and `repeat_kv` becomes a no-op. + # Required for repeat_kv(..., num_key_value_groups) self.num_key_value_groups = 1 @@ -579,6 +501,7 @@ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: position_embedding = self.position_embedding.weight.unsqueeze(0) num_positions = position_embedding.shape[1] - 1 + # always interpolate when tracing to ensure the exported model works for dynamic input shapes if not torch.jit.is_tracing() and num_patches == num_positions and height == width: return self.position_embedding(self.position_ids) @@ -594,12 +517,13 @@ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) + # TODO: check if we can drop dtype cast target_dtype = patch_pos_embed.dtype patch_pos_embed = nn.functional.interpolate( patch_pos_embed.to(torch.float32), size=(new_height, new_width), mode="bicubic", - antialias=True, + antialias=True, # different from CLIP align_corners=False, ).to(target_dtype) @@ -615,6 +539,7 @@ def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: """ batch_size, _, grid_height, grid_width = patch_embeds.shape patch_embeds = patch_embeds.flatten(2).transpose(1, 2) + class_embeds = self.class_embedding.expand(batch_size, 1, -1) embeddings = torch.cat([class_embeds, patch_embeds], dim=1) embeddings = embeddings + self.interpolate_pos_encoding( @@ -640,11 +565,11 @@ def __init__(self, config: UnlimitedOcrVisionEncoderConfig): def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: r""" patch_embeds (`torch.Tensor` of shape `(batch_size, hidden_size, grid_height, grid_width)`): - The SAM feature map used in place of the CLIP patch embeddings. + Patch embeddings. """ hidden_states = self.embeddings(patch_embeds) hidden_states = self.pre_layrnorm(hidden_states) - encoder_outputs: BaseModelOutput = self.encoder(inputs_embeds=hidden_states, **kwargs) + encoder_outputs = self.encoder(inputs_embeds=hidden_states, **kwargs) return BaseModelOutput(last_hidden_state=encoder_outputs.last_hidden_state) @@ -765,34 +690,6 @@ def decode_kv_length(self, query_length: int = 1) -> int | None: return prefill_length + generated_after -class UnlimitedOcrTextRotaryEmbedding(DeepseekOcr2TextRotaryEmbedding): - pass - - -class UnlimitedOcrTextAttention(DeepseekOcr2TextAttention): - pass - - -class UnlimitedOcrTextMLP(DeepseekOcr2TextMLP): - pass - - -class UnlimitedOcrTextExperts(DeepseekOcr2TextExperts): - pass - - -class UnlimitedOcrTextMoe(DeepseekOcr2TextMoe): - pass - - -class UnlimitedOcrTextRMSNorm(DeepseekOcr2TextRMSNorm): - pass - - -class UnlimitedOcrTextDecoderLayer(DeepseekOcr2TextDecoderLayer): - pass - - class UnlimitedOcrTextPreTrainedModel(DeepseekOcr2TextPreTrainedModel): pass @@ -933,9 +830,6 @@ def get_image_features( [local_grid, newline.expand(num_rows * num_queries_local, 1, hidden_size)], dim=1 ) local_flat = local_grid.reshape(-1, hidden_size) - # NOTE: local-then-global ordering matches the reference `forward` (the feature order the weights - # were trained on), NOT the token order built by the reference `infer`. - # TODO: verify correctness all_features.append(torch.cat([local_flat, global_flat, view_separator], dim=0)) else: all_features.append(torch.cat([global_flat, view_separator], dim=0)) @@ -1093,7 +987,6 @@ def prepare_inputs_for_generation( **kwargs, ) - # Image inputs are only needed during prefill or when the cache is disabled 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 From bec037c60a143e68566980fe978ee0565d5d3ed0 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 18:12:58 +0200 Subject: [PATCH 052/195] add architecture to docs --- docs/source/en/model_doc/unlimited_ocr.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 34d0eae429a8..ecd7239087ad 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -29,6 +29,8 @@ The abstract from the paper is the following: *Recently, end-to-end OCR models, exemplified by DeepSeek OCR, have once again thrust OCR into the spotlight. A widely held view is that employing a large language model (LLM) as the decoder allows the model to leverage the prior distribution of language, leading to improved OCR performance. However, the downside is equally evident: as the output sequence lengthens, the accumulated KV cache drives up memory consumption and progressively slows down generation. This stands in stark contrast to humans, who exhibit no such decline in efficiency during long-horizon copying tasks. In this technical report, we propose Unlimited OCR, a model designed to emulate human parsing working memory. Taking DeepSeek OCR as the baseline, we replace all attention layers in the decoder with our proposed Reference Sliding Window Attention (R-SWA), which reduces attention computation costs while maintaining a constant KV cache throughout the entire decoding process. By combining the high compression rate of DeepSeek OCR's encoder with our constant KV cache design, Unlimited OCR can transcribe dozens of pages of documents in a single forward pass under a standard maximum length of 32K. More importantly, R-SWA is a general-purpose parsing attention mechanism β€” beyond OCR, it is equally applicable to tasks such as ASR, translation, etc. Codes and model weights are publicly available at http://github.com/baidu/Unlimited-OCR* + + This model was contributed by [guarin](https://huggingface.co/guarin). The original code can be found [here](https://github.com/baidu/Unlimited-OCR). @@ -76,6 +78,7 @@ num_pages = 2 inputs = processor( images=[page1, page2], text="" * num_pages + "Multi page parsing.", + crop_to_patches=False, return_tensors="pt", ).to(model.device) From 1ea87e3b5c340ec59305de5227bd9714f5221e46 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 18:27:18 +0200 Subject: [PATCH 053/195] fix modular generation --- .../unlimited_ocr/modeling_unlimited_ocr.py | 1490 +++++++---------- .../unlimited_ocr/modular_unlimited_ocr.py | 5 + 2 files changed, 644 insertions(+), 851 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 24aa78225c59..f010c095f82a 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -39,7 +39,6 @@ use_kernelized_func, ) from ...masking_utils import create_causal_mask, create_reference_sliding_window_causal_mask -from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update @@ -268,77 +267,48 @@ def forward(self, hidden_states: torch.Tensor, output_attentions=None) -> tuple[ return attn_output, attn_weights -class UnlimitedOcrSamVisionEncoder(UnlimitedOcrPreTrainedModel): - _can_record_outputs = {"hidden_states": UnlimitedOcrSamVisionLayer, "attentions": UnlimitedOcrSamVisionAttention} - - def __init__(self, config: UnlimitedOcrSamVisionConfig): - super().__init__(config) - self.config = config - self.image_size = config.image_size - self.patch_embed = UnlimitedOcrSamPatchEmbeddings(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 = UnlimitedOcrSamVisionLayer( - config, - window_size=config.window_size if i not in config.global_attn_indexes else 0, - ) - self.layers.append(layer) - - self.neck = UnlimitedOcrSamVisionNeck(config) - - self.gradient_checkpointing = False - self.proj = UnlimitedOcrSamVisionProj(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) +@auto_docstring +class UnlimitedOcrPreTrainedModel(PreTrainedModel): + config: UnlimitedOcrConfig + base_model_prefix = "model" + input_modalities = ("image", "text") + supports_gradient_checkpointing = True + _no_split_modules = [ + "UnlimitedOcrSamVisionLayer", + "UnlimitedOcrVisionEncoderLayer", + "UnlimitedOcrTextDecoderLayer", + ] + _skip_keys_device_placement = ["past_key_values"] + # SAM uses rel-pos bias, incompatible with flash attention. + _supports_flash_attn = False + _supports_sdpa = True - 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 + _can_compile_fullgraph = True + _supports_flex_attn = True + _supports_attention_backend = True - 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 + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, UnlimitedOcrSamVisionAttention): + if module.use_rel_pos: + init.zeros_(module.rel_pos_h) + init.zeros_(module.rel_pos_w) + elif isinstance(module, UnlimitedOcrSamVisionEncoder): + if module.pos_embed is not None: + init.zeros_(module.pos_embed) + elif isinstance(module, UnlimitedOcrModel): + embed_std = 1 / math.sqrt(self.config.text_config.hidden_size) + init.normal_(module.view_separator, mean=0.0, std=embed_std) + if isinstance(module, UnlimitedOcrModel): + embed_std = 1 / math.sqrt(self.config.text_config.hidden_size) + init.normal_(module.image_newline, mean=0.0, std=embed_std) + elif isinstance(module, UnlimitedOcrVisionEmbeddings): + factor = module.config.initializer_factor + init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor) + init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor) + init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor) + init.copy_(module.position_ids, torch.arange(module.num_positions).expand((1, -1))) class UnlimitedOcrSamMLPBlock(nn.Module): @@ -596,218 +566,138 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states -class UnlimitedOcrVisionMLP(nn.Module): - def __init__(self, config): - super().__init__() +class UnlimitedOcrSamVisionEncoder(UnlimitedOcrPreTrainedModel): + _can_record_outputs = {"hidden_states": UnlimitedOcrSamVisionLayer, "attentions": UnlimitedOcrSamVisionAttention} + + def __init__(self, config: UnlimitedOcrSamVisionConfig): + super().__init__(config) 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] + self.image_size = config.image_size + self.patch_embed = UnlimitedOcrSamPatchEmbeddings(config) - def forward(self, x): - down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) - return down_proj + 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 = UnlimitedOcrSamVisionLayer( + config, + window_size=config.window_size if i not in config.global_attn_indexes else 0, + ) + self.layers.append(layer) -@use_kernel_forward_from_hub("RMSNorm") -class UnlimitedOcrVisionRMSNorm(nn.Module): - def __init__(self, hidden_size, eps: float = 1e-6) -> None: - """ - UnlimitedOcrVisionRMSNorm is equivalent to T5LayerNorm - """ - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps + self.neck = UnlimitedOcrSamVisionNeck(config) - 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) + self.gradient_checkpointing = False + self.proj = UnlimitedOcrSamVisionProj(config) + self.post_init() - def extra_repr(self): - return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + 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] + ) -class UnlimitedOcrVisionRotaryEmbedding(nn.Module): - inv_freq: torch.Tensor # fix linting for `register_buffer` + for layer_module in self.layers: + hidden_states = layer_module(hidden_states) - def __init__(self, config: UnlimitedOcrVisionConfig, device=None): - super().__init__() - self.max_seq_len_cached = config.max_position_embeddings - self.original_max_seq_len = config.max_position_embeddings + hidden_states = self.neck(hidden_states) + hidden_states = self.proj(hidden_states) + return BaseModelOutput(last_hidden_state=hidden_states) - self.config = config + 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 - 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) + 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 - 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: UnlimitedOcrVisionConfig | 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 +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], +): + attn_weights = torch.matmul(query, key.transpose(-1, -2)) * 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) - attention_factor = 1.0 # Unused in this type of RoPE + attn_output = torch.matmul(attn_weights, value) + attn_output = attn_output.transpose(1, 2).contiguous() + return attn_output, attn_weights - # 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() +class UnlimitedOcrAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" - 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 UnlimitedOcrVisionAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" - - def __init__(self, config: UnlimitedOcrVisionConfig, layer_idx: int): + def __init__(self, config: UnlimitedOcrVisionEncoderConfig): 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 + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + self.scale = self.head_dim**-0.5 + self.dropout = config.attention_dropout + self.is_causal = False + + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) + # Required for repeat_kv(..., num_key_value_groups) + self.num_key_value_groups = 1 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], + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor | None]: - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, -1, self.head_dim) + """Input shape: Batch x Time x Channel""" - 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) + input_shape = hidden_states.shape[:-1] - cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + hidden_shape = (*input_shape, -1, self.head_dim) + queries = self.q_proj(hidden_states) + keys = self.k_proj(hidden_states) + values = self.v_proj(hidden_states) - if past_key_values is not None: - key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + queries = queries.view(hidden_shape).transpose(1, 2) + keys = keys.view(hidden_shape).transpose(1, 2) + values = values.view(hidden_shape).transpose(1, 2) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward @@ -815,567 +705,116 @@ def forward( attn_output, attn_weights = attention_interface( self, - query_states, - key_states, - value_states, + queries, + keys, + values, 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 + scaling=self.scale, + dropout=0.0 if not self.training else self.dropout, **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() - attn_output = self.o_proj(attn_output) + attn_output = self.out_proj(attn_output) + return attn_output, attn_weights -class UnlimitedOcrVisionEncoderLayer(GradientCheckpointingLayer): - def __init__(self, config: UnlimitedOcrVisionConfig, layer_idx: int): +class UnlimitedOcrMLP(nn.Module): + def __init__(self, config): super().__init__() - self.hidden_size = config.hidden_size + self.config = config + self.activation_fn = ACT2FN[config.hidden_act] + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states - self.self_attn = UnlimitedOcrVisionAttention(config=config, layer_idx=layer_idx) - self.mlp = UnlimitedOcrVisionMLP(config) - self.input_layernorm = UnlimitedOcrVisionRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = UnlimitedOcrVisionRMSNorm(config.hidden_size, eps=config.rms_norm_eps) +class UnlimitedOcrEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: UnlimitedOcrVisionConfig | UnlimitedOcrTextConfig): + super().__init__() + self.embed_dim = config.hidden_size + self.self_attn = UnlimitedOcrAttention(config) + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.mlp = UnlimitedOcrMLP(config) + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_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, + attention_mask: torch.Tensor, **kwargs: Unpack[TransformersKwargs], - ) -> torch.Tensor: + ) -> torch.FloatTensor: residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - # Self Attention + + hidden_states = self.layer_norm1(hidden_states) 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.layer_norm2(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states - return hidden_states + return hidden_states -class UnlimitedOcrTextRotaryEmbedding(nn.Module): - inv_freq: torch.Tensor # fix linting for `register_buffer` - def __init__(self, config: UnlimitedOcrTextConfig, device=None): +class UnlimitedOcrVisionEmbeddings(nn.Module): + def __init__(self, config: UnlimitedOcrVisionConfig): super().__init__() - self.max_seq_len_cached = config.max_position_embeddings - self.original_max_seq_len = config.max_position_embeddings - self.config = config + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size - 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.class_embedding = nn.Parameter(torch.randn(self.embed_dim)) - self.register_buffer("inv_freq", inv_freq, persistent=False) - self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) + self.patch_embedding = nn.Conv2d( + in_channels=config.num_channels, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + bias=False, + ) - @staticmethod - def compute_default_rope_parameters( - config: UnlimitedOcrTextConfig | None = None, - device: Optional["torch.device"] = None, - seq_len: int | None = None, - ) -> tuple["torch.Tensor", float]: + self.num_patches = (self.image_size // self.patch_size) ** 2 + self.num_positions = self.num_patches + 1 + self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False) + + def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: """ - 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). + This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution + images. This method is also adapted to support torch.jit tracing. + + Adapted from: + - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and + - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 """ - base = config.rope_parameters["rope_theta"] - dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + num_patches = embeddings.shape[1] - 1 + position_embedding = self.position_embedding.weight.unsqueeze(0) + num_positions = position_embedding.shape[1] - 1 - attention_factor = 1.0 # Unused in this type of RoPE + # always interpolate when tracing to ensure the exported model works for dynamic input shapes + if not torch.jit.is_tracing() and num_patches == num_positions and height == width: + return self.position_embedding(self.position_ids) - # 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 + class_pos_embed = position_embedding[:, :1] + patch_pos_embed = position_embedding[:, 1:] - @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() + dim = embeddings.shape[-1] - 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 UnlimitedOcrTextAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" - - def __init__(self, config: UnlimitedOcrTextConfig, 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 UnlimitedOcrTextMLP(nn.Module): - def __init__(self, config: UnlimitedOcrTextConfig, 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 UnlimitedOcrTextExperts(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 UnlimitedOcrTextMoe(nn.Module): - def __init__(self, config: UnlimitedOcrTextConfig): - super().__init__() - self.config = config - self.experts = UnlimitedOcrTextExperts(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 = UnlimitedOcrTextMLP(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 UnlimitedOcrTextRMSNorm(nn.Module): - def __init__(self, hidden_size, eps: float = 1e-6) -> None: - """ - UnlimitedOcrTextRMSNorm 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 UnlimitedOcrTextDecoderLayer(GradientCheckpointingLayer): - def __init__(self, config, layer_idx: int): - super().__init__() - self.hidden_size = config.hidden_size - self.self_attn = UnlimitedOcrTextAttention(config=config, layer_idx=layer_idx) - self.mlp = ( - UnlimitedOcrTextMoe(config) - if config.mlp_layer_types[layer_idx] == "sparse" - else UnlimitedOcrTextMLP(config) - ) - - self.input_layernorm = UnlimitedOcrTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = UnlimitedOcrTextRMSNorm(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 UnlimitedOcrPreTrainedModel(PreTrainedModel): - config: UnlimitedOcrConfig - base_model_prefix = "model" - input_modalities = ("image", "text") - supports_gradient_checkpointing = True - _no_split_modules = [ - "UnlimitedOcrSamVisionLayer", - "UnlimitedOcrVisionEncoderLayer", - "UnlimitedOcrTextDecoderLayer", - ] - _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, UnlimitedOcrSamVisionAttention): - if module.use_rel_pos: - init.zeros_(module.rel_pos_h) - init.zeros_(module.rel_pos_w) - elif isinstance(module, UnlimitedOcrSamVisionEncoder): - if module.pos_embed is not None: - init.zeros_(module.pos_embed) - elif isinstance(module, UnlimitedOcrModel): - embed_std = 1 / math.sqrt(self.config.text_config.hidden_size) - init.normal_(module.view_separator, mean=0.0, std=embed_std) - if isinstance(module, UnlimitedOcrModel): - embed_std = 1 / math.sqrt(self.config.text_config.hidden_size) - init.normal_(module.image_newline, mean=0.0, std=embed_std) - elif isinstance(module, UnlimitedOcrVisionEmbeddings): - factor = module.config.initializer_factor - init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor) - init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor) - init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor) - init.copy_(module.position_ids, torch.arange(module.num_positions).expand((1, -1))) - - -class UnlimitedOcrAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" - - def __init__(self, config: UnlimitedOcrVisionEncoderConfig): - super().__init__() - self.config = config - self.embed_dim = config.hidden_size - self.num_heads = config.num_attention_heads - self.head_dim = self.embed_dim // self.num_heads - self.scale = self.head_dim**-0.5 - self.dropout = config.attention_dropout - self.is_causal = False - - self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) - self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) - self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) - self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) - # Required for repeat_kv(..., num_key_value_groups) - self.num_key_value_groups = 1 - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - **kwargs: Unpack[TransformersKwargs], - ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Input shape: Batch x Time x Channel""" - - input_shape = hidden_states.shape[:-1] - - hidden_shape = (*input_shape, -1, self.head_dim) - queries = self.q_proj(hidden_states) - keys = self.k_proj(hidden_states) - values = self.v_proj(hidden_states) - - queries = queries.view(hidden_shape).transpose(1, 2) - keys = keys.view(hidden_shape).transpose(1, 2) - values = values.view(hidden_shape).transpose(1, 2) - - attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( - self.config._attn_implementation, eager_attention_forward - ) - - attn_output, attn_weights = attention_interface( - self, - queries, - keys, - values, - attention_mask, - scaling=self.scale, - dropout=0.0 if not self.training else self.dropout, - **kwargs, - ) - - attn_output = attn_output.reshape(*input_shape, -1).contiguous() - attn_output = self.out_proj(attn_output) - - return attn_output, attn_weights - - -class UnlimitedOcrMLP(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.activation_fn = ACT2FN[config.hidden_act] - self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) - self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.fc1(hidden_states) - hidden_states = self.activation_fn(hidden_states) - hidden_states = self.fc2(hidden_states) - return hidden_states - - -class UnlimitedOcrEncoderLayer(GradientCheckpointingLayer): - def __init__(self, config: UnlimitedOcrVisionConfig | UnlimitedOcrTextConfig): - super().__init__() - self.embed_dim = config.hidden_size - self.self_attn = UnlimitedOcrAttention(config) - self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) - self.mlp = UnlimitedOcrMLP(config) - self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor, - **kwargs: Unpack[TransformersKwargs], - ) -> torch.FloatTensor: - residual = hidden_states - - hidden_states = self.layer_norm1(hidden_states) - hidden_states, _ = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - **kwargs, - ) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.layer_norm2(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - - return hidden_states - - -class UnlimitedOcrVisionEmbeddings(nn.Module): - def __init__(self, config: UnlimitedOcrVisionConfig): - super().__init__() - self.config = config - self.embed_dim = config.hidden_size - self.image_size = config.image_size - self.patch_size = config.patch_size - - self.class_embedding = nn.Parameter(torch.randn(self.embed_dim)) - - self.patch_embedding = nn.Conv2d( - in_channels=config.num_channels, - out_channels=self.embed_dim, - kernel_size=self.patch_size, - stride=self.patch_size, - bias=False, - ) - - self.num_patches = (self.image_size // self.patch_size) ** 2 - self.num_positions = self.num_patches + 1 - self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) - self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False) - - def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: - """ - This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution - images. This method is also adapted to support torch.jit tracing. - - Adapted from: - - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and - - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 - """ - num_patches = embeddings.shape[1] - 1 - position_embedding = self.position_embedding.weight.unsqueeze(0) - num_positions = position_embedding.shape[1] - 1 - - # always interpolate when tracing to ensure the exported model works for dynamic input shapes - if not torch.jit.is_tracing() and num_patches == num_positions and height == width: - return self.position_embedding(self.position_ids) - - class_pos_embed = position_embedding[:, :1] - patch_pos_embed = position_embedding[:, 1:] - - dim = embeddings.shape[-1] - - new_height = height // self.patch_size - new_width = width // self.patch_size + new_height = height // self.patch_size + new_width = width // self.patch_size sqrt_num_positions = torch_int(num_positions**0.5) patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) @@ -1512,96 +951,380 @@ def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwarg ) -class DynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): - """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens and windows - only the generated ones. +class DynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): + """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens and windows + only the generated ones. + + The stock `DynamicSlidingWindowLayer` evicts the oldest tokens once the window fills, which would + discard the image/prompt prefill. Here the prefill (every token cached before the first single-token + decode step) is kept intact and never evicted. Generated tokens are appended right after the prefill + until ``sliding_window`` of them have accumulated; from then on the newest generated token overwrites + the oldest one. + + While the window is filling, the cached length grows by one each step (exactly like the stock sliding + layer), so within the first ``sliding_window`` generated tokens the layer is indistinguishable from + full attention. Once the ring is full the cached tensors stay at a constant length + (``prefill + sliding_window``), which lets the cuDNN SDPA backend reuse a single kernel plan across the + (typically long) steady-state decode instead of re-planning on every distinct sequence length. + """ + + layer_type = "reference_sliding_attention" + + def __init__(self, config: PretrainedConfig | None = None, sliding_window: int | None = None): + super().__init__(config=config, sliding_window=sliding_window) + self.prefill_length: int | None = None + self.ring_position = 0 + + def update( + self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs + ) -> tuple[torch.Tensor, torch.Tensor]: + # Lazy initialization + if not self.is_initialized: + self.lazy_initialization(key_states, value_states) + + sequence_length = key_states.shape[-2] + self.cumulative_length += sequence_length + + # Prefill with prompt context + if self.prefill_length is None and sequence_length > 1: + self.keys = torch.cat([self.keys, key_states], dim=-2) + self.values = torch.cat([self.values, value_states], dim=-2) + return self.keys, self.values + + # First decode step + # Handle generation with empty prompt + if self.prefill_length is None: + self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 + + # Append while window grows + generated_length = self.keys.shape[-2] - self.prefill_length if self.keys.dim() > 1 else 0 + append_length = min(sequence_length, max(0, self.sliding_window - generated_length)) + if append_length > 0: + self.keys = torch.cat([self.keys, key_states[..., :append_length, :]], dim=-2) + self.values = torch.cat([self.values, value_states[..., :append_length, :]], dim=-2) + + # Overwrite if window size is reached + overwrite_length = sequence_length - append_length + if overwrite_length > 0: + # Only the most recent `sliding_window` overwrites survive + write_length = min(overwrite_length, self.sliding_window) + start = self.ring_position + overwrite_length - write_length + offsets = torch.arange(write_length, device=key_states.device) + slots = self.prefill_length + (start + offsets) % self.sliding_window + self.keys[..., slots, :] = key_states[..., sequence_length - write_length :, :] + self.values[..., slots, :] = value_states[..., sequence_length - write_length :, :] + self.ring_position = (self.ring_position + overwrite_length) % self.sliding_window + + return self.keys, self.values + + def get_mask_sizes(self, query_length: int) -> tuple[int, int]: + # Full visibility, no sliding offset: every query attends to all cached keys (all prefill plus the + # generated tokens currently held in the ring), so the prefill is never masked out. + if self.prefill_length is None: + return self.cumulative_length + query_length, 0 + return self.decode_kv_length(query_length), 0 + + def decode_kv_length(self, query_length: int = 1) -> int | None: + """Physical length of the cached buffer after the upcoming `update` of `query_length` tokens. + + Returns `None` before the layer is initialized. Used to size the all-visible decode mask, whose + width must match the key/value tensors returned by `update`. + """ + if not self.is_initialized: + return None + if self.prefill_length is not None: + prefill_length = self.prefill_length + generated_before = self.keys.shape[-2] - prefill_length if self.keys.dim() > 1 else 0 + else: + # Before the first decode step the whole buffer is prefill. + prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 + generated_before = 0 + generated_after = min(generated_before + query_length, self.sliding_window) + return prefill_length + generated_after + + +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) + + +@use_kernelized_func(apply_rotary_pos_emb) +class UnlimitedOcrTextAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: UnlimitedOcrTextConfig, 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 UnlimitedOcrTextMLP(nn.Module): + def __init__(self, config: UnlimitedOcrTextConfig, 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 UnlimitedOcrTextExperts(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 - The stock `DynamicSlidingWindowLayer` evicts the oldest tokens once the window fills, which would - discard the image/prompt prefill. Here the prefill (every token cached before the first single-token - decode step) is kept intact and never evicted. Generated tokens are appended right after the prefill - until ``sliding_window`` of them have accumulated; from then on the newest generated token overwrites - the oldest one. - While the window is filling, the cached length grows by one each step (exactly like the stock sliding - layer), so within the first ``sliding_window`` generated tokens the layer is indistinguishable from - full attention. Once the ring is full the cached tensors stay at a constant length - (``prefill + sliding_window``), which lets the cuDNN SDPA backend reuse a single kernel plan across the - (typically long) steady-state decode instead of re-planning on every distinct sequence length. - """ +class UnlimitedOcrTextMoe(nn.Module): + def __init__(self, config: UnlimitedOcrTextConfig): + super().__init__() + self.config = config + self.experts = UnlimitedOcrTextExperts(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 = UnlimitedOcrTextMLP(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 - layer_type = "reference_sliding_attention" + 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) - def __init__(self, config: PretrainedConfig | None = None, sliding_window: int | None = None): - super().__init__(config=config, sliding_window=sliding_window) - self.prefill_length: int | None = None - self.ring_position = 0 + topk_weight = topk_weight * self.routed_scaling_factor + return topk_idx, topk_weight - def update( - self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs - ) -> tuple[torch.Tensor, torch.Tensor]: - # Lazy initialization - if not self.is_initialized: - self.lazy_initialization(key_states, value_states) + 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 - sequence_length = key_states.shape[-2] - self.cumulative_length += sequence_length - # Prefill with prompt context - if self.prefill_length is None and sequence_length > 1: - self.keys = torch.cat([self.keys, key_states], dim=-2) - self.values = torch.cat([self.values, value_states], dim=-2) - return self.keys, self.values +@use_kernel_forward_from_hub("RMSNorm") +class UnlimitedOcrTextRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + UnlimitedOcrTextRMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps - # First decode step - # Handle generation with empty prompt - if self.prefill_length is None: - self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 + 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) - # Append while window grows - generated_length = self.keys.shape[-2] - self.prefill_length if self.keys.dim() > 1 else 0 - append_length = min(sequence_length, max(0, self.sliding_window - generated_length)) - if append_length > 0: - self.keys = torch.cat([self.keys, key_states[..., :append_length, :]], dim=-2) - self.values = torch.cat([self.values, value_states[..., :append_length, :]], dim=-2) + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" - # Overwrite if window size is reached - overwrite_length = sequence_length - append_length - if overwrite_length > 0: - # Only the most recent `sliding_window` overwrites survive - write_length = min(overwrite_length, self.sliding_window) - start = self.ring_position + overwrite_length - write_length - offsets = torch.arange(write_length, device=key_states.device) - slots = self.prefill_length + (start + offsets) % self.sliding_window - self.keys[..., slots, :] = key_states[..., sequence_length - write_length :, :] - self.values[..., slots, :] = value_states[..., sequence_length - write_length :, :] - self.ring_position = (self.ring_position + overwrite_length) % self.sliding_window - return self.keys, self.values +class UnlimitedOcrTextDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = UnlimitedOcrTextAttention(config=config, layer_idx=layer_idx) + self.mlp = ( + UnlimitedOcrTextMoe(config) + if config.mlp_layer_types[layer_idx] == "sparse" + else UnlimitedOcrTextMLP(config) + ) - def get_mask_sizes(self, query_length: int) -> tuple[int, int]: - # Full visibility, no sliding offset: every query attends to all cached keys (all prefill plus the - # generated tokens currently held in the ring), so the prefill is never masked out. - if self.prefill_length is None: - return self.cumulative_length + query_length, 0 - return self.decode_kv_length(query_length), 0 + self.input_layernorm = UnlimitedOcrTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = UnlimitedOcrTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - def decode_kv_length(self, query_length: int = 1) -> int | None: - """Physical length of the cached buffer after the upcoming `update` of `query_length` tokens. + 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 - Returns `None` before the layer is initialized. Used to size the all-visible decode mask, whose - width must match the key/value tensors returned by `update`. - """ - if not self.is_initialized: - return None - if self.prefill_length is not None: - prefill_length = self.prefill_length - generated_before = self.keys.shape[-2] - prefill_length if self.keys.dim() > 1 else 0 - else: - # Before the first decode step the whole buffer is prefill. - prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 - generated_before = 0 - generated_after = min(generated_before + query_length, self.sliding_window) - return prefill_length + generated_after + # 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 @@ -1630,6 +1353,71 @@ def _init_weights(self, module): init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range) +class UnlimitedOcrTextRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: UnlimitedOcrTextConfig, 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: UnlimitedOcrTextConfig | 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) + + @auto_docstring class UnlimitedOcrTextModel(UnlimitedOcrTextPreTrainedModel): def __init__(self, config: UnlimitedOcrTextConfig): diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 74d688819030..f1b01b1c70d6 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -55,6 +55,7 @@ DeepseekOcr2ModelOutputWithPast, DeepseekOcr2ModelOutputWithPooling, DeepseekOcr2PreTrainedModel, + DeepseekOcr2SamVisionEncoder, DeepseekOcr2TextModel, DeepseekOcr2TextPreTrainedModel, DeepseekOcr2VisionModel, @@ -484,6 +485,10 @@ def _init_weights(self, module): init.copy_(module.position_ids, torch.arange(module.num_positions).expand((1, -1))) +class UnlimitedOcrSamVisionEncoder(DeepseekOcr2SamVisionEncoder): + pass + + class UnlimitedOcrAttention(CLIPAttention): def __init__(self, config: UnlimitedOcrVisionEncoderConfig): super().__init__(config) From 033041e612c540c873efa8fe9ee20aae36594f26 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 18:39:46 +0200 Subject: [PATCH 054/195] remove mask helper --- src/transformers/masking_utils.py | 134 ------------------ .../configuration_unlimited_ocr.py | 6 + .../unlimited_ocr/modeling_unlimited_ocr.py | 4 +- .../unlimited_ocr/modular_unlimited_ocr.py | 13 +- 4 files changed, 16 insertions(+), 141 deletions(-) diff --git a/src/transformers/masking_utils.py b/src/transformers/masking_utils.py index 4f87217a6227..bf97836cb865 100644 --- a/src/transformers/masking_utils.py +++ b/src/transformers/masking_utils.py @@ -1243,140 +1243,6 @@ def create_sliding_window_causal_mask( return causal_mask -# NOTE: The only difference with create_causal_mask is that we select here sliding instead of full layers. -# TODO: Check if we really want to keep this. The only reason we need it is in case we have a hybrid cache -# with different layer types. In the default case where config.layer_types contains a single type of layer -# create_causal_mask is sufficient. -def create_reference_sliding_window_causal_mask( - config: PreTrainedConfig, - inputs_embeds: torch.Tensor, - attention_mask: torch.Tensor | None, - past_key_values: Cache | None, - position_ids: torch.Tensor | None = None, - or_mask_function: Callable | None = None, - and_mask_function: Callable | None = None, - block_sequence_ids: torch.Tensor | None = None, -) -> torch.Tensor | BlockMask | None: - """ - Create the mask for the reference-sliding layers. This mirrors `create_sliding_window_causal_mask` (it reads the - sizes from the sliding cache layer), but uses a plain causal mask function instead of a sliding window one: the - reference-sliding cache already holds the window physically as a ring buffer (prefill + the last `sliding_window` - generated tokens) and reports its physical length with offset 0. A sliding window mask function would wrongly - evict the prefill; a plain causal mask instead stays correctly causal during prefill and degenerates to full - visibility over the physical buffer during single-token decode (every cached key index <= the large query - offset) β€” which is exactly what the order-agnostic ring buffer needs (RoPE is baked into the cached keys). - - Args: - config (`PreTrainedConfig`): - The model config. - inputs_embeds (`torch.Tensor`): - The input embeddings of shape (batch_size, query_length, hidden_dim). This is used only to infer the - batch size, query length and dtype. - attention_mask (`torch.Tensor`, optional): - The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length). - It can also be an already prepared 4D mask, in which case it is returned as-is. - cache_position (`torch.Tensor`): - Deprecated and unused. - past_key_values (`Cache`, optional): - The past key values, if we use a cache. - position_ids (`torch.Tensor`, optional) - A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences. - or_mask_function (`Callable`, optional): - An optional mask function to combine with the causal mask function (by doing the union of both). This is - useful to easily overlay another mask on top of the causal one, for example for image tokens handling. - and_mask_function (`Callable`, optional): - An optional mask function to combine with the causal mask function (by doing the intersection of both). This is - useful to easily overlay another mask on top of the causal one, for example for image tokens handling. - block_sequence_ids (`torch.Tensor`, *optional*): - A tensor of same shape as input IDs indicating to which block or group each token belongs to. Tokens from - the same block will keep a bidirectional mask within the block, attending causally to the past. Index `-1` - can be used for blocks that have to keep complete causality within itself. - """ - # Power feature: if `is_causal` is False, then fallback to bi-directional mask for bi-directional attention. - # It allows to use decoder-only models with bi-directional attention as well - if not getattr(config, "is_causal", True): - return create_bidirectional_mask( - config, - inputs_embeds, - attention_mask, - past_key_values=past_key_values, - or_mask_function=or_mask_function, - and_mask_function=and_mask_function, - ) - - # NOTE: The only difference with create_causal_mask is that we select here sliding instead of full layers - # If we have an hybrid cache structure, here we want to create the mask for the sliding layers - if hasattr(past_key_values, "is_sliding") and True in past_key_values.is_sliding: - layer_idx = past_key_values.is_sliding.index(True) - else: - layer_idx = 0 - - early_exit, attention_mask, packed_sequence_mask, q_length, kv_length, q_offset, kv_offset = ( - _preprocess_mask_arguments(config, inputs_embeds, attention_mask, past_key_values, position_ids, layer_idx) - ) - if early_exit: - return attention_mask - - batch_size, dtype, device = inputs_embeds.shape[0], inputs_embeds.dtype, inputs_embeds.device - mask_factory_function = causal_mask_function - mask_interface = ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation] - - # Defaulting to using non-vmap based mask creations except when detecting - # users passing custom mask functions (as we cannot guarantee that they - # are properly index-based as required by our implementation). - use_vmap = False - - # Do not allow skip if we are compiling (this is to match BC) - # TODO: cyril -> probably revisit and remove this, but a lot of tests rely on it - if _is_torch_xpu_available: - # Do not allow skip if we are compiling for decoding, but for prefill, we still allow skip to optimization the perf of 1st token generation - allow_is_causal_skip = not (getattr(past_key_values, "is_compileable", False) and q_length == 1) - else: - allow_is_causal_skip = not getattr(past_key_values, "is_compileable", False) - - # Allow slight deviations from causal mask - # Note that it is very important to apply this before any other deviations of the mask (such as packed sequence mask, - # padding mask, etc) as the resulting mask may otherwise not be correct! - if or_mask_function is not None: - if not _is_torch_greater_or_equal_than_2_6: - raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6") - mask_factory_function = or_masks(mask_factory_function, or_mask_function) - allow_is_causal_skip = False - use_vmap = True - if and_mask_function is not None: - if not _is_torch_greater_or_equal_than_2_6: - raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6") - mask_factory_function = and_masks(mask_factory_function, and_mask_function) - allow_is_causal_skip = False - use_vmap = True - - # If we detected packing format or blockwise overlay - if packed_sequence_mask is not None: - mask_factory_function = and_masks(mask_factory_function, packed_sequence_mask_function(packed_sequence_mask)) - allow_is_causal_skip = False - if block_sequence_ids is not None: - block_sequence_ids = maybe_pad_block_sequence_ids(block_sequence_ids, attention_mask, kv_length, kv_offset) - mask_factory_function = or_masks(mask_factory_function, blockwise_overlay(block_sequence_ids)) - allow_is_causal_skip = False - - # We now create the mask - causal_mask = mask_interface( - batch_size=batch_size, - q_length=q_length, - kv_length=kv_length, - q_offset=q_offset, - kv_offset=kv_offset, - mask_function=mask_factory_function, - attention_mask=attention_mask, - allow_is_causal_skip=allow_is_causal_skip, # additional kwarg for sdpa - dtype=dtype, # Additional kwarg for eager - config=config, # Pass the config as well, in case someone wants to easily have their own mask_interface - use_vmap=use_vmap, # Short-circuit to non-vmap expansions for the mask - device=device, - ) - return causal_mask - - def create_bidirectional_sliding_window_mask( config: PreTrainedConfig, inputs_embeds: torch.Tensor, diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index f060974056d3..a3285962b91f 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -228,6 +228,12 @@ def __post_init__(self, **kwargs): self.layer_types = [ "full_attention" if self.sliding_window is None else "reference_sliding_attention" ] * self.num_hidden_layers + elif len(set(self.layer_types)) > 1: + # This requires a custom create_causal_mask implementation for reference_sliding_attention + # that fetches the first layer with is_sliding=True. + raise ValueError( + f"Combining multiple layer types ({set(self.layer_types)}) is not supported for this model." + ) if self.mlp_layer_types is None: # Some configs may use `first_k_dense_replace` instead of `layer_types`/`mlp_layer_types` first_k_dense_replace = kwargs.pop("first_k_dense_replace", 1) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index f010c095f82a..0984f2240ea6 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -38,7 +38,7 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...masking_utils import create_causal_mask, create_reference_sliding_window_causal_mask +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update @@ -1475,7 +1475,7 @@ def forward( } causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), - "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), + "reference_sliding_attention": create_causal_mask(**mask_kwargs), } hidden_states = inputs_embeds diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index f1b01b1c70d6..86c26becad61 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -24,10 +24,7 @@ from ...feature_extraction_utils import BatchFeature from ...image_transforms import group_images_by_shape, reorder_images from ...image_utils import ImageInput, PILImageResampling, SizeDict -from ...masking_utils import ( - create_causal_mask, - create_reference_sliding_window_causal_mask, -) +from ...masking_utils import create_causal_mask from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast from ...processing_utils import Unpack from ...tokenization_utils_base import PreTokenizedInput, TextInput @@ -423,6 +420,12 @@ def __post_init__(self, **kwargs): self.layer_types = [ "full_attention" if self.sliding_window is None else "reference_sliding_attention" ] * self.num_hidden_layers + elif len(set(self.layer_types)) > 1: + # This requires a custom create_causal_mask implementation for reference_sliding_attention + # that fetches the first layer with is_sliding=True. + raise ValueError( + f"Combining multiple layer types ({set(self.layer_types)}) is not supported for this model." + ) if self.mlp_layer_types is None: # Some configs may use `first_k_dense_replace` instead of `layer_types`/`mlp_layer_types` first_k_dense_replace = kwargs.pop("first_k_dense_replace", 1) @@ -738,7 +741,7 @@ def forward( } causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), - "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), + "reference_sliding_attention": create_causal_mask(**mask_kwargs), } hidden_states = inputs_embeds From f5b37fa166c8687adda393f02f59b778af863359 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 18:59:20 +0200 Subject: [PATCH 055/195] add image processing tests --- .../test_image_processing_unlimited_ocr.py | 134 ++++++++++++++++++ 1 file changed, 134 insertions(+) create mode 100644 tests/models/unlimited_ocr/test_image_processing_unlimited_ocr.py diff --git a/tests/models/unlimited_ocr/test_image_processing_unlimited_ocr.py b/tests/models/unlimited_ocr/test_image_processing_unlimited_ocr.py new file mode 100644 index 000000000000..0a906c1bb3d1 --- /dev/null +++ b/tests/models/unlimited_ocr/test_image_processing_unlimited_ocr.py @@ -0,0 +1,134 @@ +# Copyright 2026 The HuggingFace Inc. 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 unittest + +from transformers import is_vision_available +from transformers.testing_utils import require_torch, require_vision + +from ...test_image_processing_common import ImageProcessingTestMixin, prepare_image_inputs + + +if is_vision_available(): + pass + + +class UnlimitedOcrImageProcessingTester: + def __init__( + self, + parent, + batch_size=7, + num_channels=3, + image_size=18, + min_resolution=500, + max_resolution=800, + do_resize=True, + size=None, + tile_size=384, + do_normalize=True, + image_mean=[0.5, 0.5, 0.5], + image_std=[0.5, 0.5, 0.5], + do_convert_rgb=True, + ): + size = size if size is not None else {"height": 512, "width": 512} + self.parent = parent + self.batch_size = batch_size + self.num_channels = num_channels + self.image_size = image_size + self.min_resolution = min_resolution + self.max_resolution = max_resolution + self.do_resize = do_resize + self.size = size + self.tile_size = tile_size + self.do_normalize = do_normalize + self.image_mean = image_mean + self.image_std = image_std + self.do_convert_rgb = do_convert_rgb + + def prepare_image_processor_dict(self): + return { + "do_resize": self.do_resize, + "size": self.size, + "tile_size": self.tile_size, + "do_normalize": self.do_normalize, + "image_mean": self.image_mean, + "image_std": self.image_std, + "do_convert_rgb": self.do_convert_rgb, + } + + def expected_output_image_shape(self, images): + return self.num_channels, self.size["height"], self.size["width"] + + def prepare_image_inputs(self, equal_resolution=False, numpify=False, torchify=False): + return prepare_image_inputs( + batch_size=self.batch_size, + num_channels=self.num_channels, + min_resolution=self.min_resolution, + max_resolution=self.max_resolution, + equal_resolution=equal_resolution, + numpify=numpify, + torchify=torchify, + ) + + +@require_torch +@require_vision +class UnlimitedOcrImageProcessingTest(ImageProcessingTestMixin, unittest.TestCase): + def setUp(self): + super().setUp() + self.image_processor_tester = UnlimitedOcrImageProcessingTester(self) + + @property + def image_processor_dict(self): + return self.image_processor_tester.prepare_image_processor_dict() + + @unittest.skip(reason="Not supported") + def test_call_numpy_4_channels(self): + pass + + def test_preprocess_crop_to_patches(self): + for image_processing_class in self.image_processing_classes.values(): + image_processor = image_processing_class(**self.image_processor_dict) + num_channels = self.image_processor_tester.num_channels + size = self.image_processor_tester.size + tile_size = self.image_processor_tester.tile_size + + image = self.image_processor_tester.prepare_image_inputs(equal_resolution=False, torchify=True)[0] + processed_image = image_processor.preprocess(images=image, crop_to_patches=True) + + num_local_patches = processed_image["num_local_patches"][0] + num_columns, num_rows = processed_image["patches_grid"][0] + self.assertEqual(num_local_patches, num_columns * num_rows) + + self.assertEqual(len(processed_image["pixel_values"]), 1) + self.assertEqual(processed_image["pixel_values"][0].shape, (num_channels, size["height"], size["width"])) + + self.assertEqual(len(processed_image["pixel_values_local"]), num_local_patches) + for local_patch in processed_image["pixel_values_local"]: + self.assertEqual(local_patch.shape, (num_channels, tile_size, tile_size)) + + def test_preprocess_no_crop_to_patches(self): + for image_processing_class in self.image_processing_classes.values(): + image_processor = image_processing_class(**self.image_processor_dict) + num_channels = self.image_processor_tester.num_channels + size = self.image_processor_tester.size + + image = self.image_processor_tester.prepare_image_inputs(equal_resolution=False, torchify=True)[0] + processed_image = image_processor.preprocess(images=image, crop_to_patches=False) + + self.assertEqual(processed_image["num_local_patches"][0], 0) + self.assertNotIn("pixel_values_local", processed_image) + + self.assertEqual(len(processed_image["pixel_values"]), 1) + self.assertEqual(processed_image["pixel_values"][0].shape, (num_channels, size["height"], size["width"])) From 7b4db9210fbdcb89ac766b6a478329a2092dcfbb Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 25 Jun 2026 19:02:37 +0200 Subject: [PATCH 056/195] remove redundant var --- .../models/unlimited_ocr/image_processing_unlimited_ocr.py | 3 --- src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 3 --- 2 files changed, 6 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index 617ed1edd3aa..571f7457f98f 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -224,7 +224,6 @@ def _preprocess( **kwargs, ) -> BatchFeature: # --- Local patches (batched by shape group) --- - num_local_patches_grouped = {} local_patches_grouped = {} if crop_to_patches: @@ -245,10 +244,8 @@ def _preprocess( flat_patches, do_rescale, rescale_factor, do_normalize, image_mean, image_std ) local_patches_grouped[shape] = flat_patches.reshape(stacked_patches.shape) - num_local_patches_grouped[shape] = [n_patches] * stacked_images.shape[0] else: local_patches_grouped[shape] = [None] * stacked_images.shape[0] - num_local_patches_grouped[shape] = [0] * stacked_images.shape[0] ordered_local = reorder_images(local_patches_grouped, grouped_images_index) else: diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 86c26becad61..aba5ebbc6c42 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -128,7 +128,6 @@ def _preprocess( **kwargs, ) -> BatchFeature: # --- Local patches (batched by shape group) --- - num_local_patches_grouped = {} local_patches_grouped = {} if crop_to_patches: @@ -149,10 +148,8 @@ def _preprocess( flat_patches, do_rescale, rescale_factor, do_normalize, image_mean, image_std ) local_patches_grouped[shape] = flat_patches.reshape(stacked_patches.shape) - num_local_patches_grouped[shape] = [n_patches] * stacked_images.shape[0] else: local_patches_grouped[shape] = [None] * stacked_images.shape[0] - num_local_patches_grouped[shape] = [0] * stacked_images.shape[0] ordered_local = reorder_images(local_patches_grouped, grouped_images_index) else: From 031cfd3449662b88f50313389c8a887c35ef943d Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 26 Jun 2026 07:45:15 +0200 Subject: [PATCH 057/195] fix-repo --- .../configuration_unlimited_ocr.py | 9 +++ .../unlimited_ocr/modeling_unlimited_ocr.py | 56 +++++++++++-------- 2 files changed, 41 insertions(+), 24 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index a3285962b91f..15caa6e68fa5 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -210,6 +210,15 @@ class UnlimitedOcrTextConfig(PreTrainedConfig): attention_dropout: float | None = 0.0 mlp_bias: bool = False head_dim: int | None = None + base_model_ep_plan = { + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_tp_experts", + } + attribute_map = { + "num_experts": "n_routed_experts", + } n_group: int | None = 1 n_routed_experts: int = 64 n_shared_experts: int = 2 diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 0984f2240ea6..57541a82e860 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1178,7 +1178,7 @@ class UnlimitedOcrTextExperts(nn.Module): def __init__(self, config): super().__init__() - self.num_experts = config.n_routed_experts + self.num_experts = config.num_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)) @@ -1212,48 +1212,54 @@ def forward( return final_hidden_states -class UnlimitedOcrTextMoe(nn.Module): +class UnlimitedOcrTextTopkRouter(nn.Module): def __init__(self, config: UnlimitedOcrTextConfig): super().__init__() - self.config = config - self.experts = UnlimitedOcrTextExperts(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 = UnlimitedOcrTextMLP(config=config, intermediate_size=intermediate_size) + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_experts + self.hidden_dim = config.hidden_size + self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim)) 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) + def forward(self, hidden_states): + hidden_states = hidden_states.view(-1, self.hidden_dim) + router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) + scores = 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) + topk_weights, topk_indices = torch.topk(scores, 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_scores = scores.view(-1, self.num_group, self.num_experts // self.num_group).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) + .expand(-1, self.num_group, self.num_experts // self.num_group) + .reshape(-1, self.num_experts) ) - 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) + scores = scores.masked_fill(~score_mask.bool(), 0.0) + topk_weights, topk_indices = torch.topk(scores, k=self.top_k, dim=-1, sorted=False) + topk_weights = topk_weights * self.routed_scaling_factor + return router_logits, topk_weights, topk_indices - topk_weight = topk_weight * self.routed_scaling_factor - return topk_idx, topk_weight + +class UnlimitedOcrTextMoe(nn.Module): + def __init__(self, config: UnlimitedOcrTextConfig): + super().__init__() + self.config = config + self.experts = UnlimitedOcrTextExperts(config) + self.gate = UnlimitedOcrTextTopkRouter(config) + self.shared_experts = UnlimitedOcrTextMLP( + config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts + ) 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) + _, topk_weights, topk_indices = self.gate(hidden_states) 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) @@ -1348,7 +1354,9 @@ class UnlimitedOcrTextPreTrainedModel(PreTrainedModel): @torch.no_grad() def _init_weights(self, module): super()._init_weights(module) - if isinstance(module, UnlimitedOcrTextExperts): + if isinstance(module, UnlimitedOcrTextTopkRouter): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif isinstance(module, UnlimitedOcrTextExperts): 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) From 6b14e3d0debc37e2497e096bd9f6a169131a1a82 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 26 Jun 2026 09:32:14 +0200 Subject: [PATCH 058/195] update docstring --- .../unlimited_ocr/modeling_unlimited_ocr.py | 19 +++++-------------- .../unlimited_ocr/modular_unlimited_ocr.py | 19 +++++-------------- 2 files changed, 10 insertions(+), 28 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 57541a82e860..a47915b2c3f1 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -952,20 +952,11 @@ def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwarg class DynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): - """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens and windows - only the generated ones. - - The stock `DynamicSlidingWindowLayer` evicts the oldest tokens once the window fills, which would - discard the image/prompt prefill. Here the prefill (every token cached before the first single-token - decode step) is kept intact and never evicted. Generated tokens are appended right after the prefill - until ``sliding_window`` of them have accumulated; from then on the newest generated token overwrites - the oldest one. - - While the window is filling, the cached length grows by one each step (exactly like the stock sliding - layer), so within the first ``sliding_window`` generated tokens the layer is indistinguishable from - full attention. Once the ring is full the cached tensors stay at a constant length - (``prefill + sliding_window``), which lets the cuDNN SDPA backend reuse a single kernel plan across the - (typically long) steady-state decode instead of re-planning on every distinct sequence length. + """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens before + the first decode step and applies a sliding window to all decoded tokens. + + Once ``sliding_window`` decode tokens have accumulated, the oldest decode tokens are evicted and + replaced by the most recent ones. The prefill tokens always remain in the cache. """ layer_type = "reference_sliding_attention" diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index aba5ebbc6c42..6c66d37d2cc4 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -604,20 +604,11 @@ def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwarg class DynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): - """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens and windows - only the generated ones. - - The stock `DynamicSlidingWindowLayer` evicts the oldest tokens once the window fills, which would - discard the image/prompt prefill. Here the prefill (every token cached before the first single-token - decode step) is kept intact and never evicted. Generated tokens are appended right after the prefill - until ``sliding_window`` of them have accumulated; from then on the newest generated token overwrites - the oldest one. - - While the window is filling, the cached length grows by one each step (exactly like the stock sliding - layer), so within the first ``sliding_window`` generated tokens the layer is indistinguishable from - full attention. Once the ring is full the cached tensors stay at a constant length - (``prefill + sliding_window``), which lets the cuDNN SDPA backend reuse a single kernel plan across the - (typically long) steady-state decode instead of re-planning on every distinct sequence length. + """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens before + the first decode step and applies a sliding window to all decoded tokens. + + Once ``sliding_window`` decode tokens have accumulated, the oldest decode tokens are evicted and + replaced by the most recent ones. The prefill tokens always remain in the cache. """ layer_type = "reference_sliding_attention" From 2d0505d1caa030e41db8cd9b29e093796c9d3ea2 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 26 Jun 2026 13:54:13 +0000 Subject: [PATCH 059/195] add broken static layer --- src/transformers/cache_utils.py | 240 +++++++++++++++++- .../unlimited_ocr/modeling_unlimited_ocr.py | 132 ++++------ .../unlimited_ocr/modular_unlimited_ocr.py | 131 ++++------ 3 files changed, 327 insertions(+), 176 deletions(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index fe37838eca83..429c379490dc 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -291,6 +291,89 @@ def crop(self, max_length: int) -> None: self.cumulative_length = self.keys.shape[-2] +class DynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): + """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens before + the first decode step and applies a sliding window to all decoded tokens. + + Once ``sliding_window`` decode tokens have accumulated, the oldest decode tokens are evicted and + replaced by the most recent ones. The prefill tokens always remain in the cache. + """ + + layer_type = "reference_sliding_attention" + + def __init__(self, config: PreTrainedConfig | None = None, sliding_window: int | None = None): + super().__init__(config=config, sliding_window=sliding_window) + self.prefill_length: int | None = None + self.ring_position = 0 + + def update( + self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs + ) -> tuple[torch.Tensor, torch.Tensor]: + # Lazy initialization + if not self.is_initialized: + self.lazy_initialization(key_states, value_states) + + sequence_length = key_states.shape[-2] + self.cumulative_length += sequence_length + + # Prefill with prompt context + if self.prefill_length is None and sequence_length > 1: + self.keys = torch.cat([self.keys, key_states], dim=-2) + self.values = torch.cat([self.values, value_states], dim=-2) + return self.keys, self.values + + # First decode step + # Handle generation with empty prompt + if self.prefill_length is None: + self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 + + # Append while window grows + generated_length = self.keys.shape[-2] - self.prefill_length if self.keys.dim() > 1 else 0 + append_length = min(sequence_length, max(0, self.sliding_window - generated_length)) + if append_length > 0: + self.keys = torch.cat([self.keys, key_states[..., :append_length, :]], dim=-2) + self.values = torch.cat([self.values, value_states[..., :append_length, :]], dim=-2) + + # Overwrite if window size is reached + overwrite_length = sequence_length - append_length + if overwrite_length > 0: + # Only the most recent `sliding_window` overwrites survive + write_length = min(overwrite_length, self.sliding_window) + start = self.ring_position + overwrite_length - write_length + offsets = torch.arange(write_length, device=key_states.device) + slots = self.prefill_length + (start + offsets) % self.sliding_window + self.keys[..., slots, :] = key_states[..., sequence_length - write_length :, :] + self.values[..., slots, :] = value_states[..., sequence_length - write_length :, :] + self.ring_position = (self.ring_position + overwrite_length) % self.sliding_window + + return self.keys, self.values + + def get_mask_sizes(self, query_length: int) -> tuple[int, int]: + # Full visibility, no sliding offset: every query attends to all cached keys (all prefill plus the + # generated tokens currently held in the ring), so the prefill is never masked out. + if self.prefill_length is None: + return self.cumulative_length + query_length, 0 + return self.decode_kv_length(query_length), 0 + + def decode_kv_length(self, query_length: int = 1) -> int | None: + """Physical length of the cached buffer after the upcoming `update` of `query_length` tokens. + + Returns `None` before the layer is initialized. Used to size the all-visible decode mask, whose + width must match the key/value tensors returned by `update`. + """ + if not self.is_initialized: + return None + if self.prefill_length is not None: + prefill_length = self.prefill_length + generated_before = self.keys.shape[-2] - prefill_length if self.keys.dim() > 1 else 0 + else: + # Before the first decode step the whole buffer is prefill. + prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 + generated_before = 0 + generated_after = min(generated_before + query_length, self.sliding_window) + return prefill_length + generated_after + + class DynamicIndexedLayer(DynamicLayer): """ A cache layer that extends `DynamicLayer` with an extra indexer key cache for Dynamic Sparse Attention (DSA) @@ -600,6 +683,150 @@ def reset(self): self.cumulative_length_int = 0 +class StaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): + """Static counterpart of [`DynamicReferenceSlidingWindowLayer`], used when generating with + ``cache_implementation="static"``. + + The backing buffer has the full, fixed ``max_cache_len`` length and is split in two regions: the first + ``max_cache_len - sliding_window`` slots hold the image/prompt prefill (written once, at the front, and + never evicted) and the last ``sliding_window`` slots are a ring for the generated tokens. Exactly like + [`StaticSlidingWindowLayer`] the oldest generated token is rolled out and replaced by the newest once the + ring is full, the only difference being that the roll is restricted to that trailing window so the prefill + region stays untouched. Because the buffer length is constant on every decode step, the key/value tensors + keep a constant shape and the SDPA backend (cuDNN included) reuses a single kernel plan instead of + re-planning on each new sequence length. The unused prefill slots and the not-yet-filled window slots are + hidden by the decode mask built by the model (e.g. ``UnlimitedOcrTextModel.forward``). + """ + + layer_type = "reference_sliding_attention" + + def __init__(self, max_cache_len: int, sliding_window: int): + super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window) + # `StaticSlidingWindowLayer` shrinks its buffer to `sliding_window` and rolls the whole thing. Here the + # buffer must also hold the pinned prefill, so keep the full `max_cache_len` length and track the real + # window separately; the roll below only touches the trailing `sliding_window` slots. + self.max_cache_len = max_cache_len + self.sliding_window = sliding_window + self.prefill_length: int | None = None + + def update( + self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs + ) -> tuple[torch.Tensor, torch.Tensor]: + # Lazy initialization + if not self.is_initialized: + self.lazy_initialization(key_states, value_states) + + kv_length = key_states.shape[-2] + + # Prefill. Mirror `DynamicReferenceSlidingWindowLayer`: every multi-token chunk fed before the first + # single-token (decode) step is prefill and stays pinned at the front of the buffer. Crucially we do + # *not* pin `prefill_length` here, so chunked prefill (a prompt fed in more than one chunk) keeps + # extending the pinned region instead of having the later chunks slid out of the window. + if self.prefill_length is None and kv_length > 1: + # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs + # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten + cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length + try: + self.keys.index_copy_(2, cache_position, key_states) + self.values.index_copy_(2, cache_position, value_states) + except NotImplementedError: + # Fallback for devices like MPS where index_copy_ might not be supported. + self.keys[:, :, cache_position] = key_states + self.values[:, :, cache_position] = value_states + + # Keep both the int (control flow) and the tensor (cudagraph-safe indexing) versions in sync. + self.cumulative_length_int += kv_length + self.cumulative_length.add_(kv_length) + + # Very important to return the `self` tensors here, as they have the static dynamo address + return self.keys, self.values + + # First single-token step (or an empty prompt that skipped prefill above): finalize the pinned prefill + # region so the trailing window below is measured relative to it. + if self.prefill_length is None: + self.prefill_length = self.cumulative_length_int + + # Everything below mirrors `StaticSlidingWindowLayer.update` with the difference that the sliding window + # is in the last `sliding_window` slots instead of spanning the whole buffer. + window_start = self.max_cache_len - self.sliding_window + current_length = self.cumulative_length_int - self.prefill_length + is_full = current_length >= self.sliding_window + # Update it now that we saved the value above + self.cumulative_length_int += kv_length + + if is_full: + # In general, we should use a much simpler `cat` here as well, independently of the states size. However, + # dynamo is currently bugged when doing it - see https://github.com/pytorch/pytorch/issues/159855 for more details + if key_states.shape[-2] == 1: + # Roll the window region to the left by 1 position (the pinned prefill in front stays put) + new_keys = self.keys[:, :, window_start:, :].roll(-1, dims=-2) + new_values = self.values[:, :, window_start:, :].roll(-1, dims=-2) + # Overwrite the last position with new states + # (note: very important to use a tensor to index here, see https://github.com/pytorch/pytorch/issues/159855) + index = torch.tensor([-1], dtype=int, device=self.device) + new_keys[:, :, index] = key_states + new_values[:, :, index] = value_states + + # Copy back into `self` (do not just assign again) in order to keep the static dynamo address + self.keys[:, :, window_start:, :].copy_(new_keys) + self.values[:, :, window_start:, :].copy_(new_values) + + # Very important to return the `self` tensors here, as they have the static dynamo address + return self.keys, self.values + # Already full but using more than 1 new token (e.g. prefill caching, chat continuation, etc...) + else: + full_key_states = torch.cat((self.keys[:, :, window_start + 1 :, :], key_states), dim=-2) + full_value_states = torch.cat((self.values[:, :, window_start + 1 :, :], value_states), dim=-2) + # Not yet full, but becoming full on this update + elif current_length + kv_length > self.sliding_window: + # Fast path, no need to cat() in this case, as the window is currently empty + if current_length == 0: + full_key_states = key_states + full_value_states = value_states + else: + window = slice(window_start, window_start + current_length) + full_key_states = torch.cat((self.keys[:, :, window, :], key_states), dim=-2) + full_value_states = torch.cat((self.values[:, :, window, :], value_states), dim=-2) + else: + # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs + # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten. + # `window_start - prefill_length` is constant during decode, so it is safe to bake into the graph. + cache_position = ( + torch.arange(kv_length, device=self.device) + + self.cumulative_length + + (window_start - self.prefill_length) + ) + try: + self.keys.index_copy_(2, cache_position, key_states) + self.values.index_copy_(2, cache_position, value_states) + except NotImplementedError: + # Fallback for devices like MPS where index_copy_ might not be supported. + self.keys[:, :, cache_position] = key_states + self.values[:, :, cache_position] = value_states + + # Update the tensor version of the length in-place (we don't need to update it if we are already outside + # of this branch, as we don't need the tensor anymore) + self.cumulative_length.add_(kv_length) + + # Very important to return the `self` tensors here, as they have the static dynamo address + return self.keys, self.values + + # We only keep the last `sliding_window` tokens. Unlike `StaticSlidingWindowLayer` we return `self.keys/values` + # because `get_mask_sizes` is fixed to the full `max_cache_len`. + self.keys[:, :, window_start:, :].copy_(full_key_states[:, :, -self.sliding_window :, :]) + self.values[:, :, window_start:, :].copy_(full_value_states[:, :, -self.sliding_window :, :]) + return self.keys, self.values + + def get_mask_sizes(self, query_length: int) -> tuple[int, int]: + # The buffer always spans the full `max_cache_len`. + # The decode mask hides the unwritten slots. + return self.max_cache_len, 0 + + def reset(self) -> None: + super().reset() + self.prefill_length = None + + class StaticIndexedLayer(StaticLayer): """ A `StaticLayer` with an additional statically-allocated indexer key cache for Dynamic Sparse @@ -1638,14 +1865,21 @@ def __init__( layer = StaticSlidingWindowLayer( max_cache_len=max_cache_len, sliding_window=config.attention_chunk_size ) + # Custom layer types (e.g. M3's sparse-attention indexer cache, R-SWA's reference sliding window) + # that registered a static variant. Checked before the generic sliding branch below so a custom + # static layer whose dynamic counterpart subclasses `DynamicSlidingWindowLayer` (and thus appears in + # `sliding_layer_types`) is not silently downgraded to a plain `StaticSlidingWindowLayer`. + elif layer_type in LAYER_TYPE_STATIC_CACHE_MAPPING: + static_layer_cls = LAYER_TYPE_STATIC_CACHE_MAPPING[layer_type] + if issubclass(static_layer_cls, StaticSlidingWindowLayer): + layer = static_layer_cls(max_cache_len=max_cache_len, sliding_window=config.sliding_window) + else: + layer = static_layer_cls(max_cache_len=max_cache_len) elif layer_type in sliding_layer_types: layer = StaticSlidingWindowLayer(max_cache_len=max_cache_len, sliding_window=config.sliding_window) # LinearAttention layers are static by essence - using `"moe"` as well is a trick, see the comment about it on DynamicCache elif layer_type in ("mamba", "conv", "linear_attention", "moe"): layer = LinearAttentionLayer() - # Custom layer types (e.g. M3's sparse-attention indexer cache) that registered a static variant. - elif layer_type in LAYER_TYPE_STATIC_CACHE_MAPPING: - layer = LAYER_TYPE_STATIC_CACHE_MAPPING[layer_type](max_cache_len=max_cache_len) elif layer_type == "deepseek_sparse_attention": # Static / compile-friendly indexed layer (preallocated indexer key cache). layer = StaticIndexedLayer(max_cache_len=max_cache_len) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index a47915b2c3f1..f4347247f90d 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -29,8 +29,7 @@ from ... import initialization as init from ...activations import ACT2FN -from ...cache_utils import Cache, DynamicCache, DynamicSlidingWindowLayer -from ...configuration_utils import PretrainedConfig +from ...cache_utils import Cache, DynamicCache, StaticReferenceSlidingWindowLayer from ...generation import GenerationMixin from ...integrations import ( use_experts_implementation, @@ -38,7 +37,7 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...masking_utils import create_causal_mask +from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update @@ -951,89 +950,6 @@ def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwarg ) -class DynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): - """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens before - the first decode step and applies a sliding window to all decoded tokens. - - Once ``sliding_window`` decode tokens have accumulated, the oldest decode tokens are evicted and - replaced by the most recent ones. The prefill tokens always remain in the cache. - """ - - layer_type = "reference_sliding_attention" - - def __init__(self, config: PretrainedConfig | None = None, sliding_window: int | None = None): - super().__init__(config=config, sliding_window=sliding_window) - self.prefill_length: int | None = None - self.ring_position = 0 - - def update( - self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs - ) -> tuple[torch.Tensor, torch.Tensor]: - # Lazy initialization - if not self.is_initialized: - self.lazy_initialization(key_states, value_states) - - sequence_length = key_states.shape[-2] - self.cumulative_length += sequence_length - - # Prefill with prompt context - if self.prefill_length is None and sequence_length > 1: - self.keys = torch.cat([self.keys, key_states], dim=-2) - self.values = torch.cat([self.values, value_states], dim=-2) - return self.keys, self.values - - # First decode step - # Handle generation with empty prompt - if self.prefill_length is None: - self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 - - # Append while window grows - generated_length = self.keys.shape[-2] - self.prefill_length if self.keys.dim() > 1 else 0 - append_length = min(sequence_length, max(0, self.sliding_window - generated_length)) - if append_length > 0: - self.keys = torch.cat([self.keys, key_states[..., :append_length, :]], dim=-2) - self.values = torch.cat([self.values, value_states[..., :append_length, :]], dim=-2) - - # Overwrite if window size is reached - overwrite_length = sequence_length - append_length - if overwrite_length > 0: - # Only the most recent `sliding_window` overwrites survive - write_length = min(overwrite_length, self.sliding_window) - start = self.ring_position + overwrite_length - write_length - offsets = torch.arange(write_length, device=key_states.device) - slots = self.prefill_length + (start + offsets) % self.sliding_window - self.keys[..., slots, :] = key_states[..., sequence_length - write_length :, :] - self.values[..., slots, :] = value_states[..., sequence_length - write_length :, :] - self.ring_position = (self.ring_position + overwrite_length) % self.sliding_window - - return self.keys, self.values - - def get_mask_sizes(self, query_length: int) -> tuple[int, int]: - # Full visibility, no sliding offset: every query attends to all cached keys (all prefill plus the - # generated tokens currently held in the ring), so the prefill is never masked out. - if self.prefill_length is None: - return self.cumulative_length + query_length, 0 - return self.decode_kv_length(query_length), 0 - - def decode_kv_length(self, query_length: int = 1) -> int | None: - """Physical length of the cached buffer after the upcoming `update` of `query_length` tokens. - - Returns `None` before the layer is initialized. Used to size the all-visible decode mask, whose - width must match the key/value tensors returned by `update`. - """ - if not self.is_initialized: - return None - if self.prefill_length is not None: - prefill_length = self.prefill_length - generated_before = self.keys.shape[-2] - prefill_length if self.keys.dim() > 1 else 0 - else: - # Before the first decode step the whole buffer is prefill. - prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 - generated_before = 0 - generated_after = min(generated_before + query_length, self.sliding_window) - return prefill_length + generated_after - - def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -1472,9 +1388,51 @@ def forward( "past_key_values": past_key_values, "position_ids": position_ids, } + # The static cache ([`StaticReferenceSlidingWindowLayer`]) keeps the prefill at the front of its + # buffer (slots `[0, prefill_length)`) and the most recent `sliding_window` generated tokens in the + # trailing window (slots `[max_cache_len - sliding_window, max_cache_len)`). Those physical slots do + # not line up with the tokens' logical positions, so the reference sliding window cannot be expressed + # as a position-based causal/sliding mask; we build it directly over the physical buffer slots + # instead. The dynamic cache evicts old tokens from its (compact) buffer, so the plain causal mask + # built by the fallback below is already correct for it. + reference_attention_mask = None + if past_key_values is not None: + static_reference_layer = next( + ( + layer + for layer in past_key_values.layers + if isinstance(layer, StaticReferenceSlidingWindowLayer) + ), + None, + ) + # Only decode steps (a single query token) read from the trailing window. Prefill (more than one + # token) is plain causal and is handled by the `create_causal_mask` fallback below. + if static_reference_layer is not None and inputs_embeds.shape[1] == 1: + prefill_length = static_reference_layer.prefill_length + if prefill_length is None: + # First decode step: `prefill_length` is pinned during this step's cache update, which + # runs after the mask is built, so derive it here -- every token seen so far is prefill. + prefill_length = static_reference_layer.cumulative_length_int + sliding_window = static_reference_layer.sliding_window + window_start = static_reference_layer.max_cache_len - sliding_window + # Number of trailing-window slots holding a token once this step's token has been written. + decode_length = static_reference_layer.cumulative_length_int - prefill_length + window_end = window_start + min(decode_length + 1, sliding_window) + + def reference_window_mask_function(batch_idx, head_idx, q_idx, kv_idx): + return (kv_idx < prefill_length) | ((kv_idx >= window_start) & (kv_idx < window_end)) + + # A bidirectional (all-visible) base intersected with the physical-slot selection yields + # exactly the prefill plus the filled-window slots. A causal base would instead mask the + # whole trailing window out, since its physical indices exceed the query's logical position. + reference_attention_mask = create_bidirectional_mask( + **mask_kwargs, and_mask_function=reference_window_mask_function + ) + if reference_attention_mask is None: + reference_attention_mask = create_causal_mask(**mask_kwargs) causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), - "reference_sliding_attention": create_causal_mask(**mask_kwargs), + "reference_sliding_attention": reference_attention_mask, } hidden_states = inputs_embeds diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 6c66d37d2cc4..1203915293c4 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -19,12 +19,12 @@ from torchvision.transforms.v2 import functional as tvF from ... import initialization as init -from ...cache_utils import Cache, DynamicCache, DynamicSlidingWindowLayer +from ...cache_utils import Cache, DynamicCache, StaticReferenceSlidingWindowLayer from ...configuration_utils import PretrainedConfig from ...feature_extraction_utils import BatchFeature from ...image_transforms import group_images_by_shape, reorder_images from ...image_utils import ImageInput, PILImageResampling, SizeDict -from ...masking_utils import create_causal_mask +from ...masking_utils import create_bidirectional_mask, create_causal_mask from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast from ...processing_utils import Unpack from ...tokenization_utils_base import PreTokenizedInput, TextInput @@ -603,89 +603,6 @@ def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwarg ) -class DynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): - """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens before - the first decode step and applies a sliding window to all decoded tokens. - - Once ``sliding_window`` decode tokens have accumulated, the oldest decode tokens are evicted and - replaced by the most recent ones. The prefill tokens always remain in the cache. - """ - - layer_type = "reference_sliding_attention" - - def __init__(self, config: PretrainedConfig | None = None, sliding_window: int | None = None): - super().__init__(config=config, sliding_window=sliding_window) - self.prefill_length: int | None = None - self.ring_position = 0 - - def update( - self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs - ) -> tuple[torch.Tensor, torch.Tensor]: - # Lazy initialization - if not self.is_initialized: - self.lazy_initialization(key_states, value_states) - - sequence_length = key_states.shape[-2] - self.cumulative_length += sequence_length - - # Prefill with prompt context - if self.prefill_length is None and sequence_length > 1: - self.keys = torch.cat([self.keys, key_states], dim=-2) - self.values = torch.cat([self.values, value_states], dim=-2) - return self.keys, self.values - - # First decode step - # Handle generation with empty prompt - if self.prefill_length is None: - self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 - - # Append while window grows - generated_length = self.keys.shape[-2] - self.prefill_length if self.keys.dim() > 1 else 0 - append_length = min(sequence_length, max(0, self.sliding_window - generated_length)) - if append_length > 0: - self.keys = torch.cat([self.keys, key_states[..., :append_length, :]], dim=-2) - self.values = torch.cat([self.values, value_states[..., :append_length, :]], dim=-2) - - # Overwrite if window size is reached - overwrite_length = sequence_length - append_length - if overwrite_length > 0: - # Only the most recent `sliding_window` overwrites survive - write_length = min(overwrite_length, self.sliding_window) - start = self.ring_position + overwrite_length - write_length - offsets = torch.arange(write_length, device=key_states.device) - slots = self.prefill_length + (start + offsets) % self.sliding_window - self.keys[..., slots, :] = key_states[..., sequence_length - write_length :, :] - self.values[..., slots, :] = value_states[..., sequence_length - write_length :, :] - self.ring_position = (self.ring_position + overwrite_length) % self.sliding_window - - return self.keys, self.values - - def get_mask_sizes(self, query_length: int) -> tuple[int, int]: - # Full visibility, no sliding offset: every query attends to all cached keys (all prefill plus the - # generated tokens currently held in the ring), so the prefill is never masked out. - if self.prefill_length is None: - return self.cumulative_length + query_length, 0 - return self.decode_kv_length(query_length), 0 - - def decode_kv_length(self, query_length: int = 1) -> int | None: - """Physical length of the cached buffer after the upcoming `update` of `query_length` tokens. - - Returns `None` before the layer is initialized. Used to size the all-visible decode mask, whose - width must match the key/value tensors returned by `update`. - """ - if not self.is_initialized: - return None - if self.prefill_length is not None: - prefill_length = self.prefill_length - generated_before = self.keys.shape[-2] - prefill_length if self.keys.dim() > 1 else 0 - else: - # Before the first decode step the whole buffer is prefill. - prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 - generated_before = 0 - generated_after = min(generated_before + query_length, self.sliding_window) - return prefill_length + generated_after - - class UnlimitedOcrTextPreTrainedModel(DeepseekOcr2TextPreTrainedModel): pass @@ -727,9 +644,51 @@ def forward( "past_key_values": past_key_values, "position_ids": position_ids, } + # The static cache ([`StaticReferenceSlidingWindowLayer`]) keeps the prefill at the front of its + # buffer (slots `[0, prefill_length)`) and the most recent `sliding_window` generated tokens in the + # trailing window (slots `[max_cache_len - sliding_window, max_cache_len)`). Those physical slots do + # not line up with the tokens' logical positions, so the reference sliding window cannot be expressed + # as a position-based causal/sliding mask; we build it directly over the physical buffer slots + # instead. The dynamic cache evicts old tokens from its (compact) buffer, so the plain causal mask + # built by the fallback below is already correct for it. + reference_attention_mask = None + if past_key_values is not None: + static_reference_layer = next( + ( + layer + for layer in past_key_values.layers + if isinstance(layer, StaticReferenceSlidingWindowLayer) + ), + None, + ) + # Only decode steps (a single query token) read from the trailing window. Prefill (more than one + # token) is plain causal and is handled by the `create_causal_mask` fallback below. + if static_reference_layer is not None and inputs_embeds.shape[1] == 1: + prefill_length = static_reference_layer.prefill_length + if prefill_length is None: + # First decode step: `prefill_length` is pinned during this step's cache update, which + # runs after the mask is built, so derive it here -- every token seen so far is prefill. + prefill_length = static_reference_layer.cumulative_length_int + sliding_window = static_reference_layer.sliding_window + window_start = static_reference_layer.max_cache_len - sliding_window + # Number of trailing-window slots holding a token once this step's token has been written. + decode_length = static_reference_layer.cumulative_length_int - prefill_length + window_end = window_start + min(decode_length + 1, sliding_window) + + def reference_window_mask_function(batch_idx, head_idx, q_idx, kv_idx): + return (kv_idx < prefill_length) | ((kv_idx >= window_start) & (kv_idx < window_end)) + + # A bidirectional (all-visible) base intersected with the physical-slot selection yields + # exactly the prefill plus the filled-window slots. A causal base would instead mask the + # whole trailing window out, since its physical indices exceed the query's logical position. + reference_attention_mask = create_bidirectional_mask( + **mask_kwargs, and_mask_function=reference_window_mask_function + ) + if reference_attention_mask is None: + reference_attention_mask = create_causal_mask(**mask_kwargs) causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), - "reference_sliding_attention": create_causal_mask(**mask_kwargs), + "reference_sliding_attention": reference_attention_mask, } hidden_states = inputs_embeds From f2f5cab50117de61ab6728c6a85d30608ddedb25 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 26 Jun 2026 14:15:54 +0000 Subject: [PATCH 060/195] make static cache work --- src/transformers/cache_utils.py | 60 +++++++++++-------- .../unlimited_ocr/modeling_unlimited_ocr.py | 52 +++++++--------- .../unlimited_ocr/modular_unlimited_ocr.py | 52 +++++++--------- 3 files changed, 78 insertions(+), 86 deletions(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index 429c379490dc..5305fcd6dbe4 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -687,15 +687,17 @@ class StaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): """Static counterpart of [`DynamicReferenceSlidingWindowLayer`], used when generating with ``cache_implementation="static"``. - The backing buffer has the full, fixed ``max_cache_len`` length and is split in two regions: the first - ``max_cache_len - sliding_window`` slots hold the image/prompt prefill (written once, at the front, and - never evicted) and the last ``sliding_window`` slots are a ring for the generated tokens. Exactly like - [`StaticSlidingWindowLayer`] the oldest generated token is rolled out and replaced by the newest once the - ring is full, the only difference being that the roll is restricted to that trailing window so the prefill - region stays untouched. Because the buffer length is constant on every decode step, the key/value tensors - keep a constant shape and the SDPA backend (cuDNN included) reuses a single kernel plan instead of - re-planning on each new sequence length. The unused prefill slots and the not-yet-filled window slots are - hidden by the decode mask built by the model (e.g. ``UnlimitedOcrTextModel.forward``). + The backing buffer has the full, fixed ``max_cache_len`` length and is split in three regions: slots + ``[0, prefill_length)`` hold the image/prompt prefill (written once, at the front, and never evicted), + slots ``[prefill_length, prefill_length + sliding_window)`` are a ring for the generated tokens, and the + trailing slots ``[prefill_length + sliding_window, max_cache_len)`` are an unused empty tail (the caller + must size ``max_cache_len >= prefill_length + sliding_window``). Exactly like [`StaticSlidingWindowLayer`] + the oldest generated token is rolled out and replaced by the newest once the ring is full, the only + difference being that the roll is restricted to that window so the prefill region stays untouched. Because + the buffer length is constant on every decode step, the key/value tensors keep a constant shape and the + SDPA backend (cuDNN included) reuses a single kernel plan instead of re-planning on each new sequence + length. The not-yet-written window slots and the empty tail are hidden by the decode mask built by the + model (e.g. ``UnlimitedOcrTextModel.forward``). """ layer_type = "reference_sliding_attention" @@ -704,7 +706,8 @@ def __init__(self, max_cache_len: int, sliding_window: int): super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window) # `StaticSlidingWindowLayer` shrinks its buffer to `sliding_window` and rolls the whole thing. Here the # buffer must also hold the pinned prefill, so keep the full `max_cache_len` length and track the real - # window separately; the roll below only touches the trailing `sliding_window` slots. + # window separately; the roll below only touches the `sliding_window` slots right after the prefill + # (i.e. `[prefill_length, prefill_length + sliding_window)`). self.max_cache_len = max_cache_len self.sliding_window = sliding_window self.prefill_length: int | None = None @@ -747,8 +750,8 @@ def update( self.prefill_length = self.cumulative_length_int # Everything below mirrors `StaticSlidingWindowLayer.update` with the difference that the sliding window - # is in the last `sliding_window` slots instead of spanning the whole buffer. - window_start = self.max_cache_len - self.sliding_window + # is in the `sliding_window` slots right after the prefill instead of spanning the whole buffer. + window_start = self.prefill_length current_length = self.cumulative_length_int - self.prefill_length is_full = current_length >= self.sliding_window # Update it now that we saved the value above @@ -759,8 +762,8 @@ def update( # dynamo is currently bugged when doing it - see https://github.com/pytorch/pytorch/issues/159855 for more details if key_states.shape[-2] == 1: # Roll the window region to the left by 1 position (the pinned prefill in front stays put) - new_keys = self.keys[:, :, window_start:, :].roll(-1, dims=-2) - new_values = self.values[:, :, window_start:, :].roll(-1, dims=-2) + new_keys = self.keys[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) + new_values = self.values[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) # Overwrite the last position with new states # (note: very important to use a tensor to index here, see https://github.com/pytorch/pytorch/issues/159855) index = torch.tensor([-1], dtype=int, device=self.device) @@ -768,15 +771,19 @@ def update( new_values[:, :, index] = value_states # Copy back into `self` (do not just assign again) in order to keep the static dynamo address - self.keys[:, :, window_start:, :].copy_(new_keys) - self.values[:, :, window_start:, :].copy_(new_values) + self.keys[:, :, window_start : window_start + self.sliding_window, :].copy_(new_keys) + self.values[:, :, window_start : window_start + self.sliding_window, :].copy_(new_values) # Very important to return the `self` tensors here, as they have the static dynamo address return self.keys, self.values # Already full but using more than 1 new token (e.g. prefill caching, chat continuation, etc...) else: - full_key_states = torch.cat((self.keys[:, :, window_start + 1 :, :], key_states), dim=-2) - full_value_states = torch.cat((self.values[:, :, window_start + 1 :, :], value_states), dim=-2) + full_key_states = torch.cat( + (self.keys[:, :, window_start + 1 : window_start + self.sliding_window, :], key_states), dim=-2 + ) + full_value_states = torch.cat( + (self.values[:, :, window_start + 1 : window_start + self.sliding_window, :], value_states), dim=-2 + ) # Not yet full, but becoming full on this update elif current_length + kv_length > self.sliding_window: # Fast path, no need to cat() in this case, as the window is currently empty @@ -790,12 +797,9 @@ def update( else: # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten. - # `window_start - prefill_length` is constant during decode, so it is safe to bake into the graph. - cache_position = ( - torch.arange(kv_length, device=self.device) - + self.cumulative_length - + (window_start - self.prefill_length) - ) + # The window sits right after the prefill (`window_start == prefill_length`), so the decode tokens are + # written at their logical positions, which already start at `prefill_length`. + cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length try: self.keys.index_copy_(2, cache_position, key_states) self.values.index_copy_(2, cache_position, value_states) @@ -813,8 +817,12 @@ def update( # We only keep the last `sliding_window` tokens. Unlike `StaticSlidingWindowLayer` we return `self.keys/values` # because `get_mask_sizes` is fixed to the full `max_cache_len`. - self.keys[:, :, window_start:, :].copy_(full_key_states[:, :, -self.sliding_window :, :]) - self.values[:, :, window_start:, :].copy_(full_value_states[:, :, -self.sliding_window :, :]) + self.keys[:, :, window_start : window_start + self.sliding_window, :].copy_( + full_key_states[:, :, -self.sliding_window :, :] + ) + self.values[:, :, window_start : window_start + self.sliding_window, :].copy_( + full_value_states[:, :, -self.sliding_window :, :] + ) return self.keys, self.values def get_mask_sizes(self, query_length: int) -> tuple[int, int]: diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index f4347247f90d..27fabddc44ba 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -37,7 +37,7 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...masking_utils import create_bidirectional_mask, create_causal_mask +from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update @@ -1389,13 +1389,13 @@ def forward( "position_ids": position_ids, } # The static cache ([`StaticReferenceSlidingWindowLayer`]) keeps the prefill at the front of its - # buffer (slots `[0, prefill_length)`) and the most recent `sliding_window` generated tokens in the - # trailing window (slots `[max_cache_len - sliding_window, max_cache_len)`). Those physical slots do - # not line up with the tokens' logical positions, so the reference sliding window cannot be expressed - # as a position-based causal/sliding mask; we build it directly over the physical buffer slots - # instead. The dynamic cache evicts old tokens from its (compact) buffer, so the plain causal mask - # built by the fallback below is already correct for it. - reference_attention_mask = None + # buffer (slots `[0, prefill_length)`) followed by the most recent `sliding_window` generated tokens + # (slots `[prefill_length, prefill_length + sliding_window)`) and an empty tail. The valid region is + # therefore the contiguous prefix `[0, prefill_length + sliding_window)`: causal already hides the + # not-yet-written window slots while the window fills, and a constant `kv_idx < window_end` cap hides + # the empty tail. The dynamic cache evicts old tokens from its (compact) buffer instead, so the plain + # causal mask built without the cap is already correct for it. + reference_window_mask_function = None if past_key_values is not None: static_reference_layer = next( ( @@ -1405,34 +1405,26 @@ def forward( ), None, ) - # Only decode steps (a single query token) read from the trailing window. Prefill (more than one - # token) is plain causal and is handled by the `create_causal_mask` fallback below. - if static_reference_layer is not None and inputs_embeds.shape[1] == 1: + if static_reference_layer is not None: prefill_length = static_reference_layer.prefill_length - if prefill_length is None: + if prefill_length is None and inputs_embeds.shape[1] == 1: # First decode step: `prefill_length` is pinned during this step's cache update, which # runs after the mask is built, so derive it here -- every token seen so far is prefill. prefill_length = static_reference_layer.cumulative_length_int - sliding_window = static_reference_layer.sliding_window - window_start = static_reference_layer.max_cache_len - sliding_window - # Number of trailing-window slots holding a token once this step's token has been written. - decode_length = static_reference_layer.cumulative_length_int - prefill_length - window_end = window_start + min(decode_length + 1, sliding_window) - - def reference_window_mask_function(batch_idx, head_idx, q_idx, kv_idx): - return (kv_idx < prefill_length) | ((kv_idx >= window_start) & (kv_idx < window_end)) - - # A bidirectional (all-visible) base intersected with the physical-slot selection yields - # exactly the prefill plus the filled-window slots. A causal base would instead mask the - # whole trailing window out, since its physical indices exceed the query's logical position. - reference_attention_mask = create_bidirectional_mask( - **mask_kwargs, and_mask_function=reference_window_mask_function - ) - if reference_attention_mask is None: - reference_attention_mask = create_causal_mask(**mask_kwargs) + if prefill_length is not None: + # `window_end` is constant for the whole generation, so the cap does not trigger + # per-step recompiles. Prefill (`prefill_length is None`, more than one query token) + # keeps `reference_window_mask_function=None` and so stays plain causal. + window_end = prefill_length + self.config.sliding_window + + def reference_window_mask_function(batch_idx, head_idx, q_idx, kv_idx): + return kv_idx < window_end + causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), - "reference_sliding_attention": reference_attention_mask, + "reference_sliding_attention": create_causal_mask( + **mask_kwargs, and_mask_function=reference_window_mask_function + ), } hidden_states = inputs_embeds diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 1203915293c4..c61855face38 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -24,7 +24,7 @@ from ...feature_extraction_utils import BatchFeature from ...image_transforms import group_images_by_shape, reorder_images from ...image_utils import ImageInput, PILImageResampling, SizeDict -from ...masking_utils import create_bidirectional_mask, create_causal_mask +from ...masking_utils import create_causal_mask from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast from ...processing_utils import Unpack from ...tokenization_utils_base import PreTokenizedInput, TextInput @@ -645,13 +645,13 @@ def forward( "position_ids": position_ids, } # The static cache ([`StaticReferenceSlidingWindowLayer`]) keeps the prefill at the front of its - # buffer (slots `[0, prefill_length)`) and the most recent `sliding_window` generated tokens in the - # trailing window (slots `[max_cache_len - sliding_window, max_cache_len)`). Those physical slots do - # not line up with the tokens' logical positions, so the reference sliding window cannot be expressed - # as a position-based causal/sliding mask; we build it directly over the physical buffer slots - # instead. The dynamic cache evicts old tokens from its (compact) buffer, so the plain causal mask - # built by the fallback below is already correct for it. - reference_attention_mask = None + # buffer (slots `[0, prefill_length)`) followed by the most recent `sliding_window` generated tokens + # (slots `[prefill_length, prefill_length + sliding_window)`) and an empty tail. The valid region is + # therefore the contiguous prefix `[0, prefill_length + sliding_window)`: causal already hides the + # not-yet-written window slots while the window fills, and a constant `kv_idx < window_end` cap hides + # the empty tail. The dynamic cache evicts old tokens from its (compact) buffer instead, so the plain + # causal mask built without the cap is already correct for it. + reference_window_mask_function = None if past_key_values is not None: static_reference_layer = next( ( @@ -661,34 +661,26 @@ def forward( ), None, ) - # Only decode steps (a single query token) read from the trailing window. Prefill (more than one - # token) is plain causal and is handled by the `create_causal_mask` fallback below. - if static_reference_layer is not None and inputs_embeds.shape[1] == 1: + if static_reference_layer is not None: prefill_length = static_reference_layer.prefill_length - if prefill_length is None: + if prefill_length is None and inputs_embeds.shape[1] == 1: # First decode step: `prefill_length` is pinned during this step's cache update, which # runs after the mask is built, so derive it here -- every token seen so far is prefill. prefill_length = static_reference_layer.cumulative_length_int - sliding_window = static_reference_layer.sliding_window - window_start = static_reference_layer.max_cache_len - sliding_window - # Number of trailing-window slots holding a token once this step's token has been written. - decode_length = static_reference_layer.cumulative_length_int - prefill_length - window_end = window_start + min(decode_length + 1, sliding_window) - - def reference_window_mask_function(batch_idx, head_idx, q_idx, kv_idx): - return (kv_idx < prefill_length) | ((kv_idx >= window_start) & (kv_idx < window_end)) - - # A bidirectional (all-visible) base intersected with the physical-slot selection yields - # exactly the prefill plus the filled-window slots. A causal base would instead mask the - # whole trailing window out, since its physical indices exceed the query's logical position. - reference_attention_mask = create_bidirectional_mask( - **mask_kwargs, and_mask_function=reference_window_mask_function - ) - if reference_attention_mask is None: - reference_attention_mask = create_causal_mask(**mask_kwargs) + if prefill_length is not None: + # `window_end` is constant for the whole generation, so the cap does not trigger + # per-step recompiles. Prefill (`prefill_length is None`, more than one query token) + # keeps `reference_window_mask_function=None` and so stays plain causal. + window_end = prefill_length + self.config.sliding_window + + def reference_window_mask_function(batch_idx, head_idx, q_idx, kv_idx): + return kv_idx < window_end + causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), - "reference_sliding_attention": reference_attention_mask, + "reference_sliding_attention": create_causal_mask( + **mask_kwargs, and_mask_function=reference_window_mask_function + ), } hidden_states = inputs_embeds From c60d24c8a59cffff71d116a35892ca9eae4a1483 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Sat, 27 Jun 2026 09:33:25 +0200 Subject: [PATCH 061/195] test --- src/transformers/cache_utils.py | 87 +++++++++++++++---- .../unlimited_ocr/modeling_unlimited_ocr.py | 41 ++------- .../unlimited_ocr/modular_unlimited_ocr.py | 41 ++------- 3 files changed, 85 insertions(+), 84 deletions(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index 5305fcd6dbe4..a4e24fb48166 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -687,17 +687,21 @@ class StaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): """Static counterpart of [`DynamicReferenceSlidingWindowLayer`], used when generating with ``cache_implementation="static"``. - The backing buffer has the full, fixed ``max_cache_len`` length and is split in three regions: slots - ``[0, prefill_length)`` hold the image/prompt prefill (written once, at the front, and never evicted), - slots ``[prefill_length, prefill_length + sliding_window)`` are a ring for the generated tokens, and the - trailing slots ``[prefill_length + sliding_window, max_cache_len)`` are an unused empty tail (the caller - must size ``max_cache_len >= prefill_length + sliding_window``). Exactly like [`StaticSlidingWindowLayer`] - the oldest generated token is rolled out and replaced by the newest once the ring is full, the only - difference being that the roll is restricted to that window so the prefill region stays untouched. Because - the buffer length is constant on every decode step, the key/value tensors keep a constant shape and the - SDPA backend (cuDNN included) reuses a single kernel plan instead of re-planning on each new sequence - length. The not-yet-written window slots and the empty tail are hidden by the decode mask built by the - model (e.g. ``UnlimitedOcrTextModel.forward``). + The backing buffer is split in two regions: slots ``[0, prefill_length)`` hold the image/prompt prefill + (written once, at the front, and never evicted) and slots ``[prefill_length, prefill_length + sliding_window)`` + are a ring for the generated tokens. Exactly like [`StaticSlidingWindowLayer`] the oldest generated token is + rolled out and replaced by the newest once the ring is full, the only difference being that the roll is + restricted to that window so the prefill region stays untouched. + + The physical buffer is sized to ``prefill_length + sliding_window`` rather than the full ``max_cache_len`` + budget that ``generate`` passes (which is ``prefill_length + max_new_tokens``): the trailing + ``max_new_tokens - sliding_window`` slots would never be written, so reserving them is pure waste -- large for + long generations. ``prefill_length`` is only known at runtime, but the buffer is allocated lazily during the + eager prefill pass (and grown there for chunked prefill), which always runs *before* the first compiled/ + cudagraph decode step; from the first decode step onward the shape and data pointer are constant, so cudagraph + capture is unaffected and the SDPA backend (cuDNN included) reuses a single kernel plan. The not-yet-written + window slots are hidden by the (plain causal) decode mask built by the model (e.g. + ``UnlimitedOcrTextModel.forward``); there is no empty tail left to hide. """ layer_type = "reference_sliding_attention" @@ -705,13 +709,51 @@ class StaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): def __init__(self, max_cache_len: int, sliding_window: int): super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window) # `StaticSlidingWindowLayer` shrinks its buffer to `sliding_window` and rolls the whole thing. Here the - # buffer must also hold the pinned prefill, so keep the full `max_cache_len` length and track the real - # window separately; the roll below only touches the `sliding_window` slots right after the prefill + # buffer must also hold the pinned prefill. We keep `max_cache_len` only as the *logical* budget (for + # length bookkeeping) and size the *physical* buffer to `prefill_length + sliding_window`, tracking the + # real window separately; the roll below only touches the `sliding_window` slots right after the prefill # (i.e. `[prefill_length, prefill_length + sliding_window)`). self.max_cache_len = max_cache_len self.sliding_window = sliding_window self.prefill_length: int | None = None + def _set_buffers(self, physical_length: int, copy_existing: bool = False) -> None: + """(Re)allocate the static key/value buffers to `physical_length` slots and (re)tag the static address. + + Only ever called from the eager prefill pass (initial allocation or chunked-prefill growth), never from a + compiled decode step, so reallocating here is safe for cudagraphs. + """ + new_keys = torch.zeros( + (self.batch_size, self.num_heads, physical_length, self.k_head_dim), dtype=self.dtype, device=self.device + ) + new_values = torch.zeros( + (self.batch_size, self.num_heads, physical_length, self.v_head_dim), dtype=self.dtype, device=self.device + ) + if copy_existing: + old_length = self.keys.shape[-2] + new_keys[:, :, :old_length, :] = self.keys + new_values[:, :, :old_length, :] = self.values + self.keys = new_keys + self.values = new_values + if not is_torchdynamo_compiling(): + torch._dynamo.mark_static_address(self.keys) + torch._dynamo.mark_static_address(self.values) + + def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None: + self.dtype, self.device = key_states.dtype, key_states.device + self.batch_size, self.num_heads = key_states.shape[:2] + self.v_head_dim = value_states.shape[-1] + self.k_head_dim = key_states.shape[-1] + self.cumulative_length = self.cumulative_length.to(self.device) + if not is_torchdynamo_compiling(): + torch._dynamo.mark_static_address(self.cumulative_length) + # Size the physical buffer from the prefill instead of the full `max_cache_len` budget (see class + # docstring). A multi-token first call is prefill; a single-token first call is an empty-prompt decode + # (`prefill_length == 0`), so only the window is needed. + prefill_seen = key_states.shape[-2] if key_states.shape[-2] > 1 else 0 + self._set_buffers(min(self.max_cache_len, prefill_seen + self.sliding_window)) + self.is_initialized = True + def update( self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs ) -> tuple[torch.Tensor, torch.Tensor]: @@ -726,6 +768,11 @@ def update( # *not* pin `prefill_length` here, so chunked prefill (a prompt fed in more than one chunk) keeps # extending the pinned region instead of having the later chunks slid out of the window. if self.prefill_length is None and kv_length > 1: + # Chunked prefill: the buffer was sized for the first chunk only, so grow it (eagerly, before any + # compiled decode step) to hold all prefill seen so far plus this chunk plus the reserved window. + required_length = min(self.max_cache_len, self.cumulative_length_int + kv_length + self.sliding_window) + if self.keys.shape[-2] < required_length: + self._set_buffers(required_length, copy_existing=True) # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length @@ -826,9 +873,17 @@ def update( return self.keys, self.values def get_mask_sizes(self, query_length: int) -> tuple[int, int]: - # The buffer always spans the full `max_cache_len`. - # The decode mask hides the unwritten slots. - return self.max_cache_len, 0 + # `kv_length` must match the length of the tensor `update` will return for this step. The buffer is sized + # to `prefill_length + sliding_window`, so anticipate the size `update` will (re)allocate it to: + if not self.is_initialized: + # First forward: `lazy_initialization` will allocate from this query length. + prefill_seen = query_length if query_length > 1 else 0 + return min(self.max_cache_len, prefill_seen + self.sliding_window), 0 + if self.prefill_length is None and query_length > 1: + # Additional prefill chunk: `update` will grow the buffer to fit all prefill so far plus the window. + return min(self.max_cache_len, self.cumulative_length_int + query_length + self.sliding_window), 0 + # Decode (including the first decode step): the buffer is already at its final size. + return self.keys.shape[-2], 0 def reset(self) -> None: super().reset() diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 27fabddc44ba..afd334acb5fb 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -29,7 +29,7 @@ from ... import initialization as init from ...activations import ACT2FN -from ...cache_utils import Cache, DynamicCache, StaticReferenceSlidingWindowLayer +from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...integrations import ( use_experts_implementation, @@ -1390,41 +1390,14 @@ def forward( } # The static cache ([`StaticReferenceSlidingWindowLayer`]) keeps the prefill at the front of its # buffer (slots `[0, prefill_length)`) followed by the most recent `sliding_window` generated tokens - # (slots `[prefill_length, prefill_length + sliding_window)`) and an empty tail. The valid region is - # therefore the contiguous prefix `[0, prefill_length + sliding_window)`: causal already hides the - # not-yet-written window slots while the window fills, and a constant `kv_idx < window_end` cap hides - # the empty tail. The dynamic cache evicts old tokens from its (compact) buffer instead, so the plain - # causal mask built without the cap is already correct for it. - reference_window_mask_function = None - if past_key_values is not None: - static_reference_layer = next( - ( - layer - for layer in past_key_values.layers - if isinstance(layer, StaticReferenceSlidingWindowLayer) - ), - None, - ) - if static_reference_layer is not None: - prefill_length = static_reference_layer.prefill_length - if prefill_length is None and inputs_embeds.shape[1] == 1: - # First decode step: `prefill_length` is pinned during this step's cache update, which - # runs after the mask is built, so derive it here -- every token seen so far is prefill. - prefill_length = static_reference_layer.cumulative_length_int - if prefill_length is not None: - # `window_end` is constant for the whole generation, so the cap does not trigger - # per-step recompiles. Prefill (`prefill_length is None`, more than one query token) - # keeps `reference_window_mask_function=None` and so stays plain causal. - window_end = prefill_length + self.config.sliding_window - - def reference_window_mask_function(batch_idx, head_idx, q_idx, kv_idx): - return kv_idx < window_end - + # (slots `[prefill_length, prefill_length + sliding_window)`). The buffer is sized to exactly + # `prefill_length + sliding_window` (no empty tail), so a plain causal mask is already correct: it + # hides the not-yet-written window slots while the window fills, and there is no tail left to cap. The + # dynamic cache evicts old tokens from its (compact) buffer instead, so plain causal is correct there + # too. causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), - "reference_sliding_attention": create_causal_mask( - **mask_kwargs, and_mask_function=reference_window_mask_function - ), + "reference_sliding_attention": create_causal_mask(**mask_kwargs), } hidden_states = inputs_embeds diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index c61855face38..6b8d081ee956 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -19,7 +19,7 @@ from torchvision.transforms.v2 import functional as tvF from ... import initialization as init -from ...cache_utils import Cache, DynamicCache, StaticReferenceSlidingWindowLayer +from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PretrainedConfig from ...feature_extraction_utils import BatchFeature from ...image_transforms import group_images_by_shape, reorder_images @@ -646,41 +646,14 @@ def forward( } # The static cache ([`StaticReferenceSlidingWindowLayer`]) keeps the prefill at the front of its # buffer (slots `[0, prefill_length)`) followed by the most recent `sliding_window` generated tokens - # (slots `[prefill_length, prefill_length + sliding_window)`) and an empty tail. The valid region is - # therefore the contiguous prefix `[0, prefill_length + sliding_window)`: causal already hides the - # not-yet-written window slots while the window fills, and a constant `kv_idx < window_end` cap hides - # the empty tail. The dynamic cache evicts old tokens from its (compact) buffer instead, so the plain - # causal mask built without the cap is already correct for it. - reference_window_mask_function = None - if past_key_values is not None: - static_reference_layer = next( - ( - layer - for layer in past_key_values.layers - if isinstance(layer, StaticReferenceSlidingWindowLayer) - ), - None, - ) - if static_reference_layer is not None: - prefill_length = static_reference_layer.prefill_length - if prefill_length is None and inputs_embeds.shape[1] == 1: - # First decode step: `prefill_length` is pinned during this step's cache update, which - # runs after the mask is built, so derive it here -- every token seen so far is prefill. - prefill_length = static_reference_layer.cumulative_length_int - if prefill_length is not None: - # `window_end` is constant for the whole generation, so the cap does not trigger - # per-step recompiles. Prefill (`prefill_length is None`, more than one query token) - # keeps `reference_window_mask_function=None` and so stays plain causal. - window_end = prefill_length + self.config.sliding_window - - def reference_window_mask_function(batch_idx, head_idx, q_idx, kv_idx): - return kv_idx < window_end - + # (slots `[prefill_length, prefill_length + sliding_window)`). The buffer is sized to exactly + # `prefill_length + sliding_window` (no empty tail), so a plain causal mask is already correct: it + # hides the not-yet-written window slots while the window fills, and there is no tail left to cap. The + # dynamic cache evicts old tokens from its (compact) buffer instead, so plain causal is correct there + # too. causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), - "reference_sliding_attention": create_causal_mask( - **mask_kwargs, and_mask_function=reference_window_mask_function - ), + "reference_sliding_attention": create_causal_mask(**mask_kwargs), } hidden_states = inputs_embeds From e21c1d98beb79cc4712dcfa59e2496efbae699b7 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Mon, 29 Jun 2026 10:02:52 +0200 Subject: [PATCH 062/195] cleanup --- docs/source/en/model_doc/unlimited_ocr.md | 2 +- src/transformers/cache_utils.py | 168 +++++++++--------- .../unlimited_ocr/modeling_unlimited_ocr.py | 7 - .../unlimited_ocr/modular_unlimited_ocr.py | 7 - 4 files changed, 88 insertions(+), 96 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index ecd7239087ad..fe1694b4bcb9 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -16,7 +16,7 @@ limitations under the License. ⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer. --> -*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-06-25.* +*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-06-29.* # UnlimitedOcr diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index a4e24fb48166..defb04dd42b6 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -684,95 +684,58 @@ def reset(self): class StaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): - """Static counterpart of [`DynamicReferenceSlidingWindowLayer`], used when generating with - ``cache_implementation="static"``. - - The backing buffer is split in two regions: slots ``[0, prefill_length)`` hold the image/prompt prefill - (written once, at the front, and never evicted) and slots ``[prefill_length, prefill_length + sliding_window)`` - are a ring for the generated tokens. Exactly like [`StaticSlidingWindowLayer`] the oldest generated token is - rolled out and replaced by the newest once the ring is full, the only difference being that the roll is - restricted to that window so the prefill region stays untouched. - - The physical buffer is sized to ``prefill_length + sliding_window`` rather than the full ``max_cache_len`` - budget that ``generate`` passes (which is ``prefill_length + max_new_tokens``): the trailing - ``max_new_tokens - sliding_window`` slots would never be written, so reserving them is pure waste -- large for - long generations. ``prefill_length`` is only known at runtime, but the buffer is allocated lazily during the - eager prefill pass (and grown there for chunked prefill), which always runs *before* the first compiled/ - cudagraph decode step; from the first decode step onward the shape and data pointer are constant, so cudagraph - capture is unaffected and the SDPA backend (cuDNN included) reuses a single kernel plan. The not-yet-written - window slots are hidden by the (plain causal) decode mask built by the model (e.g. - ``UnlimitedOcrTextModel.forward``); there is no empty tail left to hide. + """ + A static cache layer that stores the key and value states as static tensors of shape + `[batch_size, num_heads, min(max_cache_len, prefill_size + sliding_window), head_dim]`. + It lazily allocates its full backing tensors, and then mutates them in-place. + Built for `torch.compile` support. + + The backing buffer is split in two regions: slots ``[0, prefill_length)`` hold the prefill (reference) slots + that are never eviced from the cache. Slots ``[prefill_length, prefill_length + sliding_window)`` hold the + sliding window decode slots where the oldest entries are always replaced by the newest ones. + + Args: + max_cache_len (`int`): + Maximum number of tokens that can be stored, used for tensor preallocation. + sliding_window (`int`): + The size of the sliding window. """ layer_type = "reference_sliding_attention" def __init__(self, max_cache_len: int, sliding_window: int): super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window) - # `StaticSlidingWindowLayer` shrinks its buffer to `sliding_window` and rolls the whole thing. Here the - # buffer must also hold the pinned prefill. We keep `max_cache_len` only as the *logical* budget (for - # length bookkeeping) and size the *physical* buffer to `prefill_length + sliding_window`, tracking the - # real window separately; the roll below only touches the `sliding_window` slots right after the prefill - # (i.e. `[prefill_length, prefill_length + sliding_window)`). + # Keep `max_cache_len` as max value for lenght bookkeeping. + # The phyiscal buffer doens't exceed `prefill_length + sliding_window`. self.max_cache_len = max_cache_len self.sliding_window = sliding_window self.prefill_length: int | None = None - def _set_buffers(self, physical_length: int, copy_existing: bool = False) -> None: - """(Re)allocate the static key/value buffers to `physical_length` slots and (re)tag the static address. - - Only ever called from the eager prefill pass (initial allocation or chunked-prefill growth), never from a - compiled decode step, so reallocating here is safe for cudagraphs. - """ - new_keys = torch.zeros( - (self.batch_size, self.num_heads, physical_length, self.k_head_dim), dtype=self.dtype, device=self.device - ) - new_values = torch.zeros( - (self.batch_size, self.num_heads, physical_length, self.v_head_dim), dtype=self.dtype, device=self.device - ) - if copy_existing: - old_length = self.keys.shape[-2] - new_keys[:, :, :old_length, :] = self.keys - new_values[:, :, :old_length, :] = self.values - self.keys = new_keys - self.values = new_values - if not is_torchdynamo_compiling(): - torch._dynamo.mark_static_address(self.keys) - torch._dynamo.mark_static_address(self.values) - - def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None: - self.dtype, self.device = key_states.dtype, key_states.device - self.batch_size, self.num_heads = key_states.shape[:2] - self.v_head_dim = value_states.shape[-1] - self.k_head_dim = key_states.shape[-1] - self.cumulative_length = self.cumulative_length.to(self.device) - if not is_torchdynamo_compiling(): - torch._dynamo.mark_static_address(self.cumulative_length) - # Size the physical buffer from the prefill instead of the full `max_cache_len` budget (see class - # docstring). A multi-token first call is prefill; a single-token first call is an empty-prompt decode - # (`prefill_length == 0`), so only the window is needed. - prefill_seen = key_states.shape[-2] if key_states.shape[-2] > 1 else 0 - self._set_buffers(min(self.max_cache_len, prefill_seen + self.sliding_window)) - self.is_initialized = True - def update( self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Update the key and value caches in-place, and return the necessary keys and value states. + + Args: + key_states (`torch.Tensor`): The new key states to cache. + value_states (`torch.Tensor`): The new value states to cache. + + Returns: + tuple[`torch.Tensor`, `torch.Tensor`]: The key and value states. + """ # Lazy initialization if not self.is_initialized: self.lazy_initialization(key_states, value_states) kv_length = key_states.shape[-2] - # Prefill. Mirror `DynamicReferenceSlidingWindowLayer`: every multi-token chunk fed before the first - # single-token (decode) step is prefill and stays pinned at the front of the buffer. Crucially we do - # *not* pin `prefill_length` here, so chunked prefill (a prompt fed in more than one chunk) keeps - # extending the pinned region instead of having the later chunks slid out of the window. + # Prefill if self.prefill_length is None and kv_length > 1: - # Chunked prefill: the buffer was sized for the first chunk only, so grow it (eagerly, before any - # compiled decode step) to hold all prefill seen so far plus this chunk plus the reserved window. + # Chunked prefill: the buffer was sized for the first chunk only, so grow it to hold all prefill seen so far plus this chunk plus the reserved window. required_length = min(self.max_cache_len, self.cumulative_length_int + kv_length + self.sliding_window) if self.keys.shape[-2] < required_length: - self._set_buffers(required_length, copy_existing=True) + self._allocate_key_value_buffers(required_length, copy_existing=True) # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length @@ -791,8 +754,7 @@ def update( # Very important to return the `self` tensors here, as they have the static dynamo address return self.keys, self.values - # First single-token step (or an empty prompt that skipped prefill above): finalize the pinned prefill - # region so the trailing window below is measured relative to it. + # First decode step (or an empty prompt that skipped prefill above): mark prefill as complete if self.prefill_length is None: self.prefill_length = self.cumulative_length_int @@ -823,7 +785,7 @@ def update( # Very important to return the `self` tensors here, as they have the static dynamo address return self.keys, self.values - # Already full but using more than 1 new token (e.g. prefill caching, chat continuation, etc...) + # Already full but using more than 1 new token (e.g. chat continuation, etc...) else: full_key_states = torch.cat( (self.keys[:, :, window_start + 1 : window_start + self.sliding_window, :], key_states), dim=-2 @@ -844,8 +806,6 @@ def update( else: # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten. - # The window sits right after the prefill (`window_start == prefill_length`), so the decode tokens are - # written at their logical positions, which already start at `prefill_length`. cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length try: self.keys.index_copy_(2, cache_position, key_states) @@ -862,33 +822,79 @@ def update( # Very important to return the `self` tensors here, as they have the static dynamo address return self.keys, self.values - # We only keep the last `sliding_window` tokens. Unlike `StaticSlidingWindowLayer` we return `self.keys/values` - # because `get_mask_sizes` is fixed to the full `max_cache_len`. + # We only cache the last `sliding_window` tokens self.keys[:, :, window_start : window_start + self.sliding_window, :].copy_( full_key_states[:, :, -self.sliding_window :, :] ) self.values[:, :, window_start : window_start + self.sliding_window, :].copy_( full_value_states[:, :, -self.sliding_window :, :] ) + # TODO: Multi-token decode is not supported yet as it would require a custom + # create_causal_mask implementation to create correct masks. return self.keys, self.values + def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None: + self.dtype, self.device = key_states.dtype, key_states.device + self.batch_size, self.num_heads = key_states.shape[:2] + self.v_head_dim = value_states.shape[-1] + self.k_head_dim = key_states.shape[-1] + + self.cumulative_length = self.cumulative_length.to(self.device) + # Note: `mark_static_address` is used to tag the tensors as a fixed data pointer, preventing compiled graph + # breaks or cudagraph skips due to inplace mutations when updating the cache. However, it is not supported when + # tracing the graph, so we skip it in this case. As prefill should never be compiled, this is not an issue and it + # will still be run (except when users compile prefill explicitly, but this should be avoided!) + # Without this, we cannot use cudagraphs!! + if not is_torchdynamo_compiling(): + torch._dynamo.mark_static_address(self.cumulative_length) + + prefill_seen = key_states.shape[-2] if key_states.shape[-2] > 1 else 0 + self._allocate_key_value_buffers(min(self.max_cache_len, prefill_seen + self.sliding_window)) + + self.is_initialized = True + def get_mask_sizes(self, query_length: int) -> tuple[int, int]: - # `kv_length` must match the length of the tensor `update` will return for this step. The buffer is sized - # to `prefill_length + sliding_window`, so anticipate the size `update` will (re)allocate it to: + """Return the length and offset of the cache, used to generate the attention mask""" + kv_offset = 0 if not self.is_initialized: # First forward: `lazy_initialization` will allocate from this query length. prefill_seen = query_length if query_length > 1 else 0 - return min(self.max_cache_len, prefill_seen + self.sliding_window), 0 - if self.prefill_length is None and query_length > 1: + kv_length = min(self.max_cache_len, prefill_seen + self.sliding_window) + elif self.prefill_length is None and query_length > 1: # Additional prefill chunk: `update` will grow the buffer to fit all prefill so far plus the window. - return min(self.max_cache_len, self.cumulative_length_int + query_length + self.sliding_window), 0 - # Decode (including the first decode step): the buffer is already at its final size. - return self.keys.shape[-2], 0 + kv_length = min(self.max_cache_len, self.cumulative_length_int + query_length + self.sliding_window) + else: + # Decode (including the first decode step): the buffer is already at its final size. + kv_length = self.keys.shape[-2] + + return kv_length, kv_offset def reset(self) -> None: super().reset() self.prefill_length = None + def _allocate_key_value_buffers(self, physical_length: int, copy_existing: bool = False) -> None: + """(Re)allocate the static key/value buffers to `physical_length` slots and (re)tag the static address. + + Only ever called from the eager prefill pass (initial allocation or chunked-prefill growth), never from a + compiled decode step, so reallocating here is safe for cudagraphs. + """ + new_keys = torch.zeros( + (self.batch_size, self.num_heads, physical_length, self.k_head_dim), dtype=self.dtype, device=self.device + ) + new_values = torch.zeros( + (self.batch_size, self.num_heads, physical_length, self.v_head_dim), dtype=self.dtype, device=self.device + ) + if copy_existing: + old_length = self.keys.shape[-2] + new_keys[:, :, :old_length, :].copy_(self.keys) + new_values[:, :, :old_length, :].copy_(self.values) + self.keys = new_keys + self.values = new_values + if not is_torchdynamo_compiling(): + torch._dynamo.mark_static_address(self.keys) + torch._dynamo.mark_static_address(self.values) + class StaticIndexedLayer(StaticLayer): """ diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index afd334acb5fb..fb98c56f1d9d 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1388,13 +1388,6 @@ def forward( "past_key_values": past_key_values, "position_ids": position_ids, } - # The static cache ([`StaticReferenceSlidingWindowLayer`]) keeps the prefill at the front of its - # buffer (slots `[0, prefill_length)`) followed by the most recent `sliding_window` generated tokens - # (slots `[prefill_length, prefill_length + sliding_window)`). The buffer is sized to exactly - # `prefill_length + sliding_window` (no empty tail), so a plain causal mask is already correct: it - # hides the not-yet-written window slots while the window fills, and there is no tail left to cap. The - # dynamic cache evicts old tokens from its (compact) buffer instead, so plain causal is correct there - # too. causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), "reference_sliding_attention": create_causal_mask(**mask_kwargs), diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 6b8d081ee956..a8977817d545 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -644,13 +644,6 @@ def forward( "past_key_values": past_key_values, "position_ids": position_ids, } - # The static cache ([`StaticReferenceSlidingWindowLayer`]) keeps the prefill at the front of its - # buffer (slots `[0, prefill_length)`) followed by the most recent `sliding_window` generated tokens - # (slots `[prefill_length, prefill_length + sliding_window)`). The buffer is sized to exactly - # `prefill_length + sliding_window` (no empty tail), so a plain causal mask is already correct: it - # hides the not-yet-written window slots while the window fills, and there is no tail left to cap. The - # dynamic cache evicts old tokens from its (compact) buffer instead, so plain causal is correct there - # too. causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), "reference_sliding_attention": create_causal_mask(**mask_kwargs), From 6dc6d88d4ad0ee2a36c242f19baf0d1747c63cdb Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Mon, 29 Jun 2026 11:21:53 +0200 Subject: [PATCH 063/195] cleanup dynamic --- src/transformers/cache_utils.py | 79 ++++++++++++--------------------- 1 file changed, 29 insertions(+), 50 deletions(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index defb04dd42b6..0f624cf65f1a 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -304,7 +304,6 @@ class DynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): def __init__(self, config: PreTrainedConfig | None = None, sliding_window: int | None = None): super().__init__(config=config, sliding_window=sliding_window) self.prefill_length: int | None = None - self.ring_position = 0 def update( self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs @@ -313,65 +312,46 @@ def update( if not self.is_initialized: self.lazy_initialization(key_states, value_states) - sequence_length = key_states.shape[-2] - self.cumulative_length += sequence_length + kv_length = key_states.shape[-2] + self.cumulative_length += kv_length - # Prefill with prompt context - if self.prefill_length is None and sequence_length > 1: - self.keys = torch.cat([self.keys, key_states], dim=-2) - self.values = torch.cat([self.values, value_states], dim=-2) + # Compute the full states + full_key_states = torch.cat([self.keys, key_states], dim=-2) + full_value_states = torch.cat([self.values, value_states], dim=-2) + + # Prefill + if self.prefill_length is None and kv_length > 1: + self.keys = full_key_states + self.values = full_value_states return self.keys, self.values - # First decode step - # Handle generation with empty prompt + # First decode step (or an empty prompt that skipped prefill above): mark prefill as complete if self.prefill_length is None: self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 - # Append while window grows - generated_length = self.keys.shape[-2] - self.prefill_length if self.keys.dim() > 1 else 0 - append_length = min(sequence_length, max(0, self.sliding_window - generated_length)) - if append_length > 0: - self.keys = torch.cat([self.keys, key_states[..., :append_length, :]], dim=-2) - self.values = torch.cat([self.values, value_states[..., :append_length, :]], dim=-2) - - # Overwrite if window size is reached - overwrite_length = sequence_length - append_length - if overwrite_length > 0: - # Only the most recent `sliding_window` overwrites survive - write_length = min(overwrite_length, self.sliding_window) - start = self.ring_position + overwrite_length - write_length - offsets = torch.arange(write_length, device=key_states.device) - slots = self.prefill_length + (start + offsets) % self.sliding_window - self.keys[..., slots, :] = key_states[..., sequence_length - write_length :, :] - self.values[..., slots, :] = value_states[..., sequence_length - write_length :, :] - self.ring_position = (self.ring_position + overwrite_length) % self.sliding_window + generated_length = full_key_states.shape[-2] - self.prefill_length + if generated_length <= self.sliding_window: + # Append while window grows + self.keys = full_key_states + self.values = full_value_states + else: + # Overwrite + self.keys[:, :, -self.sliding_window :, :].copy_(full_key_states[:, :, -self.sliding_window :, :]) + self.values[:, :, -self.sliding_window :, :].copy_(full_value_states[:, :, -self.sliding_window :, :]) + # TODO: Add multi-token decode support. This requires a custom create_causal_mask implementation. return self.keys, self.values def get_mask_sizes(self, query_length: int) -> tuple[int, int]: - # Full visibility, no sliding offset: every query attends to all cached keys (all prefill plus the - # generated tokens currently held in the ring), so the prefill is never masked out. + """Return the length and offset of the cache, used to generate the attention mask""" + kv_offset = 0 if self.prefill_length is None: - return self.cumulative_length + query_length, 0 - return self.decode_kv_length(query_length), 0 - - def decode_kv_length(self, query_length: int = 1) -> int | None: - """Physical length of the cached buffer after the upcoming `update` of `query_length` tokens. - - Returns `None` before the layer is initialized. Used to size the all-visible decode mask, whose - width must match the key/value tensors returned by `update`. - """ - if not self.is_initialized: - return None - if self.prefill_length is not None: - prefill_length = self.prefill_length - generated_before = self.keys.shape[-2] - prefill_length if self.keys.dim() > 1 else 0 + # Before the first decode step the whole buffer is prefill + kv_length = self.cumulative_length + query_length else: - # Before the first decode step the whole buffer is prefill. - prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 - generated_before = 0 - generated_after = min(generated_before + query_length, self.sliding_window) - return prefill_length + generated_after + generated_length = self.keys.shape[-2] - self.prefill_length + kv_length = self.prefill_length + min(generated_length + query_length, self.sliding_window) + return kv_length, kv_offset class DynamicIndexedLayer(DynamicLayer): @@ -829,8 +809,7 @@ def update( self.values[:, :, window_start : window_start + self.sliding_window, :].copy_( full_value_states[:, :, -self.sliding_window :, :] ) - # TODO: Multi-token decode is not supported yet as it would require a custom - # create_causal_mask implementation to create correct masks. + # TODO: Add multi-token decode support. This requires a custom create_causal_mask implementation. return self.keys, self.values def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None: From 0d85b8915dc766e969d4c2376d062c00d1390849 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 1 Jul 2026 10:47:49 +0000 Subject: [PATCH 064/195] update --- docs/source/en/model_doc/unlimited_ocr.md | 2 +- src/transformers/cache_utils.py | 64 +++++++++++++++---- .../unlimited_ocr/modeling_unlimited_ocr.py | 18 +++++- .../unlimited_ocr/modular_unlimited_ocr.py | 18 +++++- 4 files changed, 87 insertions(+), 15 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index fe1694b4bcb9..04c9d008bbcf 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -16,7 +16,7 @@ limitations under the License. ⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer. --> -*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-06-29.* +*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-07-01.* # UnlimitedOcr diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index cd5a65cde0fc..af892ab8f16d 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -330,29 +330,69 @@ def update( self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 generated_length = full_key_states.shape[-2] - self.prefill_length - if generated_length <= self.sliding_window: - # Append while window grows + if generated_length < self.sliding_window: + # Window still growing: keep every prefill and decode token. self.keys = full_key_states self.values = full_value_states else: - # Overwrite - self.keys[:, :, -self.sliding_window :, :].copy_(full_key_states[:, :, -self.sliding_window :, :]) - self.values[:, :, -self.sliding_window :, :].copy_(full_value_states[:, :, -self.sliding_window :, :]) + # TODO: copy when already full + # Window full: keep all prefill tokens plus the most recent `sliding_window - 1` decode tokens. + # The `- 1` mirrors `DynamicSlidingWindowLayer`: the incoming query token supplies the last slot, + # so the current query still attends over a full `sliding_window` of decode tokens. + self.keys = torch.cat( + [ + full_key_states[:, :, : self.prefill_length, :], + full_key_states[:, :, -self.sliding_window + 1 :, :], + ], + dim=-2, + ) + self.values = torch.cat( + [ + full_value_states[:, :, : self.prefill_length, :], + full_value_states[:, :, -self.sliding_window + 1 :, :], + ], + dim=-2, + ) - # TODO: Add multi-token decode support. This requires a custom create_causal_mask implementation. - return self.keys, self.values + # Return full states to avoid losing context in case we added more than sliding_window tokens at once + return full_key_states, full_value_states def get_mask_sizes(self, query_length: int) -> tuple[int, int]: """Return the length and offset of the cache, used to generate the attention mask""" + is_full = ( + self.prefill_length is not None and self.cumulative_length >= self.prefill_length + self.sliding_window + ) + kv_offset = 0 - if self.prefill_length is None: - # Before the first decode step the whole buffer is prefill - kv_length = self.cumulative_length + query_length + if is_full: + kv_offset = max(self.cumulative_length - self.prefill_length - self.sliding_window + 1, 0) + kv_length = self.prefill_length + self.sliding_window - 1 + query_length else: - generated_length = self.keys.shape[-2] - self.prefill_length - kv_length = self.prefill_length + min(generated_length + query_length, self.sliding_window) + kv_length = self.cumulative_length + query_length + + #Β Returned kv_offset is with respect to sliding window keys. + #Β Remove kv_offset from kv_idx to retrieve the prefill indices. return kv_length, kv_offset + def get_max_length(self) -> int: + """Return the maximum cache shape of the cache""" + if self.prefill_length is None: + return -1 + return self.prefill_length + self.sliding_window + + def crop(self, max_length: int) -> None: + """ + Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be + negative to remove `max_length` tokens. + """ + if self.prefill_length is not None and self.get_seq_length() >= self.prefill_length + self.sliding_window: + raise ValueError( + "Cannot `crop` a `DynamicReferenceSlidingWindowLayer` after it has seen more tokens than its" + "prefill + sliding window (otherwise some states are lost)" + ) + DynamicLayer.crop(self, max_length) + self.cumulative_length = self.keys.shape[-2] + class DynamicIndexedLayer(DynamicLayer): """ diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index fb98c56f1d9d..09bdf25e81ae 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1388,9 +1388,25 @@ def forward( "past_key_values": past_key_values, "position_ids": position_ids, } + + from ...masking_utils import and_masks, causal_mask_function, create_sliding_window_causal_mask + + def create_reference_sliding_window_causal_mask(**kwargs): + cache_layer = past_key_values.layers[0] + prefill_length = float("inf") if cache_layer.prefill_length is None else cache_layer.prefill_length + _, kv_offset = cache_layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) + + def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): + # Remove kv_offset to retrieve the kv_index with respect to prefill + return kv_idx - kv_offset < prefill_length + + return create_sliding_window_causal_mask( + or_mask_function=and_masks(causal_mask_function, prefill_overlay), **kwargs + ) + causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), - "reference_sliding_attention": create_causal_mask(**mask_kwargs), + "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), } hidden_states = inputs_embeds diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index a8977817d545..268581b4ea5d 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -644,9 +644,25 @@ def forward( "past_key_values": past_key_values, "position_ids": position_ids, } + + from ...masking_utils import and_masks, causal_mask_function, create_sliding_window_causal_mask + + def create_reference_sliding_window_causal_mask(**kwargs): + cache_layer = past_key_values.layers[0] + prefill_length = float("inf") if cache_layer.prefill_length is None else cache_layer.prefill_length + _, kv_offset = cache_layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) + + def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): + # Remove kv_offset to retrieve the kv_index with respect to prefill + return kv_idx - kv_offset < prefill_length + + return create_sliding_window_causal_mask( + or_mask_function=and_masks(causal_mask_function, prefill_overlay), **kwargs + ) + causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), - "reference_sliding_attention": create_causal_mask(**mask_kwargs), + "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), } hidden_states = inputs_embeds From 3fddfb568530afbcab5ea38a2a6154d2feccfcc5 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 1 Jul 2026 12:17:39 +0000 Subject: [PATCH 065/195] update static reference sliding window cache --- src/transformers/cache_utils.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index af892ab8f16d..39709748de17 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -849,8 +849,7 @@ def update( self.values[:, :, window_start : window_start + self.sliding_window, :].copy_( full_value_states[:, :, -self.sliding_window :, :] ) - # TODO: Add multi-token decode support. This requires a custom create_causal_mask implementation. - return self.keys, self.values + return full_key_states, full_value_states def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None: self.dtype, self.device = key_states.dtype, key_states.device @@ -874,16 +873,21 @@ def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tens def get_mask_sizes(self, query_length: int) -> tuple[int, int]: """Return the length and offset of the cache, used to generate the attention mask""" + is_full = self.prefill_length is not None and self.cumulative_length_int >= self.prefill_length + self.sliding_window + kv_offset = 0 - if not self.is_initialized: - # First forward: `lazy_initialization` will allocate from this query length. - prefill_seen = query_length if query_length > 1 else 0 - kv_length = min(self.max_cache_len, prefill_seen + self.sliding_window) - elif self.prefill_length is None and query_length > 1: - # Additional prefill chunk: `update` will grow the buffer to fit all prefill so far plus the window. + # Prefill + if self.prefill_length is None and query_length > 1: kv_length = min(self.max_cache_len, self.cumulative_length_int + query_length + self.sliding_window) + # Decode: cache is already full + elif is_full: + kv_offset = max(self.cumulative_length_int - self.prefill_length - self.sliding_window + 1, 0) + kv_length = self.prefill_length + self.sliding_window - 1 + query_length + # Decode: cache not yet full, but becoming full on this update + elif self.prefill_length is not None and self.cumulative_length_int + query_length > self.prefill_length + self.sliding_window: + kv_length = self.cumulative_length_int + query_length + # Decode: cache not yet full but we return the local size as it's static else: - # Decode (including the first decode step): the buffer is already at its final size. kv_length = self.keys.shape[-2] return kv_length, kv_offset From c8e1e2bcc2dc55ce63c3be625b80cc37d94f1543 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 1 Jul 2026 18:07:21 +0100 Subject: [PATCH 066/195] move create mask to top level --- src/transformers/cache_utils.py | 13 ++++--- .../unlimited_ocr/modeling_unlimited_ocr.py | 34 ++++++++++--------- .../unlimited_ocr/modular_unlimited_ocr.py | 34 ++++++++++--------- 3 files changed, 45 insertions(+), 36 deletions(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index 39709748de17..bbe1bc89c11e 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -370,8 +370,8 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: else: kv_length = self.cumulative_length + query_length - #Β Returned kv_offset is with respect to sliding window keys. - #Β Remove kv_offset from kv_idx to retrieve the prefill indices. + # Returned kv_offset is with respect to sliding window keys. + # Remove kv_offset from kv_idx to retrieve the prefill indices. return kv_length, kv_offset def get_max_length(self) -> int: @@ -873,7 +873,9 @@ def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tens def get_mask_sizes(self, query_length: int) -> tuple[int, int]: """Return the length and offset of the cache, used to generate the attention mask""" - is_full = self.prefill_length is not None and self.cumulative_length_int >= self.prefill_length + self.sliding_window + is_full = ( + self.prefill_length is not None and self.cumulative_length_int >= self.prefill_length + self.sliding_window + ) kv_offset = 0 # Prefill @@ -884,7 +886,10 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: kv_offset = max(self.cumulative_length_int - self.prefill_length - self.sliding_window + 1, 0) kv_length = self.prefill_length + self.sliding_window - 1 + query_length # Decode: cache not yet full, but becoming full on this update - elif self.prefill_length is not None and self.cumulative_length_int + query_length > self.prefill_length + self.sliding_window: + elif ( + self.prefill_length is not None + and self.cumulative_length_int + query_length > self.prefill_length + self.sliding_window + ): kv_length = self.cumulative_length_int + query_length # Decode: cache not yet full but we return the local size as it's static else: diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 09bdf25e81ae..c37b5c206f45 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -37,7 +37,7 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...masking_utils import create_causal_mask +from ...masking_utils import and_masks, causal_mask_function, create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update @@ -1333,6 +1333,23 @@ def forward(self, x, position_ids): return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) +def create_reference_sliding_window_causal_mask(**kwargs): + past_key_values = kwargs["past_key_values"] + inputs_embeds = kwargs["inputs_embeds"] + + layer = next(layer for layer in past_key_values.layers if layer.layer_type == "reference_sliding_attention") + prefill_length = float("inf") if layer.prefill_length is None else layer.prefill_length + _, kv_offset = layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) + + def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): + # Remove kv_offset to retrieve the kv_index with respect to prefill + return kv_idx - kv_offset < prefill_length + + return create_sliding_window_causal_mask( + or_mask_function=and_masks(causal_mask_function, prefill_overlay), **kwargs + ) + + @auto_docstring class UnlimitedOcrTextModel(UnlimitedOcrTextPreTrainedModel): def __init__(self, config: UnlimitedOcrTextConfig): @@ -1389,21 +1406,6 @@ def forward( "position_ids": position_ids, } - from ...masking_utils import and_masks, causal_mask_function, create_sliding_window_causal_mask - - def create_reference_sliding_window_causal_mask(**kwargs): - cache_layer = past_key_values.layers[0] - prefill_length = float("inf") if cache_layer.prefill_length is None else cache_layer.prefill_length - _, kv_offset = cache_layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) - - def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): - # Remove kv_offset to retrieve the kv_index with respect to prefill - return kv_idx - kv_offset < prefill_length - - return create_sliding_window_causal_mask( - or_mask_function=and_masks(causal_mask_function, prefill_overlay), **kwargs - ) - causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 268581b4ea5d..a5d6d4561693 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -24,7 +24,7 @@ from ...feature_extraction_utils import BatchFeature from ...image_transforms import group_images_by_shape, reorder_images from ...image_utils import ImageInput, PILImageResampling, SizeDict -from ...masking_utils import create_causal_mask +from ...masking_utils import and_masks, causal_mask_function, create_causal_mask, create_sliding_window_causal_mask from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast from ...processing_utils import Unpack from ...tokenization_utils_base import PreTokenizedInput, TextInput @@ -607,6 +607,23 @@ class UnlimitedOcrTextPreTrainedModel(DeepseekOcr2TextPreTrainedModel): pass +def create_reference_sliding_window_causal_mask(**kwargs): + past_key_values = kwargs["past_key_values"] + inputs_embeds = kwargs["inputs_embeds"] + + layer = next(layer for layer in past_key_values.layers if layer.layer_type == "reference_sliding_attention") + prefill_length = float("inf") if layer.prefill_length is None else layer.prefill_length + _, kv_offset = layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) + + def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): + # Remove kv_offset to retrieve the kv_index with respect to prefill + return kv_idx - kv_offset < prefill_length + + return create_sliding_window_causal_mask( + or_mask_function=and_masks(causal_mask_function, prefill_overlay), **kwargs + ) + + class UnlimitedOcrTextModel(DeepseekOcr2TextModel): @merge_with_config_defaults @capture_outputs @@ -645,21 +662,6 @@ def forward( "position_ids": position_ids, } - from ...masking_utils import and_masks, causal_mask_function, create_sliding_window_causal_mask - - def create_reference_sliding_window_causal_mask(**kwargs): - cache_layer = past_key_values.layers[0] - prefill_length = float("inf") if cache_layer.prefill_length is None else cache_layer.prefill_length - _, kv_offset = cache_layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) - - def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): - # Remove kv_offset to retrieve the kv_index with respect to prefill - return kv_idx - kv_offset < prefill_length - - return create_sliding_window_causal_mask( - or_mask_function=and_masks(causal_mask_function, prefill_overlay), **kwargs - ) - causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), From 85b2e1b3ba65023821a8a6b8f2cfcc93a58fc154 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 1 Jul 2026 18:24:37 +0100 Subject: [PATCH 067/195] move cache to model files --- src/transformers/cache_utils.py | 324 ----------------- .../unlimited_ocr/modeling_unlimited_ocr.py | 328 ++++++++++++++++- .../unlimited_ocr/modular_unlimited_ocr.py | 335 +++++++++++++++++- 3 files changed, 660 insertions(+), 327 deletions(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index bbe1bc89c11e..adc95cd939bd 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -291,109 +291,6 @@ def crop(self, max_length: int) -> None: self.cumulative_length = self.keys.shape[-2] -class DynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): - """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens before - the first decode step and applies a sliding window to all decoded tokens. - - Once ``sliding_window`` decode tokens have accumulated, the oldest decode tokens are evicted and - replaced by the most recent ones. The prefill tokens always remain in the cache. - """ - - layer_type = "reference_sliding_attention" - - def __init__(self, config: PreTrainedConfig | None = None, sliding_window: int | None = None): - super().__init__(config=config, sliding_window=sliding_window) - self.prefill_length: int | None = None - - def update( - self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs - ) -> tuple[torch.Tensor, torch.Tensor]: - # Lazy initialization - if not self.is_initialized: - self.lazy_initialization(key_states, value_states) - - kv_length = key_states.shape[-2] - self.cumulative_length += kv_length - - # Compute the full states - full_key_states = torch.cat([self.keys, key_states], dim=-2) - full_value_states = torch.cat([self.values, value_states], dim=-2) - - # Prefill - if self.prefill_length is None and kv_length > 1: - self.keys = full_key_states - self.values = full_value_states - return self.keys, self.values - - # First decode step (or an empty prompt that skipped prefill above): mark prefill as complete - if self.prefill_length is None: - self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 - - generated_length = full_key_states.shape[-2] - self.prefill_length - if generated_length < self.sliding_window: - # Window still growing: keep every prefill and decode token. - self.keys = full_key_states - self.values = full_value_states - else: - # TODO: copy when already full - # Window full: keep all prefill tokens plus the most recent `sliding_window - 1` decode tokens. - # The `- 1` mirrors `DynamicSlidingWindowLayer`: the incoming query token supplies the last slot, - # so the current query still attends over a full `sliding_window` of decode tokens. - self.keys = torch.cat( - [ - full_key_states[:, :, : self.prefill_length, :], - full_key_states[:, :, -self.sliding_window + 1 :, :], - ], - dim=-2, - ) - self.values = torch.cat( - [ - full_value_states[:, :, : self.prefill_length, :], - full_value_states[:, :, -self.sliding_window + 1 :, :], - ], - dim=-2, - ) - - # Return full states to avoid losing context in case we added more than sliding_window tokens at once - return full_key_states, full_value_states - - def get_mask_sizes(self, query_length: int) -> tuple[int, int]: - """Return the length and offset of the cache, used to generate the attention mask""" - is_full = ( - self.prefill_length is not None and self.cumulative_length >= self.prefill_length + self.sliding_window - ) - - kv_offset = 0 - if is_full: - kv_offset = max(self.cumulative_length - self.prefill_length - self.sliding_window + 1, 0) - kv_length = self.prefill_length + self.sliding_window - 1 + query_length - else: - kv_length = self.cumulative_length + query_length - - # Returned kv_offset is with respect to sliding window keys. - # Remove kv_offset from kv_idx to retrieve the prefill indices. - return kv_length, kv_offset - - def get_max_length(self) -> int: - """Return the maximum cache shape of the cache""" - if self.prefill_length is None: - return -1 - return self.prefill_length + self.sliding_window - - def crop(self, max_length: int) -> None: - """ - Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be - negative to remove `max_length` tokens. - """ - if self.prefill_length is not None and self.get_seq_length() >= self.prefill_length + self.sliding_window: - raise ValueError( - "Cannot `crop` a `DynamicReferenceSlidingWindowLayer` after it has seen more tokens than its" - "prefill + sliding window (otherwise some states are lost)" - ) - DynamicLayer.crop(self, max_length) - self.cumulative_length = self.keys.shape[-2] - - class DynamicIndexedLayer(DynamicLayer): """ A cache layer that extends `DynamicLayer` with an extra indexer key cache for Dynamic Sparse Attention (DSA) @@ -703,227 +600,6 @@ def reset(self): self.cumulative_length_int = 0 -class StaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): - """ - A static cache layer that stores the key and value states as static tensors of shape - `[batch_size, num_heads, min(max_cache_len, prefill_size + sliding_window), head_dim]`. - It lazily allocates its full backing tensors, and then mutates them in-place. - Built for `torch.compile` support. - - The backing buffer is split in two regions: slots ``[0, prefill_length)`` hold the prefill (reference) slots - that are never eviced from the cache. Slots ``[prefill_length, prefill_length + sliding_window)`` hold the - sliding window decode slots where the oldest entries are always replaced by the newest ones. - - Args: - max_cache_len (`int`): - Maximum number of tokens that can be stored, used for tensor preallocation. - sliding_window (`int`): - The size of the sliding window. - """ - - layer_type = "reference_sliding_attention" - - def __init__(self, max_cache_len: int, sliding_window: int): - super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window) - # Keep `max_cache_len` as max value for lenght bookkeeping. - # The phyiscal buffer doens't exceed `prefill_length + sliding_window`. - self.max_cache_len = max_cache_len - self.sliding_window = sliding_window - self.prefill_length: int | None = None - - def update( - self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs - ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Update the key and value caches in-place, and return the necessary keys and value states. - - Args: - key_states (`torch.Tensor`): The new key states to cache. - value_states (`torch.Tensor`): The new value states to cache. - - Returns: - tuple[`torch.Tensor`, `torch.Tensor`]: The key and value states. - """ - # Lazy initialization - if not self.is_initialized: - self.lazy_initialization(key_states, value_states) - - kv_length = key_states.shape[-2] - - # Prefill - if self.prefill_length is None and kv_length > 1: - # Chunked prefill: the buffer was sized for the first chunk only, so grow it to hold all prefill seen so far plus this chunk plus the reserved window. - required_length = min(self.max_cache_len, self.cumulative_length_int + kv_length + self.sliding_window) - if self.keys.shape[-2] < required_length: - self._allocate_key_value_buffers(required_length, copy_existing=True) - # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs - # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten - cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length - try: - self.keys.index_copy_(2, cache_position, key_states) - self.values.index_copy_(2, cache_position, value_states) - except NotImplementedError: - # Fallback for devices like MPS where index_copy_ might not be supported. - self.keys[:, :, cache_position] = key_states - self.values[:, :, cache_position] = value_states - - # Keep both the int (control flow) and the tensor (cudagraph-safe indexing) versions in sync. - self.cumulative_length_int += kv_length - self.cumulative_length.add_(kv_length) - - # Very important to return the `self` tensors here, as they have the static dynamo address - return self.keys, self.values - - # First decode step (or an empty prompt that skipped prefill above): mark prefill as complete - if self.prefill_length is None: - self.prefill_length = self.cumulative_length_int - - # Everything below mirrors `StaticSlidingWindowLayer.update` with the difference that the sliding window - # is in the `sliding_window` slots right after the prefill instead of spanning the whole buffer. - window_start = self.prefill_length - current_length = self.cumulative_length_int - self.prefill_length - is_full = current_length >= self.sliding_window - # Update it now that we saved the value above - self.cumulative_length_int += kv_length - - if is_full: - # In general, we should use a much simpler `cat` here as well, independently of the states size. However, - # dynamo is currently bugged when doing it - see https://github.com/pytorch/pytorch/issues/159855 for more details - if key_states.shape[-2] == 1: - # Roll the window region to the left by 1 position (the pinned prefill in front stays put) - new_keys = self.keys[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) - new_values = self.values[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) - # Overwrite the last position with new states - # (note: very important to use a tensor to index here, see https://github.com/pytorch/pytorch/issues/159855) - index = torch.tensor([-1], dtype=int, device=self.device) - new_keys[:, :, index] = key_states - new_values[:, :, index] = value_states - - # Copy back into `self` (do not just assign again) in order to keep the static dynamo address - self.keys[:, :, window_start : window_start + self.sliding_window, :].copy_(new_keys) - self.values[:, :, window_start : window_start + self.sliding_window, :].copy_(new_values) - - # Very important to return the `self` tensors here, as they have the static dynamo address - return self.keys, self.values - # Already full but using more than 1 new token (e.g. chat continuation, etc...) - else: - full_key_states = torch.cat( - (self.keys[:, :, window_start + 1 : window_start + self.sliding_window, :], key_states), dim=-2 - ) - full_value_states = torch.cat( - (self.values[:, :, window_start + 1 : window_start + self.sliding_window, :], value_states), dim=-2 - ) - # Not yet full, but becoming full on this update - elif current_length + kv_length > self.sliding_window: - # Fast path, no need to cat() in this case, as the window is currently empty - if current_length == 0: - full_key_states = key_states - full_value_states = value_states - else: - window = slice(window_start, window_start + current_length) - full_key_states = torch.cat((self.keys[:, :, window, :], key_states), dim=-2) - full_value_states = torch.cat((self.values[:, :, window, :], value_states), dim=-2) - else: - # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs - # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten. - cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length - try: - self.keys.index_copy_(2, cache_position, key_states) - self.values.index_copy_(2, cache_position, value_states) - except NotImplementedError: - # Fallback for devices like MPS where index_copy_ might not be supported. - self.keys[:, :, cache_position] = key_states - self.values[:, :, cache_position] = value_states - - # Update the tensor version of the length in-place (we don't need to update it if we are already outside - # of this branch, as we don't need the tensor anymore) - self.cumulative_length.add_(kv_length) - - # Very important to return the `self` tensors here, as they have the static dynamo address - return self.keys, self.values - - # We only cache the last `sliding_window` tokens - self.keys[:, :, window_start : window_start + self.sliding_window, :].copy_( - full_key_states[:, :, -self.sliding_window :, :] - ) - self.values[:, :, window_start : window_start + self.sliding_window, :].copy_( - full_value_states[:, :, -self.sliding_window :, :] - ) - return full_key_states, full_value_states - - def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None: - self.dtype, self.device = key_states.dtype, key_states.device - self.batch_size, self.num_heads = key_states.shape[:2] - self.v_head_dim = value_states.shape[-1] - self.k_head_dim = key_states.shape[-1] - - self.cumulative_length = self.cumulative_length.to(self.device) - # Note: `mark_static_address` is used to tag the tensors as a fixed data pointer, preventing compiled graph - # breaks or cudagraph skips due to inplace mutations when updating the cache. However, it is not supported when - # tracing the graph, so we skip it in this case. As prefill should never be compiled, this is not an issue and it - # will still be run (except when users compile prefill explicitly, but this should be avoided!) - # Without this, we cannot use cudagraphs!! - if not is_torchdynamo_compiling(): - torch._dynamo.mark_static_address(self.cumulative_length) - - prefill_seen = key_states.shape[-2] if key_states.shape[-2] > 1 else 0 - self._allocate_key_value_buffers(min(self.max_cache_len, prefill_seen + self.sliding_window)) - - self.is_initialized = True - - def get_mask_sizes(self, query_length: int) -> tuple[int, int]: - """Return the length and offset of the cache, used to generate the attention mask""" - is_full = ( - self.prefill_length is not None and self.cumulative_length_int >= self.prefill_length + self.sliding_window - ) - - kv_offset = 0 - # Prefill - if self.prefill_length is None and query_length > 1: - kv_length = min(self.max_cache_len, self.cumulative_length_int + query_length + self.sliding_window) - # Decode: cache is already full - elif is_full: - kv_offset = max(self.cumulative_length_int - self.prefill_length - self.sliding_window + 1, 0) - kv_length = self.prefill_length + self.sliding_window - 1 + query_length - # Decode: cache not yet full, but becoming full on this update - elif ( - self.prefill_length is not None - and self.cumulative_length_int + query_length > self.prefill_length + self.sliding_window - ): - kv_length = self.cumulative_length_int + query_length - # Decode: cache not yet full but we return the local size as it's static - else: - kv_length = self.keys.shape[-2] - - return kv_length, kv_offset - - def reset(self) -> None: - super().reset() - self.prefill_length = None - - def _allocate_key_value_buffers(self, physical_length: int, copy_existing: bool = False) -> None: - """(Re)allocate the static key/value buffers to `physical_length` slots and (re)tag the static address. - - Only ever called from the eager prefill pass (initial allocation or chunked-prefill growth), never from a - compiled decode step, so reallocating here is safe for cudagraphs. - """ - new_keys = torch.zeros( - (self.batch_size, self.num_heads, physical_length, self.k_head_dim), dtype=self.dtype, device=self.device - ) - new_values = torch.zeros( - (self.batch_size, self.num_heads, physical_length, self.v_head_dim), dtype=self.dtype, device=self.device - ) - if copy_existing: - old_length = self.keys.shape[-2] - new_keys[:, :, :old_length, :].copy_(self.keys) - new_values[:, :, :old_length, :].copy_(self.values) - self.keys = new_keys - self.values = new_values - if not is_torchdynamo_compiling(): - torch._dynamo.mark_static_address(self.keys) - torch._dynamo.mark_static_address(self.values) - - class StaticIndexedLayer(StaticLayer): """ A `StaticLayer` with an additional statically-allocated indexer key cache for Dynamic Sparse diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index c37b5c206f45..5421b33a5119 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -29,7 +29,8 @@ from ... import initialization as init from ...activations import ACT2FN -from ...cache_utils import Cache, DynamicCache +from ...cache_utils import Cache, DynamicCache, DynamicLayer, DynamicSlidingWindowLayer, StaticSlidingWindowLayer +from ...configuration_utils import PreTrainedConfig from ...generation import GenerationMixin from ...integrations import ( use_experts_implementation, @@ -48,6 +49,7 @@ TransformersKwargs, auto_docstring, can_return_tuple, + is_torchdynamo_compiling, logging, torch_compilable_check, torch_int, @@ -1268,6 +1270,330 @@ def _init_weights(self, module): init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range) +class UnlimitedOcrDynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): + """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens before + the first decode step and applies a sliding window to all decoded tokens. + + Once ``sliding_window`` decode tokens have accumulated, the oldest decode tokens are evicted and + replaced by the most recent ones. The prefill tokens always remain in the cache. + """ + + layer_type = "reference_sliding_attention" + + def __init__(self, config: PreTrainedConfig | None = None, sliding_window: int | None = None): + super().__init__(config=config, sliding_window=sliding_window) + self.prefill_length: int | None = None + + def update( + self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs + ) -> tuple[torch.Tensor, torch.Tensor]: + # Lazy initialization + if not self.is_initialized: + self.lazy_initialization(key_states, value_states) + + kv_length = key_states.shape[-2] + self.cumulative_length += kv_length + + # Compute the full states + full_key_states = torch.cat([self.keys, key_states], dim=-2) + full_value_states = torch.cat([self.values, value_states], dim=-2) + + # Prefill + if self.prefill_length is None and kv_length > 1: + self.keys = full_key_states + self.values = full_value_states + return self.keys, self.values + + # First decode step (or an empty prompt that skipped prefill above): mark prefill as complete + if self.prefill_length is None: + self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 + + generated_length = full_key_states.shape[-2] - self.prefill_length + if generated_length < self.sliding_window: + # Window still growing: keep every prefill and decode token. + self.keys = full_key_states + self.values = full_value_states + else: + # TODO: copy when already full + # Window full: keep all prefill tokens plus the most recent `sliding_window - 1` decode tokens. + # The `- 1` mirrors `DynamicSlidingWindowLayer`: the incoming query token supplies the last slot, + # so the current query still attends over a full `sliding_window` of decode tokens. + self.keys = torch.cat( + [ + full_key_states[:, :, : self.prefill_length, :], + full_key_states[:, :, -self.sliding_window + 1 :, :], + ], + dim=-2, + ) + self.values = torch.cat( + [ + full_value_states[:, :, : self.prefill_length, :], + full_value_states[:, :, -self.sliding_window + 1 :, :], + ], + dim=-2, + ) + + # Return full states to avoid losing context in case we added more than sliding_window tokens at once + return full_key_states, full_value_states + + def get_mask_sizes(self, query_length: int) -> tuple[int, int]: + """Return the length and offset of the cache, used to generate the attention mask""" + is_full = ( + self.prefill_length is not None and self.cumulative_length >= self.prefill_length + self.sliding_window + ) + + kv_offset = 0 + if is_full: + kv_offset = max(self.cumulative_length - self.prefill_length - self.sliding_window + 1, 0) + kv_length = self.prefill_length + self.sliding_window - 1 + query_length + else: + kv_length = self.cumulative_length + query_length + + # Returned kv_offset is with respect to sliding window keys. + # Remove kv_offset from kv_idx to retrieve the prefill indices. + return kv_length, kv_offset + + def get_max_length(self) -> int: + """Return the maximum cache shape of the cache""" + if self.prefill_length is None: + return -1 + return self.prefill_length + self.sliding_window + + def crop(self, max_length: int) -> None: + """ + Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be + negative to remove `max_length` tokens. + """ + if self.prefill_length is not None and self.get_seq_length() >= self.prefill_length + self.sliding_window: + raise ValueError( + "Cannot `crop` a `UnlimitedOcrDynamicReferenceSlidingWindowLayer` after it has seen more tokens than its" + "prefill + sliding window (otherwise some states are lost)" + ) + DynamicLayer.crop(self, max_length) + self.cumulative_length = self.keys.shape[-2] + + +class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): + """ + A static cache layer that stores the key and value states as static tensors of shape + `[batch_size, num_heads, min(max_cache_len, prefill_size + sliding_window), head_dim]`. + It lazily allocates its full backing tensors, and then mutates them in-place. + Built for `torch.compile` support. + + The backing buffer is split in two regions: slots ``[0, prefill_length)`` hold the prefill (reference) slots + that are never eviced from the cache. Slots ``[prefill_length, prefill_length + sliding_window)`` hold the + sliding window decode slots where the oldest entries are always replaced by the newest ones. + + Args: + max_cache_len (`int`): + Maximum number of tokens that can be stored, used for tensor preallocation. + sliding_window (`int`): + The size of the sliding window. + """ + + layer_type = "reference_sliding_attention" + + def __init__(self, max_cache_len: int, sliding_window: int): + super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window) + # Keep `max_cache_len` as max value for lenght bookkeeping. + # The phyiscal buffer doens't exceed `prefill_length + sliding_window`. + self.max_cache_len = max_cache_len + self.sliding_window = sliding_window + self.prefill_length: int | None = None + + def update( + self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Update the key and value caches in-place, and return the necessary keys and value states. + + Args: + key_states (`torch.Tensor`): The new key states to cache. + value_states (`torch.Tensor`): The new value states to cache. + + Returns: + tuple[`torch.Tensor`, `torch.Tensor`]: The key and value states. + """ + # Lazy initialization + if not self.is_initialized: + self.lazy_initialization(key_states, value_states) + + kv_length = key_states.shape[-2] + + # Prefill + if self.prefill_length is None and kv_length > 1: + # Chunked prefill: the buffer was sized for the first chunk only, so grow it to hold all prefill seen so far plus this chunk plus the reserved window. + required_length = min(self.max_cache_len, self.cumulative_length_int + kv_length + self.sliding_window) + if self.keys.shape[-2] < required_length: + self._allocate_key_value_buffers(required_length, copy_existing=True) + # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs + # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten + cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length + try: + self.keys.index_copy_(2, cache_position, key_states) + self.values.index_copy_(2, cache_position, value_states) + except NotImplementedError: + # Fallback for devices like MPS where index_copy_ might not be supported. + self.keys[:, :, cache_position] = key_states + self.values[:, :, cache_position] = value_states + + # Keep both the int (control flow) and the tensor (cudagraph-safe indexing) versions in sync. + self.cumulative_length_int += kv_length + self.cumulative_length.add_(kv_length) + + # Very important to return the `self` tensors here, as they have the static dynamo address + return self.keys, self.values + + # First decode step (or an empty prompt that skipped prefill above): mark prefill as complete + if self.prefill_length is None: + self.prefill_length = self.cumulative_length_int + + # Everything below mirrors `StaticSlidingWindowLayer.update` with the difference that the sliding window + # is in the `sliding_window` slots right after the prefill instead of spanning the whole buffer. + window_start = self.prefill_length + current_length = self.cumulative_length_int - self.prefill_length + is_full = current_length >= self.sliding_window + # Update it now that we saved the value above + self.cumulative_length_int += kv_length + + if is_full: + # In general, we should use a much simpler `cat` here as well, independently of the states size. However, + # dynamo is currently bugged when doing it - see https://github.com/pytorch/pytorch/issues/159855 for more details + if key_states.shape[-2] == 1: + # Roll the window region to the left by 1 position (the pinned prefill in front stays put) + new_keys = self.keys[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) + new_values = self.values[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) + # Overwrite the last position with new states + # (note: very important to use a tensor to index here, see https://github.com/pytorch/pytorch/issues/159855) + index = torch.tensor([-1], dtype=int, device=self.device) + new_keys[:, :, index] = key_states + new_values[:, :, index] = value_states + + # Copy back into `self` (do not just assign again) in order to keep the static dynamo address + self.keys[:, :, window_start : window_start + self.sliding_window, :].copy_(new_keys) + self.values[:, :, window_start : window_start + self.sliding_window, :].copy_(new_values) + + # Very important to return the `self` tensors here, as they have the static dynamo address + return self.keys, self.values + # Already full but using more than 1 new token (e.g. chat continuation, etc...) + else: + full_key_states = torch.cat( + (self.keys[:, :, window_start + 1 : window_start + self.sliding_window, :], key_states), dim=-2 + ) + full_value_states = torch.cat( + (self.values[:, :, window_start + 1 : window_start + self.sliding_window, :], value_states), dim=-2 + ) + # Not yet full, but becoming full on this update + elif current_length + kv_length > self.sliding_window: + # Fast path, no need to cat() in this case, as the window is currently empty + if current_length == 0: + full_key_states = key_states + full_value_states = value_states + else: + window = slice(window_start, window_start + current_length) + full_key_states = torch.cat((self.keys[:, :, window, :], key_states), dim=-2) + full_value_states = torch.cat((self.values[:, :, window, :], value_states), dim=-2) + else: + # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs + # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten. + cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length + try: + self.keys.index_copy_(2, cache_position, key_states) + self.values.index_copy_(2, cache_position, value_states) + except NotImplementedError: + # Fallback for devices like MPS where index_copy_ might not be supported. + self.keys[:, :, cache_position] = key_states + self.values[:, :, cache_position] = value_states + + # Update the tensor version of the length in-place (we don't need to update it if we are already outside + # of this branch, as we don't need the tensor anymore) + self.cumulative_length.add_(kv_length) + + # Very important to return the `self` tensors here, as they have the static dynamo address + return self.keys, self.values + + # We only cache the last `sliding_window` tokens + self.keys[:, :, window_start : window_start + self.sliding_window, :].copy_( + full_key_states[:, :, -self.sliding_window :, :] + ) + self.values[:, :, window_start : window_start + self.sliding_window, :].copy_( + full_value_states[:, :, -self.sliding_window :, :] + ) + return full_key_states, full_value_states + + def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None: + self.dtype, self.device = key_states.dtype, key_states.device + self.batch_size, self.num_heads = key_states.shape[:2] + self.v_head_dim = value_states.shape[-1] + self.k_head_dim = key_states.shape[-1] + + self.cumulative_length = self.cumulative_length.to(self.device) + # Note: `mark_static_address` is used to tag the tensors as a fixed data pointer, preventing compiled graph + # breaks or cudagraph skips due to inplace mutations when updating the cache. However, it is not supported when + # tracing the graph, so we skip it in this case. As prefill should never be compiled, this is not an issue and it + # will still be run (except when users compile prefill explicitly, but this should be avoided!) + # Without this, we cannot use cudagraphs!! + if not is_torchdynamo_compiling(): + torch._dynamo.mark_static_address(self.cumulative_length) + + prefill_seen = key_states.shape[-2] if key_states.shape[-2] > 1 else 0 + self._allocate_key_value_buffers(min(self.max_cache_len, prefill_seen + self.sliding_window)) + + self.is_initialized = True + + def get_mask_sizes(self, query_length: int) -> tuple[int, int]: + """Return the length and offset of the cache, used to generate the attention mask""" + is_full = ( + self.prefill_length is not None and self.cumulative_length_int >= self.prefill_length + self.sliding_window + ) + + kv_offset = 0 + # Prefill + if self.prefill_length is None and query_length > 1: + kv_length = min(self.max_cache_len, self.cumulative_length_int + query_length + self.sliding_window) + # Decode: cache is already full + elif is_full: + kv_offset = max(self.cumulative_length_int - self.prefill_length - self.sliding_window + 1, 0) + kv_length = self.prefill_length + self.sliding_window - 1 + query_length + # Decode: cache not yet full, but becoming full on this update + elif ( + self.prefill_length is not None + and self.cumulative_length_int + query_length > self.prefill_length + self.sliding_window + ): + kv_length = self.cumulative_length_int + query_length + # Decode: cache not yet full but we return the local size as it's static + else: + kv_length = self.keys.shape[-2] + + return kv_length, kv_offset + + def reset(self) -> None: + super().reset() + self.prefill_length = None + + def _allocate_key_value_buffers(self, physical_length: int, copy_existing: bool = False) -> None: + """(Re)allocate the static key/value buffers to `physical_length` slots and (re)tag the static address. + + Only ever called from the eager prefill pass (initial allocation or chunked-prefill growth), never from a + compiled decode step, so reallocating here is safe for cudagraphs. + """ + new_keys = torch.zeros( + (self.batch_size, self.num_heads, physical_length, self.k_head_dim), dtype=self.dtype, device=self.device + ) + new_values = torch.zeros( + (self.batch_size, self.num_heads, physical_length, self.v_head_dim), dtype=self.dtype, device=self.device + ) + if copy_existing: + old_length = self.keys.shape[-2] + new_keys[:, :, :old_length, :].copy_(self.keys) + new_values[:, :, :old_length, :].copy_(self.values) + self.keys = new_keys + self.values = new_values + if not is_torchdynamo_compiling(): + torch._dynamo.mark_static_address(self.keys) + torch._dynamo.mark_static_address(self.values) + + class UnlimitedOcrTextRotaryEmbedding(nn.Module): inv_freq: torch.Tensor # fix linting for `register_buffer` diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index a5d6d4561693..eeb3d4b82302 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -19,8 +19,14 @@ from torchvision.transforms.v2 import functional as tvF from ... import initialization as init -from ...cache_utils import Cache, DynamicCache -from ...configuration_utils import PretrainedConfig +from ...cache_utils import ( + Cache, + DynamicCache, + DynamicLayer, + DynamicSlidingWindowLayer, + StaticSlidingWindowLayer, +) +from ...configuration_utils import PreTrainedConfig, PretrainedConfig from ...feature_extraction_utils import BatchFeature from ...image_transforms import group_images_by_shape, reorder_images from ...image_utils import ImageInput, PILImageResampling, SizeDict @@ -33,6 +39,7 @@ TransformersKwargs, auto_docstring, can_return_tuple, + is_torchdynamo_compiling, torch_int, ) from ...utils.generic import merge_with_config_defaults @@ -607,6 +614,330 @@ class UnlimitedOcrTextPreTrainedModel(DeepseekOcr2TextPreTrainedModel): pass +class UnlimitedOcrDynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): + """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens before + the first decode step and applies a sliding window to all decoded tokens. + + Once ``sliding_window`` decode tokens have accumulated, the oldest decode tokens are evicted and + replaced by the most recent ones. The prefill tokens always remain in the cache. + """ + + layer_type = "reference_sliding_attention" + + def __init__(self, config: PreTrainedConfig | None = None, sliding_window: int | None = None): + super().__init__(config=config, sliding_window=sliding_window) + self.prefill_length: int | None = None + + def update( + self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs + ) -> tuple[torch.Tensor, torch.Tensor]: + # Lazy initialization + if not self.is_initialized: + self.lazy_initialization(key_states, value_states) + + kv_length = key_states.shape[-2] + self.cumulative_length += kv_length + + # Compute the full states + full_key_states = torch.cat([self.keys, key_states], dim=-2) + full_value_states = torch.cat([self.values, value_states], dim=-2) + + # Prefill + if self.prefill_length is None and kv_length > 1: + self.keys = full_key_states + self.values = full_value_states + return self.keys, self.values + + # First decode step (or an empty prompt that skipped prefill above): mark prefill as complete + if self.prefill_length is None: + self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 + + generated_length = full_key_states.shape[-2] - self.prefill_length + if generated_length < self.sliding_window: + # Window still growing: keep every prefill and decode token. + self.keys = full_key_states + self.values = full_value_states + else: + # TODO: copy when already full + # Window full: keep all prefill tokens plus the most recent `sliding_window - 1` decode tokens. + # The `- 1` mirrors `DynamicSlidingWindowLayer`: the incoming query token supplies the last slot, + # so the current query still attends over a full `sliding_window` of decode tokens. + self.keys = torch.cat( + [ + full_key_states[:, :, : self.prefill_length, :], + full_key_states[:, :, -self.sliding_window + 1 :, :], + ], + dim=-2, + ) + self.values = torch.cat( + [ + full_value_states[:, :, : self.prefill_length, :], + full_value_states[:, :, -self.sliding_window + 1 :, :], + ], + dim=-2, + ) + + # Return full states to avoid losing context in case we added more than sliding_window tokens at once + return full_key_states, full_value_states + + def get_mask_sizes(self, query_length: int) -> tuple[int, int]: + """Return the length and offset of the cache, used to generate the attention mask""" + is_full = ( + self.prefill_length is not None and self.cumulative_length >= self.prefill_length + self.sliding_window + ) + + kv_offset = 0 + if is_full: + kv_offset = max(self.cumulative_length - self.prefill_length - self.sliding_window + 1, 0) + kv_length = self.prefill_length + self.sliding_window - 1 + query_length + else: + kv_length = self.cumulative_length + query_length + + # Returned kv_offset is with respect to sliding window keys. + # Remove kv_offset from kv_idx to retrieve the prefill indices. + return kv_length, kv_offset + + def get_max_length(self) -> int: + """Return the maximum cache shape of the cache""" + if self.prefill_length is None: + return -1 + return self.prefill_length + self.sliding_window + + def crop(self, max_length: int) -> None: + """ + Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be + negative to remove `max_length` tokens. + """ + if self.prefill_length is not None and self.get_seq_length() >= self.prefill_length + self.sliding_window: + raise ValueError( + "Cannot `crop` a `UnlimitedOcrDynamicReferenceSlidingWindowLayer` after it has seen more tokens than its" + "prefill + sliding window (otherwise some states are lost)" + ) + DynamicLayer.crop(self, max_length) + self.cumulative_length = self.keys.shape[-2] + + +class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): + """ + A static cache layer that stores the key and value states as static tensors of shape + `[batch_size, num_heads, min(max_cache_len, prefill_size + sliding_window), head_dim]`. + It lazily allocates its full backing tensors, and then mutates them in-place. + Built for `torch.compile` support. + + The backing buffer is split in two regions: slots ``[0, prefill_length)`` hold the prefill (reference) slots + that are never eviced from the cache. Slots ``[prefill_length, prefill_length + sliding_window)`` hold the + sliding window decode slots where the oldest entries are always replaced by the newest ones. + + Args: + max_cache_len (`int`): + Maximum number of tokens that can be stored, used for tensor preallocation. + sliding_window (`int`): + The size of the sliding window. + """ + + layer_type = "reference_sliding_attention" + + def __init__(self, max_cache_len: int, sliding_window: int): + super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window) + # Keep `max_cache_len` as max value for lenght bookkeeping. + # The phyiscal buffer doens't exceed `prefill_length + sliding_window`. + self.max_cache_len = max_cache_len + self.sliding_window = sliding_window + self.prefill_length: int | None = None + + def update( + self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Update the key and value caches in-place, and return the necessary keys and value states. + + Args: + key_states (`torch.Tensor`): The new key states to cache. + value_states (`torch.Tensor`): The new value states to cache. + + Returns: + tuple[`torch.Tensor`, `torch.Tensor`]: The key and value states. + """ + # Lazy initialization + if not self.is_initialized: + self.lazy_initialization(key_states, value_states) + + kv_length = key_states.shape[-2] + + # Prefill + if self.prefill_length is None and kv_length > 1: + # Chunked prefill: the buffer was sized for the first chunk only, so grow it to hold all prefill seen so far plus this chunk plus the reserved window. + required_length = min(self.max_cache_len, self.cumulative_length_int + kv_length + self.sliding_window) + if self.keys.shape[-2] < required_length: + self._allocate_key_value_buffers(required_length, copy_existing=True) + # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs + # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten + cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length + try: + self.keys.index_copy_(2, cache_position, key_states) + self.values.index_copy_(2, cache_position, value_states) + except NotImplementedError: + # Fallback for devices like MPS where index_copy_ might not be supported. + self.keys[:, :, cache_position] = key_states + self.values[:, :, cache_position] = value_states + + # Keep both the int (control flow) and the tensor (cudagraph-safe indexing) versions in sync. + self.cumulative_length_int += kv_length + self.cumulative_length.add_(kv_length) + + # Very important to return the `self` tensors here, as they have the static dynamo address + return self.keys, self.values + + # First decode step (or an empty prompt that skipped prefill above): mark prefill as complete + if self.prefill_length is None: + self.prefill_length = self.cumulative_length_int + + # Everything below mirrors `StaticSlidingWindowLayer.update` with the difference that the sliding window + # is in the `sliding_window` slots right after the prefill instead of spanning the whole buffer. + window_start = self.prefill_length + current_length = self.cumulative_length_int - self.prefill_length + is_full = current_length >= self.sliding_window + # Update it now that we saved the value above + self.cumulative_length_int += kv_length + + if is_full: + # In general, we should use a much simpler `cat` here as well, independently of the states size. However, + # dynamo is currently bugged when doing it - see https://github.com/pytorch/pytorch/issues/159855 for more details + if key_states.shape[-2] == 1: + # Roll the window region to the left by 1 position (the pinned prefill in front stays put) + new_keys = self.keys[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) + new_values = self.values[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) + # Overwrite the last position with new states + # (note: very important to use a tensor to index here, see https://github.com/pytorch/pytorch/issues/159855) + index = torch.tensor([-1], dtype=int, device=self.device) + new_keys[:, :, index] = key_states + new_values[:, :, index] = value_states + + # Copy back into `self` (do not just assign again) in order to keep the static dynamo address + self.keys[:, :, window_start : window_start + self.sliding_window, :].copy_(new_keys) + self.values[:, :, window_start : window_start + self.sliding_window, :].copy_(new_values) + + # Very important to return the `self` tensors here, as they have the static dynamo address + return self.keys, self.values + # Already full but using more than 1 new token (e.g. chat continuation, etc...) + else: + full_key_states = torch.cat( + (self.keys[:, :, window_start + 1 : window_start + self.sliding_window, :], key_states), dim=-2 + ) + full_value_states = torch.cat( + (self.values[:, :, window_start + 1 : window_start + self.sliding_window, :], value_states), dim=-2 + ) + # Not yet full, but becoming full on this update + elif current_length + kv_length > self.sliding_window: + # Fast path, no need to cat() in this case, as the window is currently empty + if current_length == 0: + full_key_states = key_states + full_value_states = value_states + else: + window = slice(window_start, window_start + current_length) + full_key_states = torch.cat((self.keys[:, :, window, :], key_states), dim=-2) + full_value_states = torch.cat((self.values[:, :, window, :], value_states), dim=-2) + else: + # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs + # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten. + cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length + try: + self.keys.index_copy_(2, cache_position, key_states) + self.values.index_copy_(2, cache_position, value_states) + except NotImplementedError: + # Fallback for devices like MPS where index_copy_ might not be supported. + self.keys[:, :, cache_position] = key_states + self.values[:, :, cache_position] = value_states + + # Update the tensor version of the length in-place (we don't need to update it if we are already outside + # of this branch, as we don't need the tensor anymore) + self.cumulative_length.add_(kv_length) + + # Very important to return the `self` tensors here, as they have the static dynamo address + return self.keys, self.values + + # We only cache the last `sliding_window` tokens + self.keys[:, :, window_start : window_start + self.sliding_window, :].copy_( + full_key_states[:, :, -self.sliding_window :, :] + ) + self.values[:, :, window_start : window_start + self.sliding_window, :].copy_( + full_value_states[:, :, -self.sliding_window :, :] + ) + return full_key_states, full_value_states + + def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None: + self.dtype, self.device = key_states.dtype, key_states.device + self.batch_size, self.num_heads = key_states.shape[:2] + self.v_head_dim = value_states.shape[-1] + self.k_head_dim = key_states.shape[-1] + + self.cumulative_length = self.cumulative_length.to(self.device) + # Note: `mark_static_address` is used to tag the tensors as a fixed data pointer, preventing compiled graph + # breaks or cudagraph skips due to inplace mutations when updating the cache. However, it is not supported when + # tracing the graph, so we skip it in this case. As prefill should never be compiled, this is not an issue and it + # will still be run (except when users compile prefill explicitly, but this should be avoided!) + # Without this, we cannot use cudagraphs!! + if not is_torchdynamo_compiling(): + torch._dynamo.mark_static_address(self.cumulative_length) + + prefill_seen = key_states.shape[-2] if key_states.shape[-2] > 1 else 0 + self._allocate_key_value_buffers(min(self.max_cache_len, prefill_seen + self.sliding_window)) + + self.is_initialized = True + + def get_mask_sizes(self, query_length: int) -> tuple[int, int]: + """Return the length and offset of the cache, used to generate the attention mask""" + is_full = ( + self.prefill_length is not None and self.cumulative_length_int >= self.prefill_length + self.sliding_window + ) + + kv_offset = 0 + # Prefill + if self.prefill_length is None and query_length > 1: + kv_length = min(self.max_cache_len, self.cumulative_length_int + query_length + self.sliding_window) + # Decode: cache is already full + elif is_full: + kv_offset = max(self.cumulative_length_int - self.prefill_length - self.sliding_window + 1, 0) + kv_length = self.prefill_length + self.sliding_window - 1 + query_length + # Decode: cache not yet full, but becoming full on this update + elif ( + self.prefill_length is not None + and self.cumulative_length_int + query_length > self.prefill_length + self.sliding_window + ): + kv_length = self.cumulative_length_int + query_length + # Decode: cache not yet full but we return the local size as it's static + else: + kv_length = self.keys.shape[-2] + + return kv_length, kv_offset + + def reset(self) -> None: + super().reset() + self.prefill_length = None + + def _allocate_key_value_buffers(self, physical_length: int, copy_existing: bool = False) -> None: + """(Re)allocate the static key/value buffers to `physical_length` slots and (re)tag the static address. + + Only ever called from the eager prefill pass (initial allocation or chunked-prefill growth), never from a + compiled decode step, so reallocating here is safe for cudagraphs. + """ + new_keys = torch.zeros( + (self.batch_size, self.num_heads, physical_length, self.k_head_dim), dtype=self.dtype, device=self.device + ) + new_values = torch.zeros( + (self.batch_size, self.num_heads, physical_length, self.v_head_dim), dtype=self.dtype, device=self.device + ) + if copy_existing: + old_length = self.keys.shape[-2] + new_keys[:, :, :old_length, :].copy_(self.keys) + new_values[:, :, :old_length, :].copy_(self.values) + self.keys = new_keys + self.values = new_values + if not is_torchdynamo_compiling(): + torch._dynamo.mark_static_address(self.keys) + torch._dynamo.mark_static_address(self.values) + + def create_reference_sliding_window_causal_mask(**kwargs): past_key_values = kwargs["past_key_values"] inputs_embeds = kwargs["inputs_embeds"] From 1e8df96a13abaf65277681db6c68f5a0eb37aa2a Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 1 Jul 2026 18:31:50 +0100 Subject: [PATCH 068/195] remove unused config attributes --- .../models/unlimited_ocr/configuration_unlimited_ocr.py | 1 - src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 2 ++ 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index 15caa6e68fa5..37152b4fa383 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -90,7 +90,6 @@ class UnlimitedOcrVisionEncoderConfig(PreTrainedConfig): base_config_key = "encoder_config" hidden_size: int = 1024 intermediate_size: int = 4096 - projection_dim: int = 512 num_hidden_layers: int = 24 num_attention_heads: int = 16 num_channels: int = 3 diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index eeb3d4b82302..185ee43b7a68 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -355,6 +355,8 @@ class UnlimitedOcrVisionEncoderConfig(CLIPVisionConfig): num_hidden_layers: int = 24 num_attention_heads: int = 16 patch_size: int | list[int] | tuple[int, int] | None = 14 + projection_dim = AttributeError() + sliding_window = AttributeError() @auto_docstring(checkpoint="baidu/Unlimited-OCR") From 817810c6ca3a22a9f76a1604bc5a0bc786f2989c Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 2 Jul 2026 14:23:10 +0100 Subject: [PATCH 069/195] cleanup --- .../configuration_unlimited_ocr.py | 4 +++- .../unlimited_ocr/modeling_unlimited_ocr.py | 16 +++++++------- .../unlimited_ocr/modular_unlimited_ocr.py | 21 ++++++++++--------- utils/check_config_attributes.py | 1 + 4 files changed, 23 insertions(+), 19 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index 37152b4fa383..450c7705ac82 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -229,12 +229,14 @@ class UnlimitedOcrTextConfig(PreTrainedConfig): base_config_key = "text_config" mlp_layer_types: list[str] | None = None layer_types: list[str] | None = None + use_sliding_window: bool = True sliding_window: int | None = 128 def __post_init__(self, **kwargs): + self.sliding_window = self.sliding_window if self.use_sliding_window else None if self.layer_types is None: self.layer_types = [ - "full_attention" if self.sliding_window is None else "reference_sliding_attention" + "reference_sliding_attention" if self.use_sliding_window else "full_attention" ] * self.num_hidden_layers elif len(set(self.layer_types)) > 1: # This requires a custom create_causal_mask implementation for reference_sliding_attention diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 5421b33a5119..ce72707cc4dc 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1308,16 +1308,16 @@ def update( if self.prefill_length is None: self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 - generated_length = full_key_states.shape[-2] - self.prefill_length - if generated_length < self.sliding_window: - # Window still growing: keep every prefill and decode token. + # Cache growing + if self.cumulative_length <= self.prefill_length + self.sliding_window - 1: self.keys = full_key_states self.values = full_value_states + # Cache full + elif self.keys.shape[-2] == self.prefill_length + self.sliding_window - 1: + self.keys[:, :, -self.sliding_window + 1 :].copy_(full_key_states[:, :, -self.sliding_window + 1 :, :]) + self.values[:, :, -self.sliding_window + 1 :].copy_(full_value_states[:, :, -self.sliding_window + 1 :, :]) + # Cache full after this update and full_key_states > cache size else: - # TODO: copy when already full - # Window full: keep all prefill tokens plus the most recent `sliding_window - 1` decode tokens. - # The `- 1` mirrors `DynamicSlidingWindowLayer`: the incoming query token supplies the last slot, - # so the current query still attends over a full `sliding_window` of decode tokens. self.keys = torch.cat( [ full_key_states[:, :, : self.prefill_length, :], @@ -1333,7 +1333,7 @@ def update( dim=-2, ) - # Return full states to avoid losing context in case we added more than sliding_window tokens at once + # Return full states to avoid losing context in case we added multiple tokens at once return full_key_states, full_value_states def get_mask_sizes(self, query_length: int) -> tuple[int, int]: diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 185ee43b7a68..265290d705ae 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -356,7 +356,6 @@ class UnlimitedOcrVisionEncoderConfig(CLIPVisionConfig): num_attention_heads: int = 16 patch_size: int | list[int] | tuple[int, int] | None = 14 projection_dim = AttributeError() - sliding_window = AttributeError() @auto_docstring(checkpoint="baidu/Unlimited-OCR") @@ -419,12 +418,14 @@ class UnlimitedOcrTextConfig(DeepseekOcr2TextConfig): topk_group: int | None = 1 num_experts_per_tok: int | None = 6 layer_types: list[str] | None = None + use_sliding_window: bool = True sliding_window: int | None = 128 def __post_init__(self, **kwargs): + self.sliding_window = self.sliding_window if self.use_sliding_window else None if self.layer_types is None: self.layer_types = [ - "full_attention" if self.sliding_window is None else "reference_sliding_attention" + "reference_sliding_attention" if self.use_sliding_window else "full_attention" ] * self.num_hidden_layers elif len(set(self.layer_types)) > 1: # This requires a custom create_causal_mask implementation for reference_sliding_attention @@ -654,16 +655,16 @@ def update( if self.prefill_length is None: self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 - generated_length = full_key_states.shape[-2] - self.prefill_length - if generated_length < self.sliding_window: - # Window still growing: keep every prefill and decode token. + # Cache growing + if self.cumulative_length <= self.prefill_length + self.sliding_window - 1: self.keys = full_key_states self.values = full_value_states + # Cache full + elif self.keys.shape[-2] == self.prefill_length + self.sliding_window - 1: + self.keys[:, :, -self.sliding_window + 1 :].copy_(full_key_states[:, :, -self.sliding_window + 1 :, :]) + self.values[:, :, -self.sliding_window + 1 :].copy_(full_value_states[:, :, -self.sliding_window + 1 :, :]) + # Cache full after this update and full_key_states > cache size else: - # TODO: copy when already full - # Window full: keep all prefill tokens plus the most recent `sliding_window - 1` decode tokens. - # The `- 1` mirrors `DynamicSlidingWindowLayer`: the incoming query token supplies the last slot, - # so the current query still attends over a full `sliding_window` of decode tokens. self.keys = torch.cat( [ full_key_states[:, :, : self.prefill_length, :], @@ -679,7 +680,7 @@ def update( dim=-2, ) - # Return full states to avoid losing context in case we added more than sliding_window tokens at once + # Return full states to avoid losing context in case we added multiple tokens at once return full_key_states, full_value_states def get_mask_sizes(self, query_length: int) -> tuple[int, int]: diff --git a/utils/check_config_attributes.py b/utils/check_config_attributes.py index 19ce5d0b3f8f..55fff4f5ff4a 100644 --- a/utils/check_config_attributes.py +++ b/utils/check_config_attributes.py @@ -218,6 +218,7 @@ "num_last_full_attention_layers", # builder attr consumed in __post_init__ to compute num_key_value_heads_per_layer "flip_pairs", # used externally for post-processing keypoints, not in forward pass ], + "UnlimitedOcrTextConfig": ["sliding_window"], # only used for cache init } # Common and important attributes, even if they do not always appear in the modeling files (can be a regex pattern) From abdea097378fe5513103096ff1fb3fb46bf76bb1 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 2 Jul 2026 14:29:04 +0100 Subject: [PATCH 070/195] fix check-repo --- utils/check_repo.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/utils/check_repo.py b/utils/check_repo.py index 76e1025dd923..c375ce3322f7 100644 --- a/utils/check_repo.py +++ b/utils/check_repo.py @@ -267,6 +267,8 @@ "InternVLVisionModel", # Building part of bigger (tested) model "DeepseekOcr2TextModel", # Building part of bigger (tested) model "DeepseekOcr2VisionModel", # Building part of bigger (tested) model + "UnlimitedOcrTextModel", # Building part of bigger (tested) model + "UnlimitedOcrVisionModel", # Building part of bigger (tested) model "QianfanOCRVisionModel", # Building part of bigger (tested) model "JanusVisionModel", # Building part of bigger (tested) model "PPDocLayoutV3Model", # Building part of bigger (tested) model @@ -487,6 +489,8 @@ "JanusVisionModel", # Building part of bigger (tested) model "DeepseekOcr2TextModel", # Building part of bigger (tested) model "DeepseekOcr2VisionModel", # Building part of bigger (tested) model + "UnlimitedOcrTextModel", # Building part of bigger (tested) model + "UnlimitedOcrVisionModel", # Building part of bigger (tested) model "SLANetSLAHead", # Building part of bigger (tested) model "SLANetBackbone", # Building part of bigger (tested) model "SLANeXtSLAHead", # Building part of bigger (tested) model From f699734285913452f692effcb9806f780d52cfde Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 2 Jul 2026 18:20:42 +0100 Subject: [PATCH 071/195] handle weird resize/padding --- .../image_processing_unlimited_ocr.py | 37 ++++++++----- .../unlimited_ocr/modular_unlimited_ocr.py | 53 +++++++++++++++---- 2 files changed, 69 insertions(+), 21 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index 571f7457f98f..ebb29d900ca9 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -33,17 +33,21 @@ class UnlimitedOcrImageProcessorKwargs(ImagesKwargs, total=False): r""" - crop_to_patches (`bool`, *optional*, defaults to `self.crop_to_patches`): + crop_to_patches (`bool`, *optional*, defaults to `True`): Whether to crop the image to patches. Can be overridden by the `crop_to_patches` parameter in the `preprocess` method. - min_patches (`int`, *optional*, defaults to `self.min_patches`): + min_patches (`int`, *optional*, defaults to `2`): The minimum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is set to `True`. Can be overridden by the `min_patches` parameter in the `preprocess` method. - max_patches (`int`, *optional*, defaults to `self.max_patches`): + max_patches (`int`, *optional*, defaults to `32`): The maximum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is set to `True`. Can be overridden by the `max_patches` parameter in the `preprocess` method. - tile_size (`int`, *optional*, defaults to `768`): + tile_size (`int`, *optional*, defaults to `640`): The size of each local tile. Must match the model's query embedding size. + pad_if_larger_than (`int`, *optional*, defaults to `640`): + If `crop_to_patches` is `False` and `size.height/width` is larger than this value, + the image will be resized directly to `size.height/width` without padding. Otherwise, + images are resized and padded to `size.height/width` while preserving the aspect ratio. background_color (`list[int]`, *optional*, defaults to `[127, 127, 127]`): The background color for padding. """ @@ -55,6 +59,8 @@ class UnlimitedOcrImageProcessorKwargs(ImagesKwargs, total=False): tile_size: int background_color: list[int] + pad_if_larger_than: int + @lru_cache(maxsize=10) def get_all_supported_aspect_ratios(min_image_tiles: int, max_image_tiles: int) -> list[tuple[int, int]]: @@ -147,6 +153,7 @@ class UnlimitedOcrImageProcessor(TorchvisionBackend): tile_size = 640 background_color = [127, 127, 127] model_input_names = ["pixel_values", "num_local_patches", "patches_grid"] + pad_if_larger_than = 640 def __init__(self, **kwargs: Unpack[UnlimitedOcrImageProcessorKwargs]): super().__init__(**kwargs) @@ -213,6 +220,7 @@ def _preprocess( min_patches: int, max_patches: int, tile_size: int, + pad_if_larger_than: int, resample: "PILImageResampling | None", do_rescale: bool, rescale_factor: float, @@ -232,7 +240,7 @@ def _preprocess( for shape, stacked_images in grouped_images.items(): h, w = shape[-2:] if max(h, w) > tile_size: - stacked_patches, n_patches = self.crop_image_to_patches( + stacked_patches, _ = self.crop_image_to_patches( stacked_images, min_patches=min_patches, max_patches=max_patches, @@ -254,17 +262,22 @@ def _preprocess( flat_local_list = [patch for item in ordered_local if item is not None for patch in item] # --- Global view (batched by shape group) --- - global_target_size = size.height + global_target_size = max(size.height, size.width) grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) processed_global_grouped = {} for shape, stacked in grouped_images.items(): - h, w = shape[-2:] - scale = global_target_size / max(h, w) - new_h = round(h * scale) - new_w = round(w * scale) - stacked = self.resize(stacked, SizeDict(height=new_h, width=new_w), resample=resample) - stacked = self.pad_to_square(stacked, background_color=self.background_color) + if not crop_to_patches and global_target_size <= pad_if_larger_than: + stacked = self.resize( + stacked, SizeDict(height=global_target_size, width=global_target_size), resample=resample + ) + else: + h, w = shape[-2:] + scale = global_target_size / max(h, w) + new_h = round(h * scale) + new_w = round(w * scale) + stacked = self.resize(stacked, SizeDict(height=new_h, width=new_w), resample=resample) + stacked = self.pad_to_square(stacked, background_color=self.background_color) stacked = self.rescale_and_normalize( stacked, do_rescale, rescale_factor, do_normalize, image_mean, image_std ) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 265290d705ae..89385604f4f7 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -51,7 +51,11 @@ DeepseekOcr2TextConfig, DeepseekOcr2VisionConfig, ) -from ..deepseek_ocr2.image_processing_deepseek_ocr2 import DeepseekOcr2ImageProcessor, get_optimal_tiled_canvas +from ..deepseek_ocr2.image_processing_deepseek_ocr2 import ( + DeepseekOcr2ImageProcessor, + DeepseekOcr2ImageProcessorKwargs, + get_optimal_tiled_canvas, +) from ..deepseek_ocr2.modeling_deepseek_ocr2 import ( DeepseekOcr2CausalLMOutputWithPast, DeepseekOcr2ForConditionalGeneration, @@ -68,8 +72,33 @@ from ..got_ocr2.configuration_got_ocr2 import GotOcr2VisionConfig +class UnlimitedOcrImageProcessorKwargs(DeepseekOcr2ImageProcessorKwargs): + r""" + crop_to_patches (`bool`, *optional*, defaults to `True`): + Whether to crop the image to patches. Can be overridden by the `crop_to_patches` parameter in the + `preprocess` method. + min_patches (`int`, *optional*, defaults to `2`): + The minimum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is + set to `True`. Can be overridden by the `min_patches` parameter in the `preprocess` method. + max_patches (`int`, *optional*, defaults to `32`): + The maximum number of patches to be extracted from the image. Only has an effect if `crop_to_patches` is + set to `True`. Can be overridden by the `max_patches` parameter in the `preprocess` method. + tile_size (`int`, *optional*, defaults to `640`): + The size of each local tile. Must match the model's query embedding size. + pad_if_larger_than (`int`, *optional*, defaults to `640`): + If `crop_to_patches` is `False` and `size.height/width` is larger than this value, + the image will be resized directly to `size.height/width` without padding. Otherwise, + images are resized and padded to `size.height/width` while preserving the aspect ratio. + background_color (`list[int]`, *optional*, defaults to `[127, 127, 127]`): + The background color for padding. + """ + + pad_if_larger_than: int + + class UnlimitedOcrImageProcessor(DeepseekOcr2ImageProcessor): tile_size = 640 + pad_if_larger_than = 640 max_patches = 32 model_input_names = ["pixel_values", "num_local_patches", "patches_grid"] @@ -124,6 +153,7 @@ def _preprocess( min_patches: int, max_patches: int, tile_size: int, + pad_if_larger_than: int, resample: "PILImageResampling | None", do_rescale: bool, rescale_factor: float, @@ -143,7 +173,7 @@ def _preprocess( for shape, stacked_images in grouped_images.items(): h, w = shape[-2:] if max(h, w) > tile_size: - stacked_patches, n_patches = self.crop_image_to_patches( + stacked_patches, _ = self.crop_image_to_patches( stacked_images, min_patches=min_patches, max_patches=max_patches, @@ -165,17 +195,22 @@ def _preprocess( flat_local_list = [patch for item in ordered_local if item is not None for patch in item] # --- Global view (batched by shape group) --- - global_target_size = size.height + global_target_size = max(size.height, size.width) grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) processed_global_grouped = {} for shape, stacked in grouped_images.items(): - h, w = shape[-2:] - scale = global_target_size / max(h, w) - new_h = round(h * scale) - new_w = round(w * scale) - stacked = self.resize(stacked, SizeDict(height=new_h, width=new_w), resample=resample) - stacked = self.pad_to_square(stacked, background_color=self.background_color) + if not crop_to_patches and global_target_size <= pad_if_larger_than: + stacked = self.resize( + stacked, SizeDict(height=global_target_size, width=global_target_size), resample=resample + ) + else: + h, w = shape[-2:] + scale = global_target_size / max(h, w) + new_h = round(h * scale) + new_w = round(w * scale) + stacked = self.resize(stacked, SizeDict(height=new_h, width=new_w), resample=resample) + stacked = self.pad_to_square(stacked, background_color=self.background_color) stacked = self.rescale_and_normalize( stacked, do_rescale, rescale_factor, do_normalize, image_mean, image_std ) From d86f5c94ba460ca6b575c4d3437f38500c5247ab Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 2 Jul 2026 18:25:18 +0100 Subject: [PATCH 072/195] handle past_key_values is None --- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 10 +++++++--- .../models/unlimited_ocr/modular_unlimited_ocr.py | 10 +++++++--- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index ce72707cc4dc..a24c158fc359 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1663,9 +1663,13 @@ def create_reference_sliding_window_causal_mask(**kwargs): past_key_values = kwargs["past_key_values"] inputs_embeds = kwargs["inputs_embeds"] - layer = next(layer for layer in past_key_values.layers if layer.layer_type == "reference_sliding_attention") - prefill_length = float("inf") if layer.prefill_length is None else layer.prefill_length - _, kv_offset = layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) + if past_key_values is None: + prefill_length = float("inf") + kv_offset = 0 + else: + layer = next(layer for layer in past_key_values.layers if layer.layer_type == "reference_sliding_attention") + prefill_length = float("inf") if layer.prefill_length is None else layer.prefill_length + _, kv_offset = layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): # Remove kv_offset to retrieve the kv_index with respect to prefill diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 89385604f4f7..771a141171a3 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -980,9 +980,13 @@ def create_reference_sliding_window_causal_mask(**kwargs): past_key_values = kwargs["past_key_values"] inputs_embeds = kwargs["inputs_embeds"] - layer = next(layer for layer in past_key_values.layers if layer.layer_type == "reference_sliding_attention") - prefill_length = float("inf") if layer.prefill_length is None else layer.prefill_length - _, kv_offset = layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) + if past_key_values is None: + prefill_length = float("inf") + kv_offset = 0 + else: + layer = next(layer for layer in past_key_values.layers if layer.layer_type == "reference_sliding_attention") + prefill_length = float("inf") if layer.prefill_length is None else layer.prefill_length + _, kv_offset = layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): # Remove kv_offset to retrieve the kv_index with respect to prefill From 5024831ff3885bf83019ce5886e0af6d5411815c Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 09:21:14 +0200 Subject: [PATCH 073/195] update deepseek ocr2 processing --- .../deepseek_ocr2/processing_deepseek_ocr2.py | 85 ++----------------- 1 file changed, 6 insertions(+), 79 deletions(-) diff --git a/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py index dda5210559e4..86f5413977cd 100644 --- a/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py @@ -17,10 +17,8 @@ import math -from ...feature_extraction_utils import BatchFeature -from ...image_utils import ImageInput -from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack -from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...processing_utils import ProcessingKwargs, ProcessorMixin +from ...tokenization_utils_base import TextInput from ...utils import auto_docstring, logging @@ -63,24 +61,7 @@ def __init__( self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token) super().__init__(image_processor, tokenizer, chat_template=chat_template, **kwargs) - def _expand_image_tokens( - self, - text: list[TextInput], - num_crops_list: list[int], - ) -> list[str]: - """ - Expand each `` placeholder in the text to the correct number of image tokens. - - Args: - text (`list[str]`): - List of text strings, each potentially containing `` placeholders. - num_crops_list (`list[int]`): - Number of crops for each image, consumed in order as `` placeholders - are encountered across all text samples. - - Returns: - `list[str]`: Text with expanded image token placeholders. - """ + def replace_image_token(self, image_inputs: dict, image_idx: int) -> TextInput: size = self.image_processor.size["height"] tile_size = self.image_processor.tile_size @@ -90,63 +71,9 @@ def _expand_image_tokens( num_queries_local = math.ceil(tile_size / self.patch_size / self.downsample_ratio) local_tokens = num_queries_local * num_queries_local - crop_index = 0 - for i in range(len(text)): - while self.image_token in text[i]: - num_tokens = global_tokens + local_tokens * num_crops_list[crop_index] + 1 - text[i] = text[i].replace(self.image_token, "<|placeholder|>" * num_tokens, 1) - crop_index += 1 - text[i] = text[i].replace("<|placeholder|>", self.image_token) - return text - - @auto_docstring - def __call__( - self, - images: ImageInput | None = None, - text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, - **kwargs: Unpack[DeepseekOcr2ProcessorKwargs], - ) -> BatchFeature: - r""" - Returns: - [`BatchFeature`]: A [`BatchFeature`] with the following fields: - - - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. - - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when - `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not - `None`). - - **pixel_values** -- Global view pixel values. Returned when `images` is not `None`. - - **pixel_values_local** -- Local patch pixel values. Returned when `images` is not `None`. - """ - if images is None: - raise ValueError("`images` are expected as arguments to a `DeepseekOcr2Processor` instance.") - if text is None: - raise ValueError("`text` is required for `DeepseekOcr2Processor`. Example: `'\\nFree OCR.'`") - - output_kwargs = self._merge_kwargs( - DeepseekOcr2ProcessorKwargs, - tokenizer_init_kwargs=self.tokenizer.init_kwargs, - **kwargs, - ) - - if isinstance(text, str): - text = [text] - elif not (isinstance(text, (list, tuple)) and all(isinstance(t, str) for t in text)): - raise TypeError("Invalid input text. Please provide a string, or a list of strings") - - text = text.copy() # below lines change text in-place - - image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) - num_crops_list = image_inputs["num_local_patches"] - text = self._expand_image_tokens(text, num_crops_list) - - return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) - text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) - self._check_special_mm_tokens(text, text_inputs, modalities=["image"]) - - return BatchFeature( - data={**text_inputs, **image_inputs}, - tensor_type=return_tensors, - ) + num_crops = image_inputs["num_local_patches"][image_idx] + num_tokens = global_tokens + local_tokens * num_crops + 1 + return self.image_token * num_tokens __all__ = ["DeepseekOcr2Processor"] From 690cd92bf48ad3ff8df1352014fae1510d08eb13 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 09:21:33 +0200 Subject: [PATCH 074/195] update date --- docs/source/en/model_doc/unlimited_ocr.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 04c9d008bbcf..13fe41f690fd 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -16,7 +16,7 @@ limitations under the License. ⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer. --> -*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-07-01.* +*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-07-07.* # UnlimitedOcr From 9b54925326d1ef227ce61972d35ddea2dd0f6466 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 09:21:44 +0200 Subject: [PATCH 075/195] use replace_image_token --- .../unlimited_ocr/modular_unlimited_ocr.py | 88 ++------------- .../unlimited_ocr/processing_unlimited_ocr.py | 105 ++---------------- 2 files changed, 18 insertions(+), 175 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 771a141171a3..097d2a521a68 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -29,11 +29,11 @@ from ...configuration_utils import PreTrainedConfig, PretrainedConfig from ...feature_extraction_utils import BatchFeature from ...image_transforms import group_images_by_shape, reorder_images -from ...image_utils import ImageInput, PILImageResampling, SizeDict +from ...image_utils import PILImageResampling, SizeDict from ...masking_utils import and_masks, causal_mask_function, create_causal_mask, create_sliding_window_causal_mask from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast from ...processing_utils import Unpack -from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...tokenization_utils_base import TextInput from ...utils import ( TensorType, TransformersKwargs, @@ -254,89 +254,19 @@ class UnlimitedOcrProcessorKwargs(DeepseekOcr2ProcessorKwargs, total=False): class UnlimitedOcrProcessor(DeepseekOcr2Processor): - def _expand_image_tokens( - self, - text: list[TextInput], - patches_grid: torch.Tensor, - num_local_patches: list[int] | torch.Tensor, - ) -> list[str]: - num_images = len(patches_grid) - total_image_tokens = sum(t.count(self.image_token) for t in text) - if total_image_tokens != num_images: - raise ValueError( - f"Number of `{self.image_token}` tokens in the text ({total_image_tokens}) does not match the " - f"number of images passed ({num_images}). Use one `{self.image_token}` placeholder per image, " - f"e.g. `'' * len(images) + 'Multi page parsing.'`" - ) - + def replace_image_token(self, image_inputs: dict, image_idx: int) -> TextInput: size = self.image_processor.size["height"] tile_size = self.image_processor.tile_size num_queries_global = math.ceil(size / self.patch_size / self.downsample_ratio) num_queries_local = math.ceil(tile_size / self.patch_size / self.downsample_ratio) - crop_index = 0 - for i in range(len(text)): - while self.image_token in text[i]: - num_columns = int(patches_grid[crop_index][0]) - num_rows = int(patches_grid[crop_index][1]) - num_tokens = num_queries_global * (num_queries_global + 1) + 1 - if int(num_local_patches[crop_index]) > 0: - num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) - text[i] = text[i].replace(self.image_token, "<|placeholder|>" * num_tokens, 1) - crop_index += 1 - text[i] = text[i].replace("<|placeholder|>", self.image_token) - return text - - @auto_docstring - def __call__( - self, - images: ImageInput | None = None, - text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, - **kwargs: Unpack[UnlimitedOcrProcessorKwargs], - ) -> BatchFeature: - r""" - Returns: - [`BatchFeature`]: A [`BatchFeature`] with the following fields: - - - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. - - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when - `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not - `None`). - - **pixel_values** -- Global view pixel values. Returned when `images` is not `None`. - - **pixel_values_local** -- Local patch pixel values. Returned when `images` is not `None`. - - **num_local_patches** -- Number of local patches per image. Returned when `images` is not `None`. - - **patches_grid** -- Number of patch columns and rows per image. Returned when `images` is not `None`. - """ - if images is None: - raise ValueError("`images` are expected as arguments to a `UnlimitedOcrProcessor` instance.") - if text is None: - raise ValueError("`text` is required for `UnlimitedOcrProcessor`. Example: `'\\nFree OCR.'`") - - output_kwargs = self._merge_kwargs( - UnlimitedOcrProcessorKwargs, - tokenizer_init_kwargs=self.tokenizer.init_kwargs, - **kwargs, - ) - - if isinstance(text, str): - text = [text] - elif not (isinstance(text, (list, tuple)) and all(isinstance(t, str) for t in text)): - raise TypeError("Invalid input text. Please provide a string, or a list of strings") - - text = text.copy() # below lines change text in-place - - image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) - text = self._expand_image_tokens(text, image_inputs["patches_grid"], image_inputs["num_local_patches"]) - - return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) - text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) - self._check_special_mm_tokens(text, text_inputs, modalities=["image"]) - - return BatchFeature( - data={**text_inputs, **image_inputs}, - tensor_type=return_tensors, - ) + num_columns = int(image_inputs["patches_grid"][image_idx][0]) + num_rows = int(image_inputs["patches_grid"][image_idx][1]) + num_tokens = num_queries_global * (num_queries_global + 1) + 1 + if int(image_inputs["num_local_patches"][image_idx]) > 0: + num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) + return self.image_token * num_tokens @auto_docstring(checkpoint="baidu/Unlimited-OCR") diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index 0aa50c69368b..70c9397d57e6 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -19,12 +19,8 @@ # limitations under the License. import math -import torch - -from ...feature_extraction_utils import BatchFeature -from ...image_utils import ImageInput -from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack -from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...processing_utils import ProcessingKwargs, ProcessorMixin +from ...tokenization_utils_base import TextInput from ...utils import auto_docstring @@ -59,102 +55,19 @@ def __init__( self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token) super().__init__(image_processor, tokenizer, chat_template=chat_template, **kwargs) - def _expand_image_tokens( - self, - text: list[TextInput], - patches_grid: torch.Tensor, - num_local_patches: list[int] | torch.Tensor, - ) -> list[str]: - """ - Expand each `` placeholder in the text to the correct number of image tokens. - - Args: - text (`list[str]`): - List of text strings, each potentially containing `` placeholders. - num_crops_list (`list[int]`): - Number of crops for each image, consumed in order as `` placeholders - are encountered across all text samples. - - Returns: - `list[str]`: Text with expanded image token placeholders. - """ - num_images = len(patches_grid) - total_image_tokens = sum(t.count(self.image_token) for t in text) - if total_image_tokens != num_images: - raise ValueError( - f"Number of `{self.image_token}` tokens in the text ({total_image_tokens}) does not match the " - f"number of images passed ({num_images}). Use one `{self.image_token}` placeholder per image, " - f"e.g. `'' * len(images) + 'Multi page parsing.'`" - ) - + def replace_image_token(self, image_inputs: dict, image_idx: int) -> TextInput: size = self.image_processor.size["height"] tile_size = self.image_processor.tile_size num_queries_global = math.ceil(size / self.patch_size / self.downsample_ratio) num_queries_local = math.ceil(tile_size / self.patch_size / self.downsample_ratio) - crop_index = 0 - for i in range(len(text)): - while self.image_token in text[i]: - num_columns = int(patches_grid[crop_index][0]) - num_rows = int(patches_grid[crop_index][1]) - num_tokens = num_queries_global * (num_queries_global + 1) + 1 - if int(num_local_patches[crop_index]) > 0: - num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) - text[i] = text[i].replace(self.image_token, "<|placeholder|>" * num_tokens, 1) - crop_index += 1 - text[i] = text[i].replace("<|placeholder|>", self.image_token) - return text - - @auto_docstring - def __call__( - self, - images: ImageInput | None = None, - text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, - **kwargs: Unpack[UnlimitedOcrProcessorKwargs], - ) -> BatchFeature: - r""" - Returns: - [`BatchFeature`]: A [`BatchFeature`] with the following fields: - - - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. - - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when - `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not - `None`). - - **pixel_values** -- Global view pixel values. Returned when `images` is not `None`. - - **pixel_values_local** -- Local patch pixel values. Returned when `images` is not `None`. - - **num_local_patches** -- Number of local patches per image. Returned when `images` is not `None`. - - **patches_grid** -- Number of patch columns and rows per image. Returned when `images` is not `None`. - """ - if images is None: - raise ValueError("`images` are expected as arguments to a `UnlimitedOcrProcessor` instance.") - if text is None: - raise ValueError("`text` is required for `UnlimitedOcrProcessor`. Example: `'\\nFree OCR.'`") - - output_kwargs = self._merge_kwargs( - UnlimitedOcrProcessorKwargs, - tokenizer_init_kwargs=self.tokenizer.init_kwargs, - **kwargs, - ) - - if isinstance(text, str): - text = [text] - elif not (isinstance(text, (list, tuple)) and all(isinstance(t, str) for t in text)): - raise TypeError("Invalid input text. Please provide a string, or a list of strings") - - text = text.copy() # below lines change text in-place - - image_inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) - text = self._expand_image_tokens(text, image_inputs["patches_grid"], image_inputs["num_local_patches"]) - - return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) - text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) - self._check_special_mm_tokens(text, text_inputs, modalities=["image"]) - - return BatchFeature( - data={**text_inputs, **image_inputs}, - tensor_type=return_tensors, - ) + num_columns = int(image_inputs["patches_grid"][image_idx][0]) + num_rows = int(image_inputs["patches_grid"][image_idx][1]) + num_tokens = num_queries_global * (num_queries_global + 1) + 1 + if int(image_inputs["num_local_patches"][image_idx]) > 0: + num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) + return self.image_token * num_tokens __all__ = ["UnlimitedOcrProcessor"] From 10737af64cccd40d11abf25f9ddca7182a2b6aa9 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 10:05:29 +0200 Subject: [PATCH 076/195] add processing tests --- .../deepseek_ocr2/processing_deepseek_ocr2.py | 14 +++++++++++ .../unlimited_ocr/processing_unlimited_ocr.py | 14 +++++++++++ .../test_processing_unlimited_ocr.py | 25 +++++++++++++++++++ 3 files changed, 53 insertions(+) create mode 100644 tests/models/unlimited_ocr/test_processing_unlimited_ocr.py diff --git a/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py index 86f5413977cd..73d07a670366 100644 --- a/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py @@ -17,6 +17,7 @@ import math +from ...image_utils import make_flat_list_of_images from ...processing_utils import ProcessingKwargs, ProcessorMixin from ...tokenization_utils_base import TextInput from ...utils import auto_docstring, logging @@ -61,6 +62,19 @@ def __init__( self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token) super().__init__(image_processor, tokenizer, chat_template=chat_template, **kwargs) + def validate_inputs(self, images=None, text=None, videos=None, audio=None, **kwargs): + super().validate_inputs(images=images, text=text, **kwargs) + if text is not None and images is not None: + if isinstance(text, str): + text = [text] + n_tokens = sum(sample.count(self.image_token) for sample in text) + n_images = len(make_flat_list_of_images(images)) + if n_tokens != n_images: + raise ValueError( + f"Number of {self.image_token} tokens in text ({n_tokens}) does not match " + f"number of images ({n_images})." + ) + def replace_image_token(self, image_inputs: dict, image_idx: int) -> TextInput: size = self.image_processor.size["height"] tile_size = self.image_processor.tile_size diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index 70c9397d57e6..ba42c6f3ac94 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -19,6 +19,7 @@ # limitations under the License. import math +from ...image_utils import make_flat_list_of_images from ...processing_utils import ProcessingKwargs, ProcessorMixin from ...tokenization_utils_base import TextInput from ...utils import auto_docstring @@ -55,6 +56,19 @@ def __init__( self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token) super().__init__(image_processor, tokenizer, chat_template=chat_template, **kwargs) + def validate_inputs(self, images=None, text=None, videos=None, audio=None, **kwargs): + super().validate_inputs(images=images, text=text, **kwargs) + if text is not None and images is not None: + if isinstance(text, str): + text = [text] + n_tokens = sum(sample.count(self.image_token) for sample in text) + n_images = len(make_flat_list_of_images(images)) + if n_tokens != n_images: + raise ValueError( + f"Number of {self.image_token} tokens in text ({n_tokens}) does not match " + f"number of images ({n_images})." + ) + def replace_image_token(self, image_inputs: dict, image_idx: int) -> TextInput: size = self.image_processor.size["height"] tile_size = self.image_processor.tile_size diff --git a/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py b/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py new file mode 100644 index 000000000000..588138d6f31a --- /dev/null +++ b/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py @@ -0,0 +1,25 @@ +import unittest + +from transformers.testing_utils import require_vision +from transformers.utils import is_vision_available + +from ...test_processing_common import ProcessorTesterMixin + + +if is_vision_available(): + from transformers import UnlimitedOcrProcessor + + +@require_vision +class UnlimitedOcrProcessorTest(ProcessorTesterMixin, unittest.TestCase): + processor_class = UnlimitedOcrProcessor + # TODO: Change before merge + model_id = "guarin/Unlimited-OCR" + + @classmethod + def _setup_image_processor(cls): + image_processor_class = cls._get_component_class_from_processor("image_processor") + image_processor = image_processor_class() + image_processor.size = {"height": 64, "width": 64} + image_processor.tile_size = 512 + return image_processor From 23aabee62260c11e511a824bbefd0679b56d76e1 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 10:07:55 +0200 Subject: [PATCH 077/195] update tests --- .../unlimited_ocr/test_processing_unlimited_ocr.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py b/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py index 588138d6f31a..288684aa842a 100644 --- a/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py @@ -16,10 +16,8 @@ class UnlimitedOcrProcessorTest(ProcessorTesterMixin, unittest.TestCase): # TODO: Change before merge model_id = "guarin/Unlimited-OCR" - @classmethod - def _setup_image_processor(cls): - image_processor_class = cls._get_component_class_from_processor("image_processor") - image_processor = image_processor_class() - image_processor.size = {"height": 64, "width": 64} - image_processor.tile_size = 512 - return image_processor + # Defaults from mixin are too small as a single image expands to 273 image tokens + # for this checkpoint (size=1024) + image_text_kwargs_max_length = 320 + image_text_kwargs_override_max_length = 310 + image_unstructured_max_length = 300 From b7f596f9aee3f717b907611b4ad99ba0bb145a3c Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 10:25:42 +0200 Subject: [PATCH 078/195] clean up test --- tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 95605e6972fd..967a7660f489 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -172,6 +172,7 @@ def _image_features_prepare_config_and_inputs(self): @require_torch class UnlimitedOcrIntegrationTest(unittest.TestCase): + # TODO: change before merge model_id = "guarin/Unlimited-OCR" def setUp(self): @@ -242,7 +243,6 @@ def test_small_model_integration_test_document_parsing_batched(self): images=[image1, image2], text=["document parsing.", "document parsing."], return_tensors="pt", - padding=True, ).to(model.device, dtype=torch.bfloat16) with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): From ba1ed65ec763c57b795ee57b6362ce6896195261 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 11:09:33 +0200 Subject: [PATCH 079/195] add all imports to __init__ --- src/transformers/models/unlimited_ocr/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/transformers/models/unlimited_ocr/__init__.py b/src/transformers/models/unlimited_ocr/__init__.py index 6a34d47b1edc..737519feca9a 100644 --- a/src/transformers/models/unlimited_ocr/__init__.py +++ b/src/transformers/models/unlimited_ocr/__init__.py @@ -20,7 +20,9 @@ if TYPE_CHECKING: from .configuration_unlimited_ocr import * + from .image_processing_unlimited_ocr import * from .modeling_unlimited_ocr import * + from .processing_unlimited_ocr import * else: import sys From 6b0b7d909c30603a140e72f51667b63f6ec6b57c Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 11:11:54 +0200 Subject: [PATCH 080/195] updaet names --- .../image_processing_unlimited_ocr.py | 14 +++++++------- .../models/unlimited_ocr/modular_unlimited_ocr.py | 14 +++++++------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index ebb29d900ca9..766da92b49e2 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -238,8 +238,8 @@ def _preprocess( grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) for shape, stacked_images in grouped_images.items(): - h, w = shape[-2:] - if max(h, w) > tile_size: + height, width = shape[-2:] + if max(height, width) > tile_size: stacked_patches, _ = self.crop_image_to_patches( stacked_images, min_patches=min_patches, @@ -272,11 +272,11 @@ def _preprocess( stacked, SizeDict(height=global_target_size, width=global_target_size), resample=resample ) else: - h, w = shape[-2:] - scale = global_target_size / max(h, w) - new_h = round(h * scale) - new_w = round(w * scale) - stacked = self.resize(stacked, SizeDict(height=new_h, width=new_w), resample=resample) + height, width = shape[-2:] + scale = global_target_size / max(height, width) + new_height = round(height * scale) + new_width = round(width * scale) + stacked = self.resize(stacked, SizeDict(height=new_height, width=new_width), resample=resample) stacked = self.pad_to_square(stacked, background_color=self.background_color) stacked = self.rescale_and_normalize( stacked, do_rescale, rescale_factor, do_normalize, image_mean, image_std diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 097d2a521a68..ba618b77fed0 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -171,8 +171,8 @@ def _preprocess( grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) for shape, stacked_images in grouped_images.items(): - h, w = shape[-2:] - if max(h, w) > tile_size: + height, width = shape[-2:] + if max(height, width) > tile_size: stacked_patches, _ = self.crop_image_to_patches( stacked_images, min_patches=min_patches, @@ -205,11 +205,11 @@ def _preprocess( stacked, SizeDict(height=global_target_size, width=global_target_size), resample=resample ) else: - h, w = shape[-2:] - scale = global_target_size / max(h, w) - new_h = round(h * scale) - new_w = round(w * scale) - stacked = self.resize(stacked, SizeDict(height=new_h, width=new_w), resample=resample) + height, width = shape[-2:] + scale = global_target_size / max(height, width) + new_height = round(height * scale) + new_width = round(width * scale) + stacked = self.resize(stacked, SizeDict(height=new_height, width=new_width), resample=resample) stacked = self.pad_to_square(stacked, background_color=self.background_color) stacked = self.rescale_and_normalize( stacked, do_rescale, rescale_factor, do_normalize, image_mean, image_std From f2d05e9ccf92b873a90217cf1cc6b74c4dc9a8fb Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 11:13:26 +0200 Subject: [PATCH 081/195] typos --- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 4 ++-- .../models/unlimited_ocr/modular_unlimited_ocr.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index a24c158fc359..72a5f52b8b69 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1395,8 +1395,8 @@ class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): def __init__(self, max_cache_len: int, sliding_window: int): super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window) - # Keep `max_cache_len` as max value for lenght bookkeeping. - # The phyiscal buffer doens't exceed `prefill_length + sliding_window`. + # Keep `max_cache_len` as max value for length bookkeeping. + # The physical buffer doesn't exceed `prefill_length + sliding_window`. self.max_cache_len = max_cache_len self.sliding_window = sliding_window self.prefill_length: int | None = None diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index ba618b77fed0..1974f2f8b64e 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -707,8 +707,8 @@ class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): def __init__(self, max_cache_len: int, sliding_window: int): super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window) - # Keep `max_cache_len` as max value for lenght bookkeeping. - # The phyiscal buffer doens't exceed `prefill_length + sliding_window`. + # Keep `max_cache_len` as max value for length bookkeeping. + # The physical buffer doesn't exceed `prefill_length + sliding_window`. self.max_cache_len = max_cache_len self.sliding_window = sliding_window self.prefill_length: int | None = None From 7d5658f7cfe5cee49d9aa39fb658a4bef555f97e Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 11:48:13 +0200 Subject: [PATCH 082/195] disable flex attn --- .../models/deepseek_ocr2/modeling_deepseek_ocr2.py | 3 ++- src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py | 2 ++ .../models/unlimited_ocr/modeling_unlimited_ocr.py | 3 ++- 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py index ee51dab22d6d..16edfe514b09 100644 --- a/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py @@ -143,7 +143,8 @@ class DeepseekOcr2PreTrainedModel(PreTrainedModel): _supports_sdpa = True _can_compile_fullgraph = True - _supports_flex_attn = True + # SAM doesn't support flex attention + _supports_flex_attn = False _supports_attention_backend = True @torch.no_grad() diff --git a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py index cddbd6f2151e..b5a89bf7335a 100644 --- a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py @@ -682,6 +682,8 @@ class DeepseekOcr2PreTrainedModel(LlavaNextPreTrainedModel): ] # SAM uses rel-pos bias, incompatible with flash attention. _supports_flash_attn = False + # SAM doesn't support flex attention + _supports_flex_attn = False @torch.no_grad() def _init_weights(self, module): diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 72a5f52b8b69..6c3afcb89b2a 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -285,7 +285,8 @@ class UnlimitedOcrPreTrainedModel(PreTrainedModel): _supports_sdpa = True _can_compile_fullgraph = True - _supports_flex_attn = True + # SAM doesn't support flex attention + _supports_flex_attn = False _supports_attention_backend = True @torch.no_grad() From 7739c001cf916448b33940c32507c182aef69786 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 11:49:29 +0200 Subject: [PATCH 083/195] enable torch export tests --- tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 967a7660f489..b10b80b830e1 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -123,7 +123,6 @@ def get_config(self): class UnlimitedOcrModelTest(VLMModelTest, unittest.TestCase): model_tester_class = UnlimitedOcrVisionText2TextModelTester test_all_params_have_gradient = False - test_torch_exportable = False def prepare_config_and_inputs_for_generate(self, batch_size=2): config, inputs_dict = super().prepare_config_and_inputs_for_generate(batch_size=batch_size) From 68a1f7a44895f145d0f9e10c07027a71e8aa59ab Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 11:53:59 +0200 Subject: [PATCH 084/195] update no_split_modules --- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 2 +- .../models/unlimited_ocr/modular_unlimited_ocr.py | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 6c3afcb89b2a..d949bfcc15e6 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -275,8 +275,8 @@ class UnlimitedOcrPreTrainedModel(PreTrainedModel): input_modalities = ("image", "text") supports_gradient_checkpointing = True _no_split_modules = [ + "UnlimitedOcrEncoderLayer", "UnlimitedOcrSamVisionLayer", - "UnlimitedOcrVisionEncoderLayer", "UnlimitedOcrTextDecoderLayer", ] _skip_keys_device_placement = ["past_key_values"] diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 1974f2f8b64e..c57c13a1097f 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -446,6 +446,12 @@ class UnlimitedOcrCausalLMOutputWithPast(DeepseekOcr2CausalLMOutputWithPast): class UnlimitedOcrPreTrainedModel(DeepseekOcr2PreTrainedModel): + _no_split_modules = [ + "UnlimitedOcrEncoderLayer", + "UnlimitedOcrSamVisionLayer", + "UnlimitedOcrTextDecoderLayer", + ] + @torch.no_grad() def _init_weights(self, module): super()._init_weights(module) From 3d77c49b06b618b332880e632d081adee18d97b1 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 12:18:23 +0200 Subject: [PATCH 085/195] enable skipped tests --- .../test_modeling_unlimited_ocr.py | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index b10b80b830e1..8a1486056e5b 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -140,24 +140,6 @@ def prepare_config_and_inputs_for_generate(self, batch_size=2): def test_sdpa_can_dispatch_on_flash(self): pass - @unittest.skip( - reason="UnlimitedOcrVisionModel uses `self.query_*.weight` directly, causing device mismatch when offloading." - ) - def test_cpu_offload(self): - pass - - @unittest.skip( - reason="UnlimitedOcrVisionModel uses `self.query_*.weight` directly, causing device mismatch when offloading." - ) - def test_disk_offload_bin(self): - pass - - @unittest.skip( - reason="UnlimitedOcrVisionModel uses `self.query_*.weight` directly, causing device mismatch when offloading." - ) - def test_disk_offload_safetensors(self): - pass - def _image_features_prepare_config_and_inputs(self): config, inputs_dict = super()._image_features_prepare_config_and_inputs() # `get_image_features` returns the concatenation of the SAM feature map and the CLIP encoder output as From 7ab7f5c340ad41b2ca35dc8bd3089245ad3f5f82 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 12:20:43 +0200 Subject: [PATCH 086/195] update conversion mapping --- src/transformers/conversion_mapping.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index a3a7c6aa0530..a0b1b3f8040f 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -1510,7 +1510,6 @@ def _build_checkpoint_conversion_mapping(): WeightRenaming(r"blocks\.(\d+)\.attn\.", r"layers.\1.attn."), WeightRenaming(r"blocks\.(\d+)\.mlp\.", r"layers.\1.mlp."), WeightRenaming(r"patch_embed\.proj\.", "patch_embed.projection."), - WeightRenaming(r"pos_embed", "pos_embed"), WeightRenaming(r"neck\.0\.", "neck.conv1."), WeightRenaming(r"neck\.1\.", "neck.layer_norm1."), WeightRenaming(r"neck\.2\.", "neck.conv2."), From 81dc956ddd0957974c67a6af313757e99d2ee867 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 12:22:45 +0200 Subject: [PATCH 087/195] remove PretrainedConfig --- .../models/unlimited_ocr/configuration_unlimited_ocr.py | 8 ++++---- .../models/unlimited_ocr/modular_unlimited_ocr.py | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index 450c7705ac82..8735e377cd2e 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -19,7 +19,7 @@ # limitations under the License. from huggingface_hub.dataclasses import strict -from ...configuration_utils import PreTrainedConfig, PretrainedConfig +from ...configuration_utils import PreTrainedConfig from ...modeling_rope_utils import RopeParameters from ...utils import auto_docstring @@ -70,7 +70,7 @@ class UnlimitedOcrSamVisionConfig(PreTrainedConfig): def __post_init__(self, **kwargs): if self.downsample_channels is None: self.downsample_channels = [512, 1024] - return PretrainedConfig.__post_init__(self, **kwargs) + return super().__post_init__(**kwargs) @auto_docstring(checkpoint="baidu/Unlimited-OCR") @@ -141,7 +141,7 @@ def __post_init__(self, **kwargs): elif isinstance(self.encoder_config, dict): self.encoder_config = self.sub_configs["encoder_config"](**self.encoder_config) - PretrainedConfig.__post_init__(self, **kwargs) + super().__post_init__(**kwargs) @auto_docstring(checkpoint="baidu/Unlimited-OCR") @@ -297,7 +297,7 @@ def __post_init__(self, **kwargs): self.text_config = text_config_cls(**text_kwargs) elif isinstance(self.text_config, dict): self.text_config = self.sub_configs["text_config"](**self.text_config) - PretrainedConfig.__post_init__(self, **kwargs) + super().__post_init__(**kwargs) __all__ = [ diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index c57c13a1097f..447168b8b355 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -26,7 +26,7 @@ DynamicSlidingWindowLayer, StaticSlidingWindowLayer, ) -from ...configuration_utils import PreTrainedConfig, PretrainedConfig +from ...configuration_utils import PreTrainedConfig from ...feature_extraction_utils import BatchFeature from ...image_transforms import group_images_by_shape, reorder_images from ...image_utils import PILImageResampling, SizeDict @@ -297,7 +297,7 @@ class UnlimitedOcrSamVisionConfig(GotOcr2VisionConfig): def __post_init__(self, **kwargs): if self.downsample_channels is None: self.downsample_channels = [512, 1024] - return PretrainedConfig.__post_init__(self, **kwargs) + return PreTrainedConfig.__post_init__(self, **kwargs) @auto_docstring(checkpoint="baidu/Unlimited-OCR") @@ -344,7 +344,7 @@ def __post_init__(self, **kwargs): elif isinstance(self.encoder_config, dict): self.encoder_config = self.sub_configs["encoder_config"](**self.encoder_config) - PretrainedConfig.__post_init__(self, **kwargs) + PreTrainedConfig.__post_init__(self, **kwargs) @auto_docstring(checkpoint="baidu/Unlimited-OCR") @@ -430,7 +430,7 @@ def __post_init__(self, **kwargs): self.text_config = text_config_cls(**text_kwargs) elif isinstance(self.text_config, dict): self.text_config = self.sub_configs["text_config"](**self.text_config) - PretrainedConfig.__post_init__(self, **kwargs) + PreTrainedConfig.__post_init__(self, **kwargs) class UnlimitedOcrModelOutputWithPooling(DeepseekOcr2ModelOutputWithPooling): From aa3f072d74690f9bb54b611a3b0302c10f5d53aa Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 15:26:58 +0200 Subject: [PATCH 088/195] add replace_image_tokens tests --- .../unlimited_ocr/modular_unlimited_ocr.py | 6 +-- .../unlimited_ocr/processing_unlimited_ocr.py | 6 +-- .../test_processing_unlimited_ocr.py | 45 +++++++++++++++++++ 3 files changed, 51 insertions(+), 6 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 447168b8b355..0a93b1d62c45 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -255,11 +255,11 @@ class UnlimitedOcrProcessorKwargs(DeepseekOcr2ProcessorKwargs, total=False): class UnlimitedOcrProcessor(DeepseekOcr2Processor): def replace_image_token(self, image_inputs: dict, image_idx: int) -> TextInput: - size = self.image_processor.size["height"] + size = max(self.image_processor.size["height"], self.image_processor.size["width"]) tile_size = self.image_processor.tile_size - num_queries_global = math.ceil(size / self.patch_size / self.downsample_ratio) - num_queries_local = math.ceil(tile_size / self.patch_size / self.downsample_ratio) + num_queries_global = math.ceil(size // self.patch_size / self.downsample_ratio) + num_queries_local = math.ceil(tile_size // self.patch_size / self.downsample_ratio) num_columns = int(image_inputs["patches_grid"][image_idx][0]) num_rows = int(image_inputs["patches_grid"][image_idx][1]) diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index ba42c6f3ac94..1f8ca95c6404 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -70,11 +70,11 @@ def validate_inputs(self, images=None, text=None, videos=None, audio=None, **kwa ) def replace_image_token(self, image_inputs: dict, image_idx: int) -> TextInput: - size = self.image_processor.size["height"] + size = max(self.image_processor.size["height"], self.image_processor.size["width"]) tile_size = self.image_processor.tile_size - num_queries_global = math.ceil(size / self.patch_size / self.downsample_ratio) - num_queries_local = math.ceil(tile_size / self.patch_size / self.downsample_ratio) + num_queries_global = math.ceil(size // self.patch_size / self.downsample_ratio) + num_queries_local = math.ceil(tile_size // self.patch_size / self.downsample_ratio) num_columns = int(image_inputs["patches_grid"][image_idx][0]) num_rows = int(image_inputs["patches_grid"][image_idx][1]) diff --git a/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py b/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py index 288684aa842a..a1c426c6ead5 100644 --- a/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py @@ -1,5 +1,7 @@ import unittest +import torch + from transformers.testing_utils import require_vision from transformers.utils import is_vision_available @@ -21,3 +23,46 @@ class UnlimitedOcrProcessorTest(ProcessorTesterMixin, unittest.TestCase): image_text_kwargs_max_length = 320 image_text_kwargs_override_max_length = 310 image_unstructured_max_length = 300 + + def test_replace_image_tokens(self): + processor = self.get_processor() + + images = torch.randint(0, 256, (1, 3, 200, 300), dtype=torch.uint8) + prompt = "document parsing." + + inputs = processor(images=images, text=prompt, return_tensors="pt") + num_image_tokens = (inputs["input_ids"] == processor.image_token_id).sum().item() + + # image resized to 1024, followed by patch size 16 and 4x downsampling = 16 x 16 patches + # 273 = 16 rows * (16 cols + 1 newline) + 1 view separator + self.assertEqual(num_image_tokens, 273) + self.assertNotIn("pixel_values_local", inputs) + + def test_replace_image_tokens_with_local(self): + processor = self.get_processor() + + images = torch.randint(0, 256, (1, 3, 500, 700), dtype=torch.uint8) + prompt = "document parsing." + + inputs = processor(images=images, text=prompt, return_tensors="pt") + num_image_tokens = (inputs["input_ids"] == processor.image_token_id).sum().item() + + # global is same as in test above + # 500 x 700 image is split into 3x4 tiles + # local tiles are 640, followed by patch size 16 and 4x downsample = 10 x 10 patches + # 1503 = 273 global + (3 rows * 10) * (4 cols * 10 + 1)) local + self.assertEqual(num_image_tokens, 1503) + self.assertIn("pixel_values_local", inputs) + + def test_replace_image_tokens_no_crop(self): + processor = self.get_processor() + + images = torch.randint(0, 256, (1, 3, 500, 700), dtype=torch.uint8) + prompt = "document parsing." + + inputs = processor(images=images, text=prompt, return_tensors="pt", crop_to_patches=False) + num_image_tokens = (inputs["input_ids"] == processor.image_token_id).sum().item() + + # same as in test_replace_image_tokens + self.assertEqual(num_image_tokens, 273) + self.assertNotIn("pixel_values_local", inputs) From 3fe8f6fda6e8d90e1fa2882137692e2f6f8ca937 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 15:34:58 +0200 Subject: [PATCH 089/195] remove can_return_tuple --- src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py | 1 - src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 1 - 2 files changed, 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index d949bfcc15e6..8ba6d23db602 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -911,7 +911,6 @@ def __init__(self, config: UnlimitedOcrVisionEncoderConfig): self.encoder = UnlimitedOcrEncoder(config) self.post_init() - @can_return_tuple @capture_outputs @auto_docstring def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 0a93b1d62c45..c1ccd497831b 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -545,7 +545,6 @@ def __init__(self, config: UnlimitedOcrVisionEncoderConfig): super().__init__(config) del self.post_layernorm - @can_return_tuple @capture_outputs @auto_docstring def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: From a729ebc088a106ca6b4a426afdf063bbb644d231 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 15:44:26 +0200 Subject: [PATCH 090/195] add _input_embed_layer --- .../models/deepseek_ocr2/modeling_deepseek_ocr2.py | 4 +--- .../models/deepseek_ocr2/modular_deepseek_ocr2.py | 5 +++++ .../models/unlimited_ocr/modeling_unlimited_ocr.py | 4 +--- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py index 16edfe514b09..015136cf56a6 100644 --- a/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py @@ -552,6 +552,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class DeepseekOcr2SamVisionEncoder(DeepseekOcr2PreTrainedModel): _can_record_outputs = {"hidden_states": DeepseekOcr2SamVisionLayer, "attentions": DeepseekOcr2SamVisionAttention} + _input_embed_layer = "patch_embed" def __init__(self, config: DeepseekOcr2SamVisionConfig): super().__init__(config) @@ -585,9 +586,6 @@ def __init__(self, config: DeepseekOcr2SamVisionConfig): self.proj = DeepseekOcr2SamVisionProj(config) self.post_init() - def get_input_embeddings(self): - return self.patch_embed - @merge_with_config_defaults @capture_outputs @auto_docstring diff --git a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py index b5a89bf7335a..d92ddc90609b 100644 --- a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py @@ -747,6 +747,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class DeepseekOcr2SamVisionEncoder(SamVisionEncoder, DeepseekOcr2PreTrainedModel): + _input_embed_layer = "patch_embed" + def __init__(self, config: DeepseekOcr2SamVisionConfig): super().__init__(config) self.proj = DeepseekOcr2SamVisionProj(config) @@ -785,6 +787,9 @@ def forward(self, pixel_values: torch.FloatTensor, **kwargs) -> BaseModelOutput: hidden_states = self.proj(hidden_states) return BaseModelOutput(last_hidden_state=hidden_states) + def get_input_embeddings(self): + raise AttributeError() + class DeepseekOcr2VisionMLP(Qwen2MLP): pass diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 8ba6d23db602..f7580739aa83 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -570,6 +570,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class UnlimitedOcrSamVisionEncoder(UnlimitedOcrPreTrainedModel): _can_record_outputs = {"hidden_states": UnlimitedOcrSamVisionLayer, "attentions": UnlimitedOcrSamVisionAttention} + _input_embed_layer = "patch_embed" def __init__(self, config: UnlimitedOcrSamVisionConfig): super().__init__(config) @@ -603,9 +604,6 @@ def __init__(self, config: UnlimitedOcrSamVisionConfig): self.proj = UnlimitedOcrSamVisionProj(config) self.post_init() - def get_input_embeddings(self): - return self.patch_embed - @merge_with_config_defaults @capture_outputs @auto_docstring From d5c0e01e5e922ddacbe40138b51cdfcae0a28f19 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 16:23:49 +0200 Subject: [PATCH 091/195] update tests --- .../unlimited_ocr/test_modeling_unlimited_ocr.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 8a1486056e5b..b8532044c841 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -178,30 +178,28 @@ def test_small_model_integration_test_document_parsing(self): EXPECTED_DECODED_TEXT = Expectations( { ("cuda", None): "image [383, 87, 497, 171]\ntitle [333", + ("cpu", None): "image [383, 87, 497, 171]\ntitle [333", } ).get_expectation() # fmt: skip self.assertEqual(decoded, EXPECTED_DECODED_TEXT) @slow @require_torch_accelerator - def test_small_model_integration_test_document_parsing_grounding(self): + def test_small_model_integration_test_document_parsing_no_skip_special_tokens(self): model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, device_map=torch_device).eval() image = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" ) ) - inputs = self.processor( - images=image, - text="document parsing.", - return_tensors="pt", - ).to(model.device) + inputs = self.processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=False) EXPECTED_DECODED_TEXT = Expectations( { ("cuda", None): "<|det|>image [383, 87, 497, 171]<|/det|>\n<|det|>title [333", + ("cpu", None): "<|det|>image [383, 87, 497, 171]<|/det|>\n<|det|>title [333", } ).get_expectation() # fmt: skip self.assertEqual(decoded, EXPECTED_DECODED_TEXT) @@ -238,6 +236,10 @@ def test_small_model_integration_test_document_parsing_batched(self): "image [383, 87, 497, 171]\ntitle [333", "header [53, 23, 365, 41]Advanced Template and Styl", ], + ("cpu", None): [ + "image [383, 87, 497, 171]\ntitle [333", + "header [53, 23, 365, 41]Advanced Template and Styl", + ], } ).get_expectation() self.assertEqual(decoded, EXPECTED_DECODED_TEXT) @@ -270,6 +272,7 @@ def test_small_model_integration_test_multi_page_document_parsing(self): EXPECTED_DECODED_TEXT = Expectations( { ("cuda", None): "image [382, 87, 489, 174]\n", + ("cpu", None): "image [382, 87, 489, 174]\n", } ).get_expectation() self.assertEqual(decoded, EXPECTED_DECODED_TEXT) From fd4bb928901a7827ce23283a859ed6cfd6537bb4 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 17:02:22 +0200 Subject: [PATCH 092/195] update docs --- docs/source/en/model_doc/unlimited_ocr.md | 26 +++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 13fe41f690fd..ed1ec6a14809 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -35,7 +35,7 @@ This model was contributed by [guarin](https://huggingface.co/guarin). The original code can be found [here](https://github.com/baidu/Unlimited-OCR). > [!TIP] -> Unlimited-OCR supports two inference configurations: the default "gundam" mode uses 640x640 tiles with dynamic cropping for high-resolution documents, and "base" mode uses a single 1024x1024 global view for standard-resolution inputs. Gundam mode is enabled by default via the image processor (`crop_to_patches=True`). +> Unlimited-OCR supports two inference configurations: the default "gundam" mode uses 640x640 tiles with dynamic cropping for high-resolution documents, and "base" mode uses a single 1024x1024 global view for standard-resolution inputs. To enable base mode set `crop_to_patches=False`. > [!TIP] > For multi-page documents, pass all page images together with one `` token per page in the text prompt. The model processes all pages jointly within a single context window. @@ -52,7 +52,7 @@ from transformers import AutoProcessor, AutoModelForImageTextToText model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", device_map="auto") processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") -image = "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" +image = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" inputs = processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) output = model.generate(**inputs, max_new_tokens=4096) @@ -60,6 +60,23 @@ processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=T # "R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM\nName/Phone Ext. : (...)" ``` +Pass multiple images and prompts at once for batch processing. + +```python +from transformers import AutoProcessor, AutoModelForImageTextToText + +model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", device_map="auto") +processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") + +image1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" +image2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_receipt.jpeg" +inputs = processor(images=[image1, image2], text=["document parsing.", "document parsing."], return_tensors="pt").to(model.device) + +output = model.generate(**inputs, max_new_tokens=4096) +processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) +# "R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM\nName/Phone Ext. : (...)" +``` + Multi-page documents can be parsed jointly in a single forward pass by passing all page images together. Include one `` token per page in the text prompt so the model processes all pages as a continuous document. @@ -71,8 +88,8 @@ from transformers import AutoProcessor, AutoModelForImageTextToText model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", device_map="auto") processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") -page1 = Image.open("page1.png") -page2 = Image.open("page2.png") +page1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" +page2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_receipt.jpeg" num_pages = 2 inputs = processor( @@ -84,6 +101,7 @@ inputs = processor( output = model.generate(**inputs, max_new_tokens=32768) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) +# ``` From 6cd16683cc09620805a089ccc13ae3f156c5da7f Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 17:42:51 +0200 Subject: [PATCH 093/195] update docs and handle image dtype --- docs/source/en/model_doc/unlimited_ocr.md | 21 ++++++++++++++----- .../unlimited_ocr/modeling_unlimited_ocr.py | 2 ++ .../unlimited_ocr/modular_unlimited_ocr.py | 2 ++ 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index ed1ec6a14809..2dbdd832cd53 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -57,10 +57,11 @@ inputs = processor(images=image, text="document parsing.", return_tensors output = model.generate(**inputs, max_new_tokens=4096) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) -# "R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM\nName/Phone Ext. : (...)" +# image [383, 87, 497, 171]\ntext [333, 201, 558, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... ``` -Pass multiple images and prompts at once for batch processing. +For batch processing, pass multiple images and prompts at once. Set `padding=True` for the processor +if images have different sizes. ```python from transformers import AutoProcessor, AutoModelForImageTextToText @@ -70,11 +71,21 @@ processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") image1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" image2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_receipt.jpeg" -inputs = processor(images=[image1, image2], text=["document parsing.", "document parsing."], return_tensors="pt").to(model.device) +inputs = processor( + images=[image1, image2], + text=["document parsing.", "document parsing."], + padding=True, + return_tensors="pt", +).to(model.device) output = model.generate(**inputs, max_new_tokens=4096) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) -# "R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM\nName/Phone Ext. : (...)" +# image [383, 87, 497, 171]\ntext [333, 201, 558, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... +``` + +Set `skip_special_tokens=False` to wrap all detections and region types in `<|det|>...<|/det|>` markers. This if useful for further post-processing of the output which will look like this: +``` +<|det|>image [383, 87, 497, 171]<|/det|>\n<|det|>text [333, 201, 558, 230]<|/det|>R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... ``` @@ -101,7 +112,7 @@ inputs = processor( output = model.generate(**inputs, max_new_tokens=32768) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) -# +# image [382, 87, 489, 174]\ntitle [333, 201, 556, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SCLUTION FORM... ``` diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index f7580739aa83..76f0ef23d30b 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1810,11 +1810,13 @@ def get_image_features( batch_size = pixel_values.shape[0] + pixel_values = pixel_values.to(self.vision_tower.dtype) 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: + pixel_values_local = pixel_values_local.to(self.vision_tower.dtype) 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) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index c1ccd497831b..f292916ad344 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -1028,11 +1028,13 @@ def get_image_features( batch_size = pixel_values.shape[0] + pixel_values = pixel_values.to(self.vision_tower.dtype) 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: + pixel_values_local = pixel_values_local.to(self.vision_tower.dtype) 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) From 55e6f5451552cac563fe83f13eacd2e4fd73e68f Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 7 Jul 2026 17:58:56 +0200 Subject: [PATCH 094/195] fix conversion mapping --- src/transformers/conversion_mapping.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 7bb1df8f5511..cd89f70f4bde 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -1527,6 +1527,7 @@ def _build_checkpoint_conversion_mapping(): "k_proj.weight", "v_proj.weight", ], + operations=[Chunk(dim=0)], ), WeightConverter( source_patterns="qkv_proj.bias", @@ -1535,6 +1536,7 @@ def _build_checkpoint_conversion_mapping(): "k_proj.bias", "v_proj.bias", ], + operations=[Chunk(dim=0)], ), ], "UnlimitedOcrTextModel": [ From ad5cdabff868ac62ab975d81665d7653296591c6 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 08:16:43 +0200 Subject: [PATCH 095/195] fix cache --- src/transformers/cache_utils.py | 6 +++++- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 6 +++--- .../models/unlimited_ocr/modular_unlimited_ocr.py | 6 +++--- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index 1d6dc164cc36..ef3079c02903 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -1492,7 +1492,11 @@ def get_layer_types_and_kwargs(config: PreTrainedConfig) -> tuple[list[str], dic # Prepare additional kwargs that may be needed to __init__ the cache layers layer_kwargs = {} - if "sliding_attention" in layer_types or "hybrid_sliding" in layer_types: + if ( + "sliding_attention" in layer_types + or "hybrid_sliding" in layer_types + or "reference_sliding_attention" in layer_types + ): layer_kwargs["sliding_window"] = config.sliding_window if "chunked_attention" in layer_types: layer_kwargs["sliding_window"] = config.attention_chunk_size diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 76f0ef23d30b..fddf6fb4132f 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1276,7 +1276,7 @@ class UnlimitedOcrDynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): replaced by the most recent ones. The prefill tokens always remain in the cache. """ - layer_type = "reference_sliding_attention" + _layer_type = "reference_sliding_attention" def __init__(self, config: PreTrainedConfig | None = None, sliding_window: int | None = None): super().__init__(config=config, sliding_window=sliding_window) @@ -1389,7 +1389,7 @@ class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): The size of the sliding window. """ - layer_type = "reference_sliding_attention" + _layer_type = "reference_sliding_attention" def __init__(self, max_cache_len: int, sliding_window: int): super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window) @@ -1665,7 +1665,7 @@ def create_reference_sliding_window_causal_mask(**kwargs): prefill_length = float("inf") kv_offset = 0 else: - layer = next(layer for layer in past_key_values.layers if layer.layer_type == "reference_sliding_attention") + layer = next(layer for layer in past_key_values.layers if layer._layer_type == "reference_sliding_attention") prefill_length = float("inf") if layer.prefill_length is None else layer.prefill_length _, kv_offset = layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index f292916ad344..aadcfc25a971 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -595,7 +595,7 @@ class UnlimitedOcrDynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): replaced by the most recent ones. The prefill tokens always remain in the cache. """ - layer_type = "reference_sliding_attention" + _layer_type = "reference_sliding_attention" def __init__(self, config: PreTrainedConfig | None = None, sliding_window: int | None = None): super().__init__(config=config, sliding_window=sliding_window) @@ -708,7 +708,7 @@ class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): The size of the sliding window. """ - layer_type = "reference_sliding_attention" + _layer_type = "reference_sliding_attention" def __init__(self, max_cache_len: int, sliding_window: int): super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window) @@ -919,7 +919,7 @@ def create_reference_sliding_window_causal_mask(**kwargs): prefill_length = float("inf") kv_offset = 0 else: - layer = next(layer for layer in past_key_values.layers if layer.layer_type == "reference_sliding_attention") + layer = next(layer for layer in past_key_values.layers if layer._layer_type == "reference_sliding_attention") prefill_length = float("inf") if layer.prefill_length is None else layer.prefill_length _, kv_offset = layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) From 6e37e672879b4ea70f11908c90b73be7a113d388 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 08:46:11 +0200 Subject: [PATCH 096/195] update docs --- docs/source/en/model_doc/unlimited_ocr.md | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 2dbdd832cd53..b525a17337ef 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -25,24 +25,26 @@ limitations under the License. The UnlimitedOcr model was proposed in [Unlimited OCR Works](https://huggingface.co/papers/2606.23050) by Youyang Yin, Huanhuan Liu, Qunyi Xie, Chaorun Liu, Shiqi Yang, Shaohua Wang, Zhanlong Liu, Hao Zou, Jinyue Chen, Shu Wei, Jingjing Wu, Mingxin Huang, Zhen Wu, Guibin Wang, Tengyu Du, and Lei Jia from Baidu Inc. +It is a single 3B parameter model with a standard context length of 32,768 tokens. + The abstract from the paper is the following: *Recently, end-to-end OCR models, exemplified by DeepSeek OCR, have once again thrust OCR into the spotlight. A widely held view is that employing a large language model (LLM) as the decoder allows the model to leverage the prior distribution of language, leading to improved OCR performance. However, the downside is equally evident: as the output sequence lengthens, the accumulated KV cache drives up memory consumption and progressively slows down generation. This stands in stark contrast to humans, who exhibit no such decline in efficiency during long-horizon copying tasks. In this technical report, we propose Unlimited OCR, a model designed to emulate human parsing working memory. Taking DeepSeek OCR as the baseline, we replace all attention layers in the decoder with our proposed Reference Sliding Window Attention (R-SWA), which reduces attention computation costs while maintaining a constant KV cache throughout the entire decoding process. By combining the high compression rate of DeepSeek OCR's encoder with our constant KV cache design, Unlimited OCR can transcribe dozens of pages of documents in a single forward pass under a standard maximum length of 32K. More importantly, R-SWA is a general-purpose parsing attention mechanism β€” beyond OCR, it is equally applicable to tasks such as ASR, translation, etc. Codes and model weights are publicly available at http://github.com/baidu/Unlimited-OCR* +Unlimited-OCR supports two inference configurations: the default "gundam" mode uses 640x640 tiles with dynamic cropping for high-resolution documents, and "base" mode uses a single 1024x1024 global view for standard-resolution inputs. To enable base mode set `crop_to_patches=False`. + +The vision tower follows the two-stage approach from [DeepSeek-OCR-2](./deepseek_ocr2): a SAM ViT-B encoder feeds into a CLIP ViT encoder. Unlike DeepSeek-OCR-2, the CLIP features are additionally concatenated with the SAM features to yield the final image tokens. Unlimited-OCR also omits the learnable patch queries from DeepSeek-OCR-2. + +The text model is identical to DeepSeek-OCR-2 with the additional Reference Sliding Window Attention (R-SWA). R-SWA applies only to generated tokens. All image and prompt tokens remain fully visible throughout decoding, so long documents do not lose context from earlier pages. + This model was contributed by [guarin](https://huggingface.co/guarin). The original code can be found [here](https://github.com/baidu/Unlimited-OCR). -> [!TIP] -> Unlimited-OCR supports two inference configurations: the default "gundam" mode uses 640x640 tiles with dynamic cropping for high-resolution documents, and "base" mode uses a single 1024x1024 global view for standard-resolution inputs. To enable base mode set `crop_to_patches=False`. - > [!TIP] > For multi-page documents, pass all page images together with one `` token per page in the text prompt. The model processes all pages jointly within a single context window. -> [!TIP] -> The Reference Sliding Window Attention (R-SWA) applies only to generated tokens. All image and prompt tokens from the prefill remain fully visible throughout decoding, so long documents do not lose context from earlier pages. - @@ -118,12 +120,6 @@ processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=T -## Notes - -- [`UnlimitedOcrForConditionalGeneration`] extends [`DeepseekOcr2ForConditionalGeneration`](deepseek_ocr2) with a two-stage vision pipeline: a SAM ViT-B encoder feeds into a CLIP ViT encoder, and their concatenated features are projected to the language model hidden size. The 3B-parameter model supports up to 32,768 context tokens. -- The Reference Sliding Window Attention (R-SWA) in the decoder keeps the KV cache constant throughout decoding. The prefill (image tokens and prompt) is retained in full; the sliding window applies only across generated tokens. -- Image inputs are only forwarded during the first generation step. Subsequent decode steps skip `pixel_values` to avoid reprocessing the image. - ## UnlimitedOcrConfig [[autodoc]] UnlimitedOcrConfig From 667666a3aaa211573f73e50be9a89677bbba0dbb Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 08:59:20 +0200 Subject: [PATCH 097/195] cleanup --- .../unlimited_ocr/configuration_unlimited_ocr.py | 2 +- .../unlimited_ocr/modeling_unlimited_ocr.py | 12 ++++++------ .../models/unlimited_ocr/modular_unlimited_ocr.py | 15 ++++++++------- 3 files changed, 15 insertions(+), 14 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index 8735e377cd2e..6be4bb0a28fe 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -236,7 +236,7 @@ def __post_init__(self, **kwargs): self.sliding_window = self.sliding_window if self.use_sliding_window else None if self.layer_types is None: self.layer_types = [ - "reference_sliding_attention" if self.use_sliding_window else "full_attention" + "reference_sliding_attention" if self.sliding_window is not None else "full_attention" ] * self.num_hidden_layers elif len(set(self.layer_types)) > 1: # This requires a custom create_causal_mask implementation for reference_sliding_attention diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index fddf6fb4132f..dd8a26793cae 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -935,16 +935,16 @@ def __init__(self, config: UnlimitedOcrVisionConfig): @auto_docstring def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: sam_encoder_outputs = self.sam_encoder(pixel_values, **kwargs) - sam_feature_map = sam_encoder_outputs.last_hidden_state + sam_hidden_states = sam_encoder_outputs.last_hidden_state - vision_encoder_outputs = self.vision_encoder(sam_feature_map, **kwargs) - vision_encoder_hidden_state = vision_encoder_outputs.last_hidden_state + vision_encoder_outputs = self.vision_encoder(sam_hidden_states, **kwargs) + hidden_states = vision_encoder_outputs.last_hidden_state - sam_hidden_state = sam_feature_map.flatten(2).transpose(1, 2) - hidden_state = torch.cat([vision_encoder_hidden_state[:, 1:], sam_hidden_state], dim=-1) + sam_hidden_states = sam_hidden_states.flatten(2).transpose(1, 2) + hidden_states = torch.cat([hidden_states[:, 1:], sam_hidden_states], dim=-1) return BaseModelOutput( - last_hidden_state=hidden_state, + last_hidden_state=hidden_states, hidden_states=vision_encoder_outputs.hidden_states, attentions=vision_encoder_outputs.attentions, ) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index aadcfc25a971..9a2b78a6dea0 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -390,9 +390,10 @@ def __post_init__(self, **kwargs): self.sliding_window = self.sliding_window if self.use_sliding_window else None if self.layer_types is None: self.layer_types = [ - "reference_sliding_attention" if self.use_sliding_window else "full_attention" + "reference_sliding_attention" if self.sliding_window is not None else "full_attention" ] * self.num_hidden_layers elif len(set(self.layer_types)) > 1: + #Β TODO: check that this works now # This requires a custom create_causal_mask implementation for reference_sliding_attention # that fetches the first layer with is_sliding=True. raise ValueError( @@ -568,16 +569,16 @@ def __init__(self, config: UnlimitedOcrVisionConfig): @auto_docstring def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: sam_encoder_outputs = self.sam_encoder(pixel_values, **kwargs) - sam_feature_map = sam_encoder_outputs.last_hidden_state + sam_hidden_states = sam_encoder_outputs.last_hidden_state - vision_encoder_outputs = self.vision_encoder(sam_feature_map, **kwargs) - vision_encoder_hidden_state = vision_encoder_outputs.last_hidden_state + vision_encoder_outputs = self.vision_encoder(sam_hidden_states, **kwargs) + hidden_states = vision_encoder_outputs.last_hidden_state - sam_hidden_state = sam_feature_map.flatten(2).transpose(1, 2) - hidden_state = torch.cat([vision_encoder_hidden_state[:, 1:], sam_hidden_state], dim=-1) + sam_hidden_states = sam_hidden_states.flatten(2).transpose(1, 2) + hidden_states = torch.cat([hidden_states[:, 1:], sam_hidden_states], dim=-1) return BaseModelOutput( - last_hidden_state=hidden_state, + last_hidden_state=hidden_states, hidden_states=vision_encoder_outputs.hidden_states, attentions=vision_encoder_outputs.attentions, ) From cb7bc6acdb104eac8e1795d03eb5f28b42dd74dd Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 09:11:53 +0200 Subject: [PATCH 098/195] drop UnlimitedOcrTextModel conversion mapping --- src/transformers/conversion_mapping.py | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index cd89f70f4bde..f16616a5637c 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -127,6 +127,7 @@ "Qwen2_5_VLModel": "Qwen2VLModel", "Qwen2_5_VLForConditionalGeneration": "Qwen2VLForConditionalGeneration", "Tipsv2VisionBackbone": "Tipsv2VisionModel", + "UnlimitedOcrTextModel": "qwen2_moe", # ViT-style vision models (old HuggingFace checkpoint format β†’ new modular format) "ASTModel": "ViTModel", "BeitModel": "ViTModel", @@ -1539,21 +1540,6 @@ def _build_checkpoint_conversion_mapping(): operations=[Chunk(dim=0)], ), ], - "UnlimitedOcrTextModel": [ - WeightConverter( - source_patterns=[ - "mlp.experts.*.gate_proj.weight", - "mlp.experts.*.up_proj.weight", - ], - target_patterns="mlp.experts.gate_up_proj", - operations=[MergeModulelist(dim=0), Concatenate(dim=1)], - ), - WeightConverter( - source_patterns="mlp.experts.*.down_proj.weight", - target_patterns="mlp.experts.down_proj", - operations=[MergeModulelist(dim=0)], - ), - ], "tipsv2": [ WeightRenaming("text_encoder", "text_model"), WeightRenaming("vision_encoder", "vision_model"), From e1cfa1534b80d1c309f5a0da17223370e97f5138 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 09:16:36 +0200 Subject: [PATCH 099/195] allow mixed layer types --- .../models/unlimited_ocr/configuration_unlimited_ocr.py | 6 ------ .../models/unlimited_ocr/modular_unlimited_ocr.py | 7 ------- 2 files changed, 13 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index 6be4bb0a28fe..a3c41ac4b7c0 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -238,12 +238,6 @@ def __post_init__(self, **kwargs): self.layer_types = [ "reference_sliding_attention" if self.sliding_window is not None else "full_attention" ] * self.num_hidden_layers - elif len(set(self.layer_types)) > 1: - # This requires a custom create_causal_mask implementation for reference_sliding_attention - # that fetches the first layer with is_sliding=True. - raise ValueError( - f"Combining multiple layer types ({set(self.layer_types)}) is not supported for this model." - ) if self.mlp_layer_types is None: # Some configs may use `first_k_dense_replace` instead of `layer_types`/`mlp_layer_types` first_k_dense_replace = kwargs.pop("first_k_dense_replace", 1) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 9a2b78a6dea0..a10ad874d9ca 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -392,13 +392,6 @@ def __post_init__(self, **kwargs): self.layer_types = [ "reference_sliding_attention" if self.sliding_window is not None else "full_attention" ] * self.num_hidden_layers - elif len(set(self.layer_types)) > 1: - #Β TODO: check that this works now - # This requires a custom create_causal_mask implementation for reference_sliding_attention - # that fetches the first layer with is_sliding=True. - raise ValueError( - f"Combining multiple layer types ({set(self.layer_types)}) is not supported for this model." - ) if self.mlp_layer_types is None: # Some configs may use `first_k_dense_replace` instead of `layer_types`/`mlp_layer_types` first_k_dense_replace = kwargs.pop("first_k_dense_replace", 1) From 66a63f11f1f1b71b8f7e73904ae775993ee1173e Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 09:52:01 +0200 Subject: [PATCH 100/195] add UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor --- .../unlimited_ocr/generation_unlimited_ocr.py | 58 +++++++++++++++++++ .../unlimited_ocr/modeling_unlimited_ocr.py | 28 +++++++++ .../unlimited_ocr/modular_unlimited_ocr.py | 29 ++++++++++ 3 files changed, 115 insertions(+) create mode 100644 src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py diff --git a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py new file mode 100644 index 000000000000..872748fbbb3a --- /dev/null +++ b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py @@ -0,0 +1,58 @@ +# 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 torch + +from ...generation.logits_process import LogitsProcessor + + +class UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor(LogitsProcessor): + r""" + [`LogitsProcessor`] that blocks n-gram repetitions within a sliding window over the most recently generated + tokens, rather than the full sequence. Aligned with SGLang's `DeepseekOCRNoRepeatNGramLogitProcessor`. + + Args: + no_repeat_ngram_size (`int`): + Size of the n-grams that are not allowed to repeat. + no_repeat_ngram_window_size (`int`): + Number of trailing tokens to search for repeated n-grams. + """ + + def __init__(self, no_repeat_ngram_size: int, no_repeat_ngram_window_size: int): + self.no_repeat_ngram_size = no_repeat_ngram_size + self.no_repeat_ngram_window_size = no_repeat_ngram_window_size + + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + for batch_idx in range(input_ids.shape[0]): + sequence = input_ids[batch_idx].tolist() + if len(sequence) < self.no_repeat_ngram_size: + continue + search_start = max(0, len(sequence) - self.no_repeat_ngram_window_size) + search_end = len(sequence) - self.no_repeat_ngram_size + 1 + if search_end <= search_start: + continue + if self.no_repeat_ngram_size > 1: + current_prefix = tuple(sequence[-(self.no_repeat_ngram_size - 1) :]) + else: + current_prefix = () + banned = set() + for idx in range(search_start, search_end): + ngram = sequence[idx : idx + self.no_repeat_ngram_size] + if self.no_repeat_ngram_size == 1 or tuple(ngram[:-1]) == current_prefix: + banned.add(ngram[-1]) + for token_id in banned: + scores[batch_idx, token_id] = float("-inf") + return scores + + +__all__ = ["UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor"] diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index dd8a26793cae..b2f558ae6818 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -63,6 +63,7 @@ UnlimitedOcrVisionConfig, UnlimitedOcrVisionEncoderConfig, ) +from .generation_unlimited_ocr import UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor logger = logging.get_logger(__name__) @@ -2080,6 +2081,33 @@ def prepare_inputs_for_generation( return model_inputs + def _get_logits_processor(self, generation_config, logits_processor=None, **kwargs): + no_repeat_ngram_size = generation_config.no_repeat_ngram_size + no_repeat_ngram_window_size = getattr(generation_config, "no_repeat_ngram_window_size", None) + + use_sliding_window_processor = ( + no_repeat_ngram_window_size is not None and no_repeat_ngram_size is not None and no_repeat_ngram_size > 0 + ) + if use_sliding_window_processor: + logits_processor = list(logits_processor or []) + [ + UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor( + no_repeat_ngram_size=no_repeat_ngram_size, + no_repeat_ngram_window_size=no_repeat_ngram_window_size, + ) + ] + + # Set to None to avoid adding the default NoRepeatNgramLogitsProcessor + generation_config.no_repeat_ngram_size = None + + try: + processors = super()._get_logits_processor( + generation_config=generation_config, logits_processor=logits_processor, **kwargs + ) + finally: + if use_sliding_window_processor: + generation_config.no_repeat_ngram_size = no_repeat_ngram_size + return processors + __all__ = [ "UnlimitedOcrForConditionalGeneration", diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index a10ad874d9ca..7fbd19673274 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -28,6 +28,7 @@ ) from ...configuration_utils import PreTrainedConfig from ...feature_extraction_utils import BatchFeature +from ...generation import GenerationMixin from ...image_transforms import group_images_by_shape, reorder_images from ...image_utils import PILImageResampling, SizeDict from ...masking_utils import and_masks, causal_mask_function, create_causal_mask, create_sliding_window_causal_mask @@ -70,6 +71,7 @@ ) from ..deepseek_ocr2.processing_deepseek_ocr2 import DeepseekOcr2Processor, DeepseekOcr2ProcessorKwargs from ..got_ocr2.configuration_got_ocr2 import GotOcr2VisionConfig +from .generation_unlimited_ocr import UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor class UnlimitedOcrImageProcessorKwargs(DeepseekOcr2ImageProcessorKwargs): @@ -1198,6 +1200,33 @@ def forward( image_hidden_states=outputs.image_hidden_states, ) + def _get_logits_processor(self, generation_config, logits_processor=None, **kwargs): + no_repeat_ngram_size = generation_config.no_repeat_ngram_size + no_repeat_ngram_window_size = getattr(generation_config, "no_repeat_ngram_window_size", None) + + use_sliding_window_processor = ( + no_repeat_ngram_window_size is not None and no_repeat_ngram_size is not None and no_repeat_ngram_size > 0 + ) + if use_sliding_window_processor: + logits_processor = list(logits_processor or []) + [ + UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor( + no_repeat_ngram_size=no_repeat_ngram_size, + no_repeat_ngram_window_size=no_repeat_ngram_window_size, + ) + ] + + # Set to None to avoid adding the default NoRepeatNgramLogitsProcessor + generation_config.no_repeat_ngram_size = None + + try: + processors = GenerationMixin._get_logits_processor( + self, generation_config=generation_config, logits_processor=logits_processor, **kwargs + ) + finally: + if use_sliding_window_processor: + generation_config.no_repeat_ngram_size = no_repeat_ngram_size + return processors + def prepare_inputs_for_generation( self, input_ids, From 0fdbf3cda4a6d267e83ffe1dc2ff05516d7668fc Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 10:09:13 +0200 Subject: [PATCH 101/195] add UnlimitedOcrGenerationConfig --- .../unlimited_ocr/generation_unlimited_ocr.py | 18 +++++++++++++++++- .../unlimited_ocr/modeling_unlimited_ocr.py | 6 +++++- .../unlimited_ocr/modular_unlimited_ocr.py | 7 ++++++- 3 files changed, 28 insertions(+), 3 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py index 872748fbbb3a..898d86a89170 100644 --- a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py @@ -13,9 +13,25 @@ # limitations under the License. import torch +from ...generation.configuration_utils import GenerationConfig from ...generation.logits_process import LogitsProcessor +class UnlimitedOcrGenerationConfig(GenerationConfig): + r""" + A GenerationConfig class with parameterization customized for UnlimitedOcr. + + Args: + no_repeat_ngram_window_size (`int`, *optional*): + If set together with `no_repeat_ngram_size`, n-gram repetitions are blocked only within this many + trailing tokens instead of over the whole sequence. + """ + + def __init__(self, no_repeat_ngram_window_size: int | None = None, **kwargs): + super().__init__(**kwargs) + self.no_repeat_ngram_window_size = no_repeat_ngram_window_size + + class UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor(LogitsProcessor): r""" [`LogitsProcessor`] that blocks n-gram repetitions within a sliding window over the most recently generated @@ -55,4 +71,4 @@ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> to return scores -__all__ = ["UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor"] +__all__ = ["UnlimitedOcrGenerationConfig", "UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor"] diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index b2f558ae6818..2b10c5dbcac8 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -63,7 +63,10 @@ UnlimitedOcrVisionConfig, UnlimitedOcrVisionEncoderConfig, ) -from .generation_unlimited_ocr import UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor +from .generation_unlimited_ocr import ( + UnlimitedOcrGenerationConfig, + UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor, +) logger = logging.get_logger(__name__) @@ -1949,6 +1952,7 @@ def forward( @auto_docstring class UnlimitedOcrForConditionalGeneration(UnlimitedOcrPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} + generation_config_class = UnlimitedOcrGenerationConfig def __init__(self, config: UnlimitedOcrConfig): super().__init__(config) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 7fbd19673274..f4c487040004 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -71,7 +71,10 @@ ) from ..deepseek_ocr2.processing_deepseek_ocr2 import DeepseekOcr2Processor, DeepseekOcr2ProcessorKwargs from ..got_ocr2.configuration_got_ocr2 import GotOcr2VisionConfig -from .generation_unlimited_ocr import UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor +from .generation_unlimited_ocr import ( + UnlimitedOcrGenerationConfig, + UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor, +) class UnlimitedOcrImageProcessorKwargs(DeepseekOcr2ImageProcessorKwargs): @@ -1137,6 +1140,8 @@ def forward( class UnlimitedOcrForConditionalGeneration(DeepseekOcr2ForConditionalGeneration): + generation_config_class = UnlimitedOcrGenerationConfig + @can_return_tuple @auto_docstring def forward( From 17966e54e15e51c3142d6fe6a0be66305c24cded Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 08:42:33 +0000 Subject: [PATCH 102/195] allow custom generation config --- src/transformers/generation/configuration_utils.py | 2 +- src/transformers/generation/utils.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/transformers/generation/configuration_utils.py b/src/transformers/generation/configuration_utils.py index 641201ef384f..c671fb7fd885 100644 --- a/src/transformers/generation/configuration_utils.py +++ b/src/transformers/generation/configuration_utils.py @@ -1267,7 +1267,7 @@ def from_model_config(cls, model_config: Union["PreTrainedConfig", dict]) -> "Ge model_config = model_config.get_text_config(decoder=True) model_config = model_config.to_dict() - default_generation_config = GenerationConfig() + default_generation_config = cls() for attr in generation_config.to_dict(): is_unset = getattr(generation_config, attr) == getattr(default_generation_config, attr) if attr in model_config and is_unset: diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index fbbff776713f..daa53b964fb1 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -399,7 +399,7 @@ def adjust_generation_fn( } # Load generation config try: - self.generation_config = GenerationConfig.from_pretrained( + self.generation_config = self.generation_config_class.from_pretrained( pretrained_model_name_or_path, _from_auto=from_auto_class, _from_pipeline=from_pipeline, @@ -412,7 +412,7 @@ def adjust_generation_fn( logger.info( "Generation config file not found, using a generation config created from the model config." ) - self.generation_config = GenerationConfig.from_pretrained( + self.generation_config = self.generation_config_class.from_pretrained( pretrained_model_name_or_path, config_file_name="config.json", _from_auto=from_auto_class, From e38892311b912076c7c47f70c6a1279985932280 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 11:11:19 +0200 Subject: [PATCH 103/195] refactor UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor --- .../unlimited_ocr/generation_unlimited_ocr.py | 49 ++++++------------- .../unlimited_ocr/modeling_unlimited_ocr.py | 4 +- .../unlimited_ocr/modular_unlimited_ocr.py | 4 +- 3 files changed, 20 insertions(+), 37 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py index 898d86a89170..af1a5cbc62e9 100644 --- a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py @@ -14,12 +14,12 @@ import torch from ...generation.configuration_utils import GenerationConfig -from ...generation.logits_process import LogitsProcessor +from ...generation.logits_process import LOGITS_PROCESSOR_INPUTS_DOCSTRING, NoRepeatNGramLogitsProcessor +from ...utils import add_start_docstrings class UnlimitedOcrGenerationConfig(GenerationConfig): - r""" - A GenerationConfig class with parameterization customized for UnlimitedOcr. + r"""A GenerationConfig class with parameterization customized for UnlimitedOcr. Args: no_repeat_ngram_window_size (`int`, *optional*): @@ -32,43 +32,26 @@ def __init__(self, no_repeat_ngram_window_size: int | None = None, **kwargs): self.no_repeat_ngram_window_size = no_repeat_ngram_window_size -class UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor(LogitsProcessor): - r""" - [`LogitsProcessor`] that blocks n-gram repetitions within a sliding window over the most recently generated - tokens, rather than the full sequence. Aligned with SGLang's `DeepseekOCRNoRepeatNGramLogitProcessor`. +class UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor(NoRepeatNGramLogitsProcessor): + r"""Identical to [`NoRepeatNGramLogitsProcessor`] but blocks n-gram repetitions only within the last + `window_size` generated tokens, rather than the full sequence. Args: - no_repeat_ngram_size (`int`): - Size of the n-grams that are not allowed to repeat. - no_repeat_ngram_window_size (`int`): + ngram_size (`int`): + All ngrams of size `ngram_size` can only occur once in `window_size`. + window_size (`int`): Number of trailing tokens to search for repeated n-grams. """ - def __init__(self, no_repeat_ngram_size: int, no_repeat_ngram_window_size: int): - self.no_repeat_ngram_size = no_repeat_ngram_size - self.no_repeat_ngram_window_size = no_repeat_ngram_window_size + def __init__(self, ngram_size: int, window_size: int): + super().__init__(ngram_size=ngram_size) + if not isinstance(window_size, int) or window_size <= 0: + raise ValueError(f"`window_size` has to be a strictly positive integer, but is {window_size}") + self.window_size = window_size + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: - for batch_idx in range(input_ids.shape[0]): - sequence = input_ids[batch_idx].tolist() - if len(sequence) < self.no_repeat_ngram_size: - continue - search_start = max(0, len(sequence) - self.no_repeat_ngram_window_size) - search_end = len(sequence) - self.no_repeat_ngram_size + 1 - if search_end <= search_start: - continue - if self.no_repeat_ngram_size > 1: - current_prefix = tuple(sequence[-(self.no_repeat_ngram_size - 1) :]) - else: - current_prefix = () - banned = set() - for idx in range(search_start, search_end): - ngram = sequence[idx : idx + self.no_repeat_ngram_size] - if self.no_repeat_ngram_size == 1 or tuple(ngram[:-1]) == current_prefix: - banned.add(ngram[-1]) - for token_id in banned: - scores[batch_idx, token_id] = float("-inf") - return scores + return super().__call__(input_ids[:, -self.window_size :], scores) __all__ = ["UnlimitedOcrGenerationConfig", "UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor"] diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 2b10c5dbcac8..7caeb977a160 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -2095,8 +2095,8 @@ def _get_logits_processor(self, generation_config, logits_processor=None, **kwar if use_sliding_window_processor: logits_processor = list(logits_processor or []) + [ UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor( - no_repeat_ngram_size=no_repeat_ngram_size, - no_repeat_ngram_window_size=no_repeat_ngram_window_size, + ngram_size=no_repeat_ngram_size, + window_size=no_repeat_ngram_window_size, ) ] diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index f4c487040004..0c1d915e88a1 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -1215,8 +1215,8 @@ def _get_logits_processor(self, generation_config, logits_processor=None, **kwar if use_sliding_window_processor: logits_processor = list(logits_processor or []) + [ UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor( - no_repeat_ngram_size=no_repeat_ngram_size, - no_repeat_ngram_window_size=no_repeat_ngram_window_size, + ngram_size=no_repeat_ngram_size, + window_size=no_repeat_ngram_window_size, ) ] From 301622b6cc7ff73b6451b00a6b05bff0d6ff6a86 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 11:11:19 +0200 Subject: [PATCH 104/195] refactor UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor --- .../unlimited_ocr/generation_unlimited_ocr.py | 45 +++++++------------ .../unlimited_ocr/modeling_unlimited_ocr.py | 4 +- .../unlimited_ocr/modular_unlimited_ocr.py | 4 +- 3 files changed, 19 insertions(+), 34 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py index 898d86a89170..aaf8dbd9b534 100644 --- a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py @@ -14,7 +14,8 @@ import torch from ...generation.configuration_utils import GenerationConfig -from ...generation.logits_process import LogitsProcessor +from ...generation.logits_process import LOGITS_PROCESSOR_INPUTS_DOCSTRING, NoRepeatNGramLogitsProcessor +from ...utils import add_start_docstrings class UnlimitedOcrGenerationConfig(GenerationConfig): @@ -32,43 +33,27 @@ def __init__(self, no_repeat_ngram_window_size: int | None = None, **kwargs): self.no_repeat_ngram_window_size = no_repeat_ngram_window_size -class UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor(LogitsProcessor): +class UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor(NoRepeatNGramLogitsProcessor): r""" - [`LogitsProcessor`] that blocks n-gram repetitions within a sliding window over the most recently generated - tokens, rather than the full sequence. Aligned with SGLang's `DeepseekOCRNoRepeatNGramLogitProcessor`. + Identical to [`NoRepeatNGramLogitsProcessor`] but blocks n-gram repetitions only within the last + `window_size` generated tokens, rather than the full sequence. Args: - no_repeat_ngram_size (`int`): - Size of the n-grams that are not allowed to repeat. - no_repeat_ngram_window_size (`int`): + ngram_size (`int`): + All ngrams of size `ngram_size` can only occur once in `window_size`. + window_size (`int`): Number of trailing tokens to search for repeated n-grams. """ - def __init__(self, no_repeat_ngram_size: int, no_repeat_ngram_window_size: int): - self.no_repeat_ngram_size = no_repeat_ngram_size - self.no_repeat_ngram_window_size = no_repeat_ngram_window_size + def __init__(self, ngram_size: int, window_size: int): + super().__init__(ngram_size=ngram_size) + if not isinstance(window_size, int) or window_size <= 0: + raise ValueError(f"`window_size` has to be a strictly positive integer, but is {window_size}") + self.window_size = window_size + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: - for batch_idx in range(input_ids.shape[0]): - sequence = input_ids[batch_idx].tolist() - if len(sequence) < self.no_repeat_ngram_size: - continue - search_start = max(0, len(sequence) - self.no_repeat_ngram_window_size) - search_end = len(sequence) - self.no_repeat_ngram_size + 1 - if search_end <= search_start: - continue - if self.no_repeat_ngram_size > 1: - current_prefix = tuple(sequence[-(self.no_repeat_ngram_size - 1) :]) - else: - current_prefix = () - banned = set() - for idx in range(search_start, search_end): - ngram = sequence[idx : idx + self.no_repeat_ngram_size] - if self.no_repeat_ngram_size == 1 or tuple(ngram[:-1]) == current_prefix: - banned.add(ngram[-1]) - for token_id in banned: - scores[batch_idx, token_id] = float("-inf") - return scores + return super().__call__(input_ids[:, -self.window_size :], scores) __all__ = ["UnlimitedOcrGenerationConfig", "UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor"] diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 2b10c5dbcac8..7caeb977a160 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -2095,8 +2095,8 @@ def _get_logits_processor(self, generation_config, logits_processor=None, **kwar if use_sliding_window_processor: logits_processor = list(logits_processor or []) + [ UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor( - no_repeat_ngram_size=no_repeat_ngram_size, - no_repeat_ngram_window_size=no_repeat_ngram_window_size, + ngram_size=no_repeat_ngram_size, + window_size=no_repeat_ngram_window_size, ) ] diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index f4c487040004..0c1d915e88a1 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -1215,8 +1215,8 @@ def _get_logits_processor(self, generation_config, logits_processor=None, **kwar if use_sliding_window_processor: logits_processor = list(logits_processor or []) + [ UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor( - no_repeat_ngram_size=no_repeat_ngram_size, - no_repeat_ngram_window_size=no_repeat_ngram_window_size, + ngram_size=no_repeat_ngram_size, + window_size=no_repeat_ngram_window_size, ) ] From 6b1fab5f8dcbd68f19ea3d2d6b4d00ea1977e395 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 11:44:48 +0200 Subject: [PATCH 105/195] update config docstrings --- .../configuration_unlimited_ocr.py | 53 +++++++++++++---- .../unlimited_ocr/modular_unlimited_ocr.py | 59 +++++++++++++++---- 2 files changed, 89 insertions(+), 23 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index a3c41ac4b7c0..166809f4b540 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -80,10 +80,11 @@ class UnlimitedOcrVisionEncoderConfig(PreTrainedConfig): Example: ```python - >>> from transformers import UnlimitedOcrConfig + >>> from transformers import UnlimitedOcrVisionEncoderConfig, UnlimitedOcrVisionEncoder - >>> config = UnlimitedOcrConfig() - >>> encoder_config = config.vision_config.encoder_config + >>> configuration = UnlimitedOcrVisionEncoderConfig() + >>> model = UnlimitedOcrVisionEncoder(configuration) + >>> configuration = model.config ```""" model_type = "unlimited_ocr_vision_encoder" @@ -117,8 +118,17 @@ class UnlimitedOcrVisionConfig(PreTrainedConfig): sam_config (`dict` or `UnlimitedOcrSamVisionConfig`, *optional*): Configuration for the SAM vision encoder. Defaults to `UnlimitedOcrSamVisionConfig()`. encoder_config (`dict` or `UnlimitedOcrVisionEncoderConfig`, *optional*): - Configuration for the DeepSeek-OCR-2 vision encoder. Defaults to `UnlimitedOcrVisionEncoderConfig()`. - """ + Configuration for the vision encoder. Defaults to `UnlimitedOcrVisionEncoderConfig()`. + + Example: + + ```python + >>> from transformers import UnlimitedOcrVisionConfig, UnlimitedOcrVisionModel + + >>> configuration = UnlimitedOcrVisionConfig() + >>> model = UnlimitedOcrVisionModel(configuration) + >>> configuration = model.config + ```""" model_type = "unlimited_ocr_vision" base_config_key = "vision_config" @@ -155,14 +165,22 @@ class UnlimitedOcrTextConfig(PreTrainedConfig): mlp_layer_types (`list[str]`, *optional*): MLP type (`"dense"` or `"sparse"`) for each decoder layer, e.g. `["dense", "sparse", "sparse", ...]`. layer_types (`list[str]`, *optional*): - Attention type for each decoder layer. Defaults to `"full_attention"` on every layer so the KV cache - retains all tokens; the sliding window (`sliding_window`) is applied as a mask over generated tokens - only, not by truncating the cache. + Attention type for each decoder layer. Defaults to `"reference_sliding_attention"` on every layer. + use_sliding_window (`int`, defaults to `True`): + Whether to use reference sliding window attention. sliding_window (`int`, *optional*, defaults to 128): - If set, each token additionally attends only to the last `sliding_window` tokens. The image and prompt - tokens processed during prefill stay fully visible (they are never evicted); the window only applies - across generated tokens. Set to `None` for full causal attention. - """ + Sliding window size for reference sliding window attention. If set, every token attends to the last + `sliding_window` and all image and prompt tokens. + + Example: + + ```python + >>> from transformers import UnlimitedOcrTextConfig, UnlimitedOcrTextModel + + >>> configuration = UnlimitedOcrTextConfig() + >>> model = UnlimitedOcrTextModel(configuration) + >>> configuration = model.config + ```""" model_type = "unlimited_ocr_text" keys_to_ignore_at_inference = ["past_key_values"] @@ -265,7 +283,16 @@ class UnlimitedOcrConfig(PreTrainedConfig): r""" vision_config (`dict` or `UnlimitedOcrVisionConfig`, *optional*): Configuration for the vision encoders. Defaults to `UnlimitedOcrVisionConfig()`. - """ + + Example: + + ```python + >>> from transformers import UnlimitedOcrConfig, UnlimitedOcrModel + + >>> configuration = UnlimitedOcrConfig() + >>> model = UnlimitedOcrModel(configuration) + >>> configuration = model.config + ```""" model_type = "unlimited_ocr" sub_configs = { diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 0c1d915e88a1..c3e96984bf39 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -312,10 +312,11 @@ class UnlimitedOcrVisionEncoderConfig(CLIPVisionConfig): Example: ```python - >>> from transformers import UnlimitedOcrConfig + >>> from transformers import UnlimitedOcrVisionEncoderConfig, UnlimitedOcrVisionEncoder - >>> config = UnlimitedOcrConfig() - >>> encoder_config = config.vision_config.encoder_config + >>> configuration = UnlimitedOcrVisionEncoderConfig() + >>> model = UnlimitedOcrVisionEncoder(configuration) + >>> configuration = model.config ```""" model_type = "unlimited_ocr_vision_encoder" @@ -331,6 +332,22 @@ class UnlimitedOcrVisionEncoderConfig(CLIPVisionConfig): @auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict class UnlimitedOcrVisionConfig(DeepseekOcr2VisionConfig): + r""" + sam_config (`dict` or `UnlimitedOcrSamVisionConfig`, *optional*): + Configuration for the SAM vision encoder. Defaults to `UnlimitedOcrSamVisionConfig()`. + encoder_config (`dict` or `UnlimitedOcrVisionEncoderConfig`, *optional*): + Configuration for the vision encoder. Defaults to `UnlimitedOcrVisionEncoderConfig()`. + + Example: + + ```python + >>> from transformers import UnlimitedOcrVisionConfig, UnlimitedOcrVisionModel + + >>> configuration = UnlimitedOcrVisionConfig() + >>> model = UnlimitedOcrVisionModel(configuration) + >>> configuration = model.config + ```""" + model_type = "unlimited_ocr_vision" base_config_key = "vision_config" sub_configs = { @@ -363,14 +380,22 @@ class UnlimitedOcrTextConfig(DeepseekOcr2TextConfig): mlp_layer_types (`list[str]`, *optional*): MLP type (`"dense"` or `"sparse"`) for each decoder layer, e.g. `["dense", "sparse", "sparse", ...]`. layer_types (`list[str]`, *optional*): - Attention type for each decoder layer. Defaults to `"full_attention"` on every layer so the KV cache - retains all tokens; the sliding window (`sliding_window`) is applied as a mask over generated tokens - only, not by truncating the cache. + Attention type for each decoder layer. Defaults to `"reference_sliding_attention"` on every layer. + use_sliding_window (`int`, defaults to `True`): + Whether to use reference sliding window attention. sliding_window (`int`, *optional*, defaults to 128): - If set, each token additionally attends only to the last `sliding_window` tokens. The image and prompt - tokens processed during prefill stay fully visible (they are never evicted); the window only applies - across generated tokens. Set to `None` for full causal attention. - """ + Sliding window size for reference sliding window attention. If set, every token attends to the last + `sliding_window` and all image and prompt tokens. + + Example: + + ```python + >>> from transformers import UnlimitedOcrTextConfig, UnlimitedOcrTextModel + + >>> configuration = UnlimitedOcrTextConfig() + >>> model = UnlimitedOcrTextModel(configuration) + >>> configuration = model.config + ```""" model_type = "unlimited_ocr_text" base_config_key = "text_config" @@ -410,6 +435,20 @@ def __post_init__(self, **kwargs): @auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict class UnlimitedOcrConfig(DeepseekOcr2Config): + r""" + vision_config (`dict` or `UnlimitedOcrVisionConfig`, *optional*): + Configuration for the vision encoders. Defaults to `UnlimitedOcrVisionConfig()`. + + Example: + + ```python + >>> from transformers import UnlimitedOcrConfig, UnlimitedOcrModel + + >>> configuration = UnlimitedOcrConfig() + >>> model = UnlimitedOcrModel(configuration) + >>> configuration = model.config + ```""" + model_type = "unlimited_ocr" sub_configs = { "vision_config": UnlimitedOcrVisionConfig, From ac7e12c3a38a10156aa58805f5e76073378f99d9 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 11:59:43 +0200 Subject: [PATCH 106/195] cleanup UnlimitedOcrVisionEmbeddings --- .../unlimited_ocr/modeling_unlimited_ocr.py | 14 ++------------ .../unlimited_ocr/modular_unlimited_ocr.py | 18 ++++++++++++++---- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 7caeb977a160..e999b5e73218 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -312,7 +312,6 @@ def _init_weights(self, module): elif isinstance(module, UnlimitedOcrVisionEmbeddings): factor = module.config.initializer_factor init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor) - init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor) init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor) init.copy_(module.position_ids, torch.arange(module.num_positions).expand((1, -1))) @@ -782,14 +781,6 @@ def __init__(self, config: UnlimitedOcrVisionConfig): self.class_embedding = nn.Parameter(torch.randn(self.embed_dim)) - self.patch_embedding = nn.Conv2d( - in_channels=config.num_channels, - out_channels=self.embed_dim, - kernel_size=self.patch_size, - stride=self.patch_size, - bias=False, - ) - self.num_patches = (self.image_size // self.patch_size) ** 2 self.num_positions = self.num_patches + 1 self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) @@ -824,7 +815,6 @@ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) - # TODO: check if we can drop dtype cast target_dtype = patch_pos_embed.dtype patch_pos_embed = nn.functional.interpolate( patch_pos_embed.to(torch.float32), @@ -841,8 +831,7 @@ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: r""" patch_embeds (`torch.Tensor` of shape `(batch_size, hidden_size, grid_height, grid_width)`): - The SAM feature map, injected directly as the CLIP patch embeddings. The CLIP patch convolution - (`self.patch_embedding`) is intentionally bypassed. + The SAM feature map, injected directly as patch embeddings. """ batch_size, _, grid_height, grid_width = patch_embeds.shape patch_embeds = patch_embeds.flatten(2).transpose(1, 2) @@ -903,6 +892,7 @@ class UnlimitedOcrVisionEncoder(UnlimitedOcrPreTrainedModel): "hidden_states": UnlimitedOcrEncoderLayer, "attentions": UnlimitedOcrAttention, } + _keys_to_ignore_on_load_unexpected = {"patch_embedding"} # unused def __init__(self, config: UnlimitedOcrVisionEncoderConfig): super().__init__(config) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index c3e96984bf39..42811a052c7f 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -499,7 +499,6 @@ def _init_weights(self, module): elif isinstance(module, UnlimitedOcrVisionEmbeddings): factor = module.config.initializer_factor init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor) - init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor) init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor) init.copy_(module.position_ids, torch.arange(module.num_positions).expand((1, -1))) @@ -520,7 +519,19 @@ class UnlimitedOcrEncoderLayer(CLIPEncoderLayer): class UnlimitedOcrVisionEmbeddings(CLIPVisionEmbeddings): + def __init__(self, config: UnlimitedOcrVisionConfig): + super().__init__(config) + del self.patch_embedding + def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: + """ + This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution + images. This method is also adapted to support torch.jit tracing. + + Adapted from: + - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and + - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211 + """ num_patches = embeddings.shape[1] - 1 position_embedding = self.position_embedding.weight.unsqueeze(0) num_positions = position_embedding.shape[1] - 1 @@ -541,7 +552,6 @@ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) - # TODO: check if we can drop dtype cast target_dtype = patch_pos_embed.dtype patch_pos_embed = nn.functional.interpolate( patch_pos_embed.to(torch.float32), @@ -558,8 +568,7 @@ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: r""" patch_embeds (`torch.Tensor` of shape `(batch_size, hidden_size, grid_height, grid_width)`): - The SAM feature map, injected directly as the CLIP patch embeddings. The CLIP patch convolution - (`self.patch_embedding`) is intentionally bypassed. + The SAM feature map, injected directly as patch embeddings. """ batch_size, _, grid_height, grid_width = patch_embeds.shape patch_embeds = patch_embeds.flatten(2).transpose(1, 2) @@ -578,6 +587,7 @@ class UnlimitedOcrVisionEncoder(CLIPVisionModel): "hidden_states": UnlimitedOcrEncoderLayer, "attentions": UnlimitedOcrAttention, } + _keys_to_ignore_on_load_unexpected = {"patch_embedding"} # unused def __init__(self, config: UnlimitedOcrVisionEncoderConfig): super().__init__(config) From 08eeed47d3f9e73e24d777f2f37b3f8693404dc2 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:16:01 +0200 Subject: [PATCH 107/195] cleanup UnlimitedOcrVisionModel --- .../unlimited_ocr/modeling_unlimited_ocr.py | 27 +++++++++++++++-- .../unlimited_ocr/modular_unlimited_ocr.py | 30 +++++++++++++++++++ 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index e999b5e73218..4fd672479d14 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -880,7 +880,7 @@ def forward( @auto_docstring( custom_intro=""" - The vision model from UNLIMITED_OCR without any head or projection on top. + The vision encoder from Unlimited OCR without any head or projection on top. """ ) class UnlimitedOcrVisionEncoder(UnlimitedOcrPreTrainedModel): @@ -917,7 +917,7 @@ def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwarg class UnlimitedOcrVisionModel(UnlimitedOcrPreTrainedModel): - """Vision pipeline: SAM ViT-B (with neck)""" + """Vision model encoding images first with SAM followed by an additional (CLIP) model.""" def __init__(self, config: UnlimitedOcrVisionConfig): super().__init__(config) @@ -928,6 +928,29 @@ def __init__(self, config: UnlimitedOcrVisionConfig): @can_return_tuple @auto_docstring def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: + r""" + Example: + + ```python + >>> import torch + >>> from transformers import AutoConfig, AutoImageProcessor, AutoModel + >>> from transformers.image_utils import load_image + + >>> model_id = "baidu/Unlimited-OCR" + >>> config = AutoConfig.from_pretrained(model_id) + >>> model = AutoModel.from_pretrained(model_id, config=config.vision_config, device_map="auto") + >>> image_processor = AutoImageProcessor.from_pretrained(model_id) + + >>> image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg") + >>> inputs = image_processor(images=image, return_tensors="pt").to(model.device) + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> # Patch tokens without class tokens (batch_size, height * width, vision_hidden_size + sam_hidden_size) + >>> patch_tokens = outputs.last_hidden_state + ```""" + sam_encoder_outputs = self.sam_encoder(pixel_values, **kwargs) sam_hidden_states = sam_encoder_outputs.last_hidden_state diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 42811a052c7f..e43a586dac46 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -581,6 +581,11 @@ def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: return embeddings +@auto_docstring( + custom_intro=""" + The vision encoder from Unlimited OCR without any head or projection on top. + """ +) class UnlimitedOcrVisionEncoder(CLIPVisionModel): main_input_name = "patch_embeds" _can_record_outputs = { @@ -607,6 +612,8 @@ def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwarg class UnlimitedOcrVisionModel(DeepseekOcr2VisionModel): + """Vision model encoding images first with SAM followed by an additional (CLIP) model.""" + def __init__(self, config: UnlimitedOcrVisionConfig): super().__init__(config) del self.query_768_resolution @@ -615,6 +622,29 @@ def __init__(self, config: UnlimitedOcrVisionConfig): @can_return_tuple @auto_docstring def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: + r""" + Example: + + ```python + >>> import torch + >>> from transformers import AutoConfig, AutoImageProcessor, AutoModel + >>> from transformers.image_utils import load_image + + >>> model_id = "baidu/Unlimited-OCR" + >>> config = AutoConfig.from_pretrained(model_id) + >>> model = AutoModel.from_pretrained(model_id, config=config.vision_config, device_map="auto") + >>> image_processor = AutoImageProcessor.from_pretrained(model_id) + + >>> image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg") + >>> inputs = image_processor(images=image, return_tensors="pt").to(model.device) + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> # Patch tokens without class tokens (batch_size, height * width, vision_hidden_size + sam_hidden_size) + >>> patch_tokens = outputs.last_hidden_state + ```""" + sam_encoder_outputs = self.sam_encoder(pixel_values, **kwargs) sam_hidden_states = sam_encoder_outputs.last_hidden_state From 0acc0cee3ca5434557e0886b6afadeda59ddccd4 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:32:32 +0200 Subject: [PATCH 108/195] update docs and docstrings --- docs/source/en/model_doc/unlimited_ocr.md | 21 +++- .../unlimited_ocr/modeling_unlimited_ocr.py | 51 ++++++++- .../unlimited_ocr/modular_unlimited_ocr.py | 105 +++++++++++++----- 3 files changed, 145 insertions(+), 32 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index b525a17337ef..a57a5e00a15d 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -57,7 +57,12 @@ processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") image = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" inputs = processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) -output = model.generate(**inputs, max_new_tokens=4096) +output = model.generate( + **inputs, + max_new_tokens=4096, + no_repeat_ngram_size=35, + no_repeat_ngram_window_size=128, +) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) # image [383, 87, 497, 171]\ntext [333, 201, 558, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... ``` @@ -80,7 +85,12 @@ inputs = processor( return_tensors="pt", ).to(model.device) -output = model.generate(**inputs, max_new_tokens=4096) +output = model.generate( + **inputs, + max_new_tokens=4096, + no_repeat_ngram_size=35, + no_repeat_ngram_window_size=128, +) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) # image [383, 87, 497, 171]\ntext [333, 201, 558, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... ``` @@ -112,7 +122,12 @@ inputs = processor( return_tensors="pt", ).to(model.device) -output = model.generate(**inputs, max_new_tokens=32768) +output = model.generate( + **inputs, + max_new_tokens=32768, + no_repeat_ngram_size=35, + no_repeat_ngram_window_size=1024, +) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) # image [382, 87, 489, 174]\ntitle [333, 201, 556, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SCLUTION FORM... ``` diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 4fd672479d14..0a985cc87624 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -2028,7 +2028,56 @@ def forward( Number of local patches per image in the batch. patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): The patches grid `(num_columns, num_rows)` per image. - """ + + Example single-page OCR: + + ```python + >>> from transformers import AutoProcessor, AutoModelForImageTextToText + + >>> model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", device_map="auto") + >>> processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") + + >>> image = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" + >>> inputs = processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) + + >>> output = model.generate( + ... **inputs, + ... max_new_tokens=4096, + ... no_repeat_ngram_size=35, + ... no_repeat_ngram_window_size=128, + ... ) + >>> processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) + >>> # image [383, 87, 497, 171]\ntext [333, 201, 558, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... + ``` + + Example multi-page OCR: + + ```python + >>> from transformers import AutoProcessor, AutoModelForImageTextToText + + >>> model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", device_map="auto") + >>> processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") + + >>> page1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" + >>> page2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_receipt.jpeg" + >>> num_pages = 2 + + >>> inputs = processor( + ... images=[page1, page2], + ... text="" * num_pages + "Multi page parsing.", + ... crop_to_patches=False, + ... return_tensors="pt", + ... ).to(model.device) + + >>> output = model.generate( + ... **inputs, + ... max_new_tokens=32768, + ... no_repeat_ngram_size=35, + ... no_repeat_ngram_window_size=1024, + ... ) + >>> processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) + >>> # image [382, 87, 489, 174]\ntitle [333, 201, 556, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SCLUTION FORM... + ```""" outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index e43a586dac46..1f4b0257fcc6 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -1246,7 +1246,56 @@ def forward( Number of local patches per image in the batch. patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): The patches grid `(num_columns, num_rows)` per image. - """ + + Example single-page OCR: + + ```python + >>> from transformers import AutoProcessor, AutoModelForImageTextToText + + >>> model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", device_map="auto") + >>> processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") + + >>> image = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" + >>> inputs = processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) + + >>> output = model.generate( + ... **inputs, + ... max_new_tokens=4096, + ... no_repeat_ngram_size=35, + ... no_repeat_ngram_window_size=128, + ... ) + >>> processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) + >>> # image [383, 87, 497, 171]\ntext [333, 201, 558, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... + ``` + + Example multi-page OCR: + + ```python + >>> from transformers import AutoProcessor, AutoModelForImageTextToText + + >>> model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", device_map="auto") + >>> processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") + + >>> page1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" + >>> page2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_receipt.jpeg" + >>> num_pages = 2 + + >>> inputs = processor( + ... images=[page1, page2], + ... text="" * num_pages + "Multi page parsing.", + ... crop_to_patches=False, + ... return_tensors="pt", + ... ).to(model.device) + + >>> output = model.generate( + ... **inputs, + ... max_new_tokens=32768, + ... no_repeat_ngram_size=35, + ... no_repeat_ngram_window_size=1024, + ... ) + >>> processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) + >>> # image [382, 87, 489, 174]\ntitle [333, 201, 556, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SCLUTION FORM... + ```""" outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, @@ -1284,33 +1333,6 @@ def forward( image_hidden_states=outputs.image_hidden_states, ) - def _get_logits_processor(self, generation_config, logits_processor=None, **kwargs): - no_repeat_ngram_size = generation_config.no_repeat_ngram_size - no_repeat_ngram_window_size = getattr(generation_config, "no_repeat_ngram_window_size", None) - - use_sliding_window_processor = ( - no_repeat_ngram_window_size is not None and no_repeat_ngram_size is not None and no_repeat_ngram_size > 0 - ) - if use_sliding_window_processor: - logits_processor = list(logits_processor or []) + [ - UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor( - ngram_size=no_repeat_ngram_size, - window_size=no_repeat_ngram_window_size, - ) - ] - - # Set to None to avoid adding the default NoRepeatNgramLogitsProcessor - generation_config.no_repeat_ngram_size = None - - try: - processors = GenerationMixin._get_logits_processor( - self, generation_config=generation_config, logits_processor=logits_processor, **kwargs - ) - finally: - if use_sliding_window_processor: - generation_config.no_repeat_ngram_size = no_repeat_ngram_size - return processors - def prepare_inputs_for_generation( self, input_ids, @@ -1343,6 +1365,33 @@ def prepare_inputs_for_generation( return model_inputs + def _get_logits_processor(self, generation_config, logits_processor=None, **kwargs): + no_repeat_ngram_size = generation_config.no_repeat_ngram_size + no_repeat_ngram_window_size = getattr(generation_config, "no_repeat_ngram_window_size", None) + + use_sliding_window_processor = ( + no_repeat_ngram_window_size is not None and no_repeat_ngram_size is not None and no_repeat_ngram_size > 0 + ) + if use_sliding_window_processor: + logits_processor = list(logits_processor or []) + [ + UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor( + ngram_size=no_repeat_ngram_size, + window_size=no_repeat_ngram_window_size, + ) + ] + + # Set to None to avoid adding the default NoRepeatNgramLogitsProcessor + generation_config.no_repeat_ngram_size = None + + try: + processors = GenerationMixin._get_logits_processor( + self, generation_config=generation_config, logits_processor=logits_processor, **kwargs + ) + finally: + if use_sliding_window_processor: + generation_config.no_repeat_ngram_size = no_repeat_ngram_size + return processors + __all__ = [ "UnlimitedOcrConfig", From 93b979871735ce895eae5bc1b404da39d5676dfb Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:38:36 +0200 Subject: [PATCH 109/195] update tests --- .../test_modeling_unlimited_ocr.py | 32 ++++++++++++++++--- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index b8532044c841..c9b425809e64 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -173,7 +173,13 @@ def test_small_model_integration_test_document_parsing(self): ) inputs = self.processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): - generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) + generate_ids = model.generate( + **inputs, + do_sample=False, + max_new_tokens=20, + no_repeat_ngram_size=35, + no_repeat_ngram_window_size=128, + ) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True) EXPECTED_DECODED_TEXT = Expectations( { @@ -194,7 +200,13 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se ) inputs = self.processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): - generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) + generate_ids = model.generate( + **inputs, + do_sample=False, + max_new_tokens=20, + no_repeat_ngram_size=35, + no_repeat_ngram_window_size=128, + ) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=False) EXPECTED_DECODED_TEXT = Expectations( { @@ -225,7 +237,13 @@ def test_small_model_integration_test_document_parsing_batched(self): ).to(model.device, dtype=torch.bfloat16) with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): - generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) + generate_ids = model.generate( + **inputs, + do_sample=False, + max_new_tokens=20, + no_repeat_ngram_size=35, + no_repeat_ngram_window_size=128, + ) decoded = self.processor.batch_decode( generate_ids[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True ) @@ -266,7 +284,13 @@ def test_small_model_integration_test_multi_page_document_parsing(self): ).to(model.device) with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): - generate_ids = model.generate(**inputs, do_sample=False, max_new_tokens=20) + generate_ids = model.generate( + **inputs, + do_sample=False, + max_new_tokens=20, + no_repeat_ngram_size=35, + no_repeat_ngram_window_size=128, + ) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True) EXPECTED_DECODED_TEXT = Expectations( From 7dc270906752a4f72d1948be4450fd848edea145 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:39:30 +0200 Subject: [PATCH 110/195] update init --- src/transformers/models/unlimited_ocr/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/models/unlimited_ocr/__init__.py b/src/transformers/models/unlimited_ocr/__init__.py index 737519feca9a..a5fac7aee414 100644 --- a/src/transformers/models/unlimited_ocr/__init__.py +++ b/src/transformers/models/unlimited_ocr/__init__.py @@ -20,6 +20,7 @@ if TYPE_CHECKING: from .configuration_unlimited_ocr import * + from .generation_unlimited_ocr import * from .image_processing_unlimited_ocr import * from .modeling_unlimited_ocr import * from .processing_unlimited_ocr import * From 3c9ab3800c87dac94a8c47620f63fd1451d9c3db Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:39:36 +0200 Subject: [PATCH 111/195] update docs --- docs/source/en/model_doc/unlimited_ocr.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index a57a5e00a15d..1c0fb3c612c6 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -155,6 +155,10 @@ processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=T [[autodoc]] UnlimitedOcrSamVisionConfig +## UnlimitedOcrGenerationConfig + +[[autodoc]] UnlimitedOcrGenerationConfig + ## UnlimitedOcrImageProcessor [[autodoc]] UnlimitedOcrImageProcessor @@ -190,3 +194,7 @@ processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=T [[autodoc]] UnlimitedOcrForConditionalGeneration - forward + +## UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor + +[[autodoc]] UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor From 7519723f1d3e9d7598527e1c6765d29702192df7 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 13:46:20 +0200 Subject: [PATCH 112/195] add generation test --- .../test_generation_unlimited_ocr.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 tests/models/unlimited_ocr/test_generation_unlimited_ocr.py diff --git a/tests/models/unlimited_ocr/test_generation_unlimited_ocr.py b/tests/models/unlimited_ocr/test_generation_unlimited_ocr.py new file mode 100644 index 000000000000..4be03717a7b4 --- /dev/null +++ b/tests/models/unlimited_ocr/test_generation_unlimited_ocr.py @@ -0,0 +1,24 @@ +import unittest + +import torch + +from transformers.models.unlimited_ocr.generation_unlimited_ocr import ( + UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor, +) +from transformers.testing_utils import require_torch, torch_device + + +@require_torch +class UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessorTest(unittest.TestCase): + def test_window_limits_ngram_lookup(self): + vocab_size = 3 + # The (0, 1) bigram appears at the start, so a full-sequence processor would forbid token 1 + # after the trailing 0. A small window should not see that early bigram. + input_ids = torch.tensor([[0, 1, 2, 0]], device=torch_device, dtype=torch.long) + scores = torch.zeros((1, vocab_size), device=torch_device, dtype=torch.float) + + small_window = UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor(ngram_size=2, window_size=2) + full_window = UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor(ngram_size=2, window_size=4) + + self.assertListEqual(torch.isinf(small_window(input_ids, scores.clone())).tolist(), [[False, False, False]]) + self.assertListEqual(torch.isinf(full_window(input_ids, scores.clone())).tolist(), [[False, True, False]]) From 52360595dc28b1c26f0660590b654a910a52d626 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 15:58:13 +0200 Subject: [PATCH 113/195] remove getattr --- src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py | 2 +- src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 0a985cc87624..f1785fe66380 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -2149,7 +2149,7 @@ def prepare_inputs_for_generation( def _get_logits_processor(self, generation_config, logits_processor=None, **kwargs): no_repeat_ngram_size = generation_config.no_repeat_ngram_size - no_repeat_ngram_window_size = getattr(generation_config, "no_repeat_ngram_window_size", None) + no_repeat_ngram_window_size = generation_config.no_repeat_ngram_window_size use_sliding_window_processor = ( no_repeat_ngram_window_size is not None and no_repeat_ngram_size is not None and no_repeat_ngram_size > 0 diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 1f4b0257fcc6..78e7b8929f31 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -1367,7 +1367,7 @@ def prepare_inputs_for_generation( def _get_logits_processor(self, generation_config, logits_processor=None, **kwargs): no_repeat_ngram_size = generation_config.no_repeat_ngram_size - no_repeat_ngram_window_size = getattr(generation_config, "no_repeat_ngram_window_size", None) + no_repeat_ngram_window_size = generation_config.no_repeat_ngram_window_size use_sliding_window_processor = ( no_repeat_ngram_window_size is not None and no_repeat_ngram_size is not None and no_repeat_ngram_size > 0 From af7f77a363ba1962f1c843c93fa1d7a1b8be7595 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 16:01:05 +0200 Subject: [PATCH 114/195] update tests to use revision --- .../unlimited_ocr/test_modeling_unlimited_ocr.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index c9b425809e64..32ab79f9fbcc 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -153,11 +153,12 @@ def _image_features_prepare_config_and_inputs(self): @require_torch class UnlimitedOcrIntegrationTest(unittest.TestCase): - # TODO: change before merge - model_id = "guarin/Unlimited-OCR" + model_id = "baidu/Unlimited-OCR" + #Β TODO: remove revision before merge + revision = "refs/pr/13" def setUp(self): - self.processor = AutoProcessor.from_pretrained(self.model_id) + self.processor = AutoProcessor.from_pretrained(self.model_id, revision=self.revision) def tearDown(self): cleanup(torch_device, gc_collect=True) @@ -165,7 +166,7 @@ def tearDown(self): @slow @require_torch_accelerator def test_small_model_integration_test_document_parsing(self): - model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, device_map=torch_device).eval() + model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, revision=self.revision, device_map=torch_device).eval() image = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -192,7 +193,7 @@ def test_small_model_integration_test_document_parsing(self): @slow @require_torch_accelerator def test_small_model_integration_test_document_parsing_no_skip_special_tokens(self): - model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, device_map=torch_device).eval() + model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, revision=self.revision, device_map=torch_device).eval() image = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -219,7 +220,7 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se @slow @require_torch_accelerator def test_small_model_integration_test_document_parsing_batched(self): - model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, device_map=torch_device).eval() + model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, revision=self.revision, device_map=torch_device).eval() image1 = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -265,7 +266,7 @@ def test_small_model_integration_test_document_parsing_batched(self): @slow @require_torch_accelerator def test_small_model_integration_test_multi_page_document_parsing(self): - model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, device_map=torch_device).eval() + model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, revision=self.revision, device_map=torch_device).eval() image1 = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" From 0ec84b691f27cabc9335d8d0f3b67fddedea9d26 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 14:14:46 +0000 Subject: [PATCH 115/195] update tests --- tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 32ab79f9fbcc..b199f5baa686 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -252,11 +252,11 @@ def test_small_model_integration_test_document_parsing_batched(self): EXPECTED_DECODED_TEXT = Expectations( { ("cuda", None): [ - "image [383, 87, 497, 171]\ntitle [333", + "image [383, 88, 497, 171]\ntitle [333", "header [53, 23, 365, 41]Advanced Template and Styl", ], ("cpu", None): [ - "image [383, 87, 497, 171]\ntitle [333", + "image [383, 88, 497, 171]\ntitle [333", "header [53, 23, 365, 41]Advanced Template and Styl", ], } @@ -296,8 +296,8 @@ def test_small_model_integration_test_multi_page_document_parsing(self): EXPECTED_DECODED_TEXT = Expectations( { - ("cuda", None): "image [382, 87, 489, 174]\n", - ("cpu", None): "image [382, 87, 489, 174]\n", + ("cuda", None): "image [384, 87, 489, 171]\n", + ("cpu", None): "image [384, 87, 489, 171]\n", } ).get_expectation() self.assertEqual(decoded, EXPECTED_DECODED_TEXT) From 95c34af90305355521436c0e8e221ac6a15903d6 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 8 Jul 2026 16:15:28 +0200 Subject: [PATCH 116/195] fix-repo --- .../test_modeling_unlimited_ocr.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index b199f5baa686..8fbf0a3ddb3e 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -154,7 +154,7 @@ def _image_features_prepare_config_and_inputs(self): @require_torch class UnlimitedOcrIntegrationTest(unittest.TestCase): model_id = "baidu/Unlimited-OCR" - #Β TODO: remove revision before merge + # TODO: remove revision before merge revision = "refs/pr/13" def setUp(self): @@ -166,7 +166,9 @@ def tearDown(self): @slow @require_torch_accelerator def test_small_model_integration_test_document_parsing(self): - model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, revision=self.revision, device_map=torch_device).eval() + model = UnlimitedOcrForConditionalGeneration.from_pretrained( + self.model_id, revision=self.revision, device_map=torch_device + ).eval() image = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -193,7 +195,9 @@ def test_small_model_integration_test_document_parsing(self): @slow @require_torch_accelerator def test_small_model_integration_test_document_parsing_no_skip_special_tokens(self): - model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, revision=self.revision, device_map=torch_device).eval() + model = UnlimitedOcrForConditionalGeneration.from_pretrained( + self.model_id, revision=self.revision, device_map=torch_device + ).eval() image = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -220,7 +224,9 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se @slow @require_torch_accelerator def test_small_model_integration_test_document_parsing_batched(self): - model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, revision=self.revision, device_map=torch_device).eval() + model = UnlimitedOcrForConditionalGeneration.from_pretrained( + self.model_id, revision=self.revision, device_map=torch_device + ).eval() image1 = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -266,7 +272,9 @@ def test_small_model_integration_test_document_parsing_batched(self): @slow @require_torch_accelerator def test_small_model_integration_test_multi_page_document_parsing(self): - model = UnlimitedOcrForConditionalGeneration.from_pretrained(self.model_id, revision=self.revision, device_map=torch_device).eval() + model = UnlimitedOcrForConditionalGeneration.from_pretrained( + self.model_id, revision=self.revision, device_map=torch_device + ).eval() image1 = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" From f0504a39e1c852db52ec9db1d608e8e152dad408 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 9 Jul 2026 07:51:27 +0000 Subject: [PATCH 117/195] cleanup tests --- .../test_modeling_unlimited_ocr.py | 17 ++++------------- 1 file changed, 4 insertions(+), 13 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 8fbf0a3ddb3e..ce4b28d78d80 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -77,6 +77,8 @@ def __init__(self, parent, **kwargs): kwargs.setdefault("mlp_layer_types", ["dense", "sparse"]) kwargs.setdefault("moe_intermediate_size", 64) kwargs.setdefault("num_experts_per_tok", 2) + kwargs.setdefault("sliding_window", 4) + kwargs.setdefault("use_sliding_window", True) super().__init__(parent, **kwargs) self.sam_config = { @@ -124,16 +126,6 @@ class UnlimitedOcrModelTest(VLMModelTest, unittest.TestCase): model_tester_class = UnlimitedOcrVisionText2TextModelTester test_all_params_have_gradient = False - def prepare_config_and_inputs_for_generate(self, batch_size=2): - config, inputs_dict = super().prepare_config_and_inputs_for_generate(batch_size=batch_size) - # `is_moe_model` (used by generation tests to pick a looser tolerance for MoE routing noise) checks - # `config._experts_implementation`, which is only populated when a model is instantiated. Some tests - # (e.g. `test_generate_with_static_cache`) check it on the bare config, so mark the model's default - # experts implementation here to get the MoE tolerance. The ring-buffer vs. static sliding cache - # differ by ~1e-3 once amplified through the sparse expert routing, which is well within that bound. - config._experts_implementation = "grouped_mm" - return config, inputs_dict - @unittest.skip( reason="UnlimitedOcrVisionModel builds a hybrid bidirectional+causal mask internally, so SDPA is always called with a non-null `attn_mask`." ) @@ -141,10 +133,9 @@ def test_sdpa_can_dispatch_on_flash(self): pass def _image_features_prepare_config_and_inputs(self): + #Β `test_get_image_features_output` requires `vision_config.hidden` size to be set. + # This is not the case by default as the vision model is a combination of two submodels (SAM + CLIP vision encoder). config, inputs_dict = super()._image_features_prepare_config_and_inputs() - # `get_image_features` returns the concatenation of the SAM feature map and the CLIP encoder output as - # `last_hidden_state`, so its hidden size is the sum of the two. `vision_config` has no `hidden_size` of - # its own, so set it here for `test_get_image_features_output`. config.vision_config.hidden_size = ( config.vision_config.sam_config.downsample_channels[-1] + config.vision_config.encoder_config.hidden_size ) From 73c808fe8045d44fd0756c9b8d7e00281615ff4a Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 9 Jul 2026 10:55:49 +0000 Subject: [PATCH 118/195] update get_max_length --- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 6 ++++++ .../models/unlimited_ocr/modular_unlimited_ocr.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index f1785fe66380..2b24715303fc 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1582,6 +1582,12 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: return kv_length, kv_offset + def get_max_length(self) -> int: + """Return the maximum cache shape of the cache""" + if self.prefill_length is None: + return self.max_cache_len + return self.prefill_length + self.sliding_window + def reset(self) -> None: super().reset() self.prefill_length = None diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 78e7b8929f31..4f800d0b49a7 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -962,6 +962,12 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: return kv_length, kv_offset + def get_max_length(self) -> int: + """Return the maximum cache shape of the cache""" + if self.prefill_length is None: + return self.max_cache_len + return self.prefill_length + self.sliding_window + def reset(self) -> None: super().reset() self.prefill_length = None From ca8c3a7a0600cef0640e4e42473280881a45f6ca Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 9 Jul 2026 10:55:56 +0000 Subject: [PATCH 119/195] update date --- docs/source/en/model_doc/unlimited_ocr.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 1c0fb3c612c6..d08d20e8b5b6 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -16,7 +16,7 @@ limitations under the License. ⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer. --> -*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-07-07.* +*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-07-09.* # UnlimitedOcr From 8e0fbc9fa3ce5c8bc83b98722329f1f812a26723 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 9 Jul 2026 13:48:13 +0000 Subject: [PATCH 120/195] add cache sliding window too small test --- .../unlimited_ocr/modeling_unlimited_ocr.py | 26 ++++++---- .../unlimited_ocr/modular_unlimited_ocr.py | 26 ++++++---- .../test_modeling_unlimited_ocr.py | 51 ++++++++++++++++++- 3 files changed, 82 insertions(+), 21 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 2b24715303fc..03d1a3e504cc 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1493,21 +1493,27 @@ def update( # Already full but using more than 1 new token (e.g. chat continuation, etc...) else: full_key_states = torch.cat( - (self.keys[:, :, window_start + 1 : window_start + self.sliding_window, :], key_states), dim=-2 + ( + self.keys[:, :, :window_start, :], + self.keys[:, :, window_start + 1 : window_start + self.sliding_window, :], + key_states, + ), + dim=-2, ) full_value_states = torch.cat( - (self.values[:, :, window_start + 1 : window_start + self.sliding_window, :], value_states), dim=-2 + ( + self.values[:, :, :window_start, :], + self.values[:, :, window_start + 1 : window_start + self.sliding_window, :], + value_states, + ), + dim=-2, ) # Not yet full, but becoming full on this update elif current_length + kv_length > self.sliding_window: - # Fast path, no need to cat() in this case, as the window is currently empty - if current_length == 0: - full_key_states = key_states - full_value_states = value_states - else: - window = slice(window_start, window_start + current_length) - full_key_states = torch.cat((self.keys[:, :, window, :], key_states), dim=-2) - full_value_states = torch.cat((self.values[:, :, window, :], value_states), dim=-2) + full_key_states = torch.cat((self.keys[:, :, : window_start + current_length, :], key_states), dim=-2) + full_value_states = torch.cat( + (self.values[:, :, : window_start + current_length, :], value_states), dim=-2 + ) else: # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten. diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 4f800d0b49a7..a2455694736b 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -873,21 +873,27 @@ def update( # Already full but using more than 1 new token (e.g. chat continuation, etc...) else: full_key_states = torch.cat( - (self.keys[:, :, window_start + 1 : window_start + self.sliding_window, :], key_states), dim=-2 + ( + self.keys[:, :, :window_start, :], + self.keys[:, :, window_start + 1 : window_start + self.sliding_window, :], + key_states, + ), + dim=-2, ) full_value_states = torch.cat( - (self.values[:, :, window_start + 1 : window_start + self.sliding_window, :], value_states), dim=-2 + ( + self.values[:, :, :window_start, :], + self.values[:, :, window_start + 1 : window_start + self.sliding_window, :], + value_states, + ), + dim=-2, ) # Not yet full, but becoming full on this update elif current_length + kv_length > self.sliding_window: - # Fast path, no need to cat() in this case, as the window is currently empty - if current_length == 0: - full_key_states = key_states - full_value_states = value_states - else: - window = slice(window_start, window_start + current_length) - full_key_states = torch.cat((self.keys[:, :, window, :], key_states), dim=-2) - full_value_states = torch.cat((self.values[:, :, window, :], value_states), dim=-2) + full_key_states = torch.cat((self.keys[:, :, : window_start + current_length, :], key_states), dim=-2) + full_value_states = torch.cat( + (self.values[:, :, : window_start + current_length, :], value_states), dim=-2 + ) else: # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten. diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index ce4b28d78d80..3f92a6cf173f 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -133,7 +133,7 @@ def test_sdpa_can_dispatch_on_flash(self): pass def _image_features_prepare_config_and_inputs(self): - #Β `test_get_image_features_output` requires `vision_config.hidden` size to be set. + # `test_get_image_features_output` requires `vision_config.hidden` size to be set. # This is not the case by default as the vision model is a combination of two submodels (SAM + CLIP vision encoder). config, inputs_dict = super()._image_features_prepare_config_and_inputs() config.vision_config.hidden_size = ( @@ -141,6 +141,55 @@ def _image_features_prepare_config_and_inputs(self): ) return config, inputs_dict + def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_length, config): + # max cache length can be smaller than sequence length + max_length = past_key_values.get_max_length() + seq_length = min(seq_length, max_length) if max_length >= 0 else seq_length + super()._check_past_key_values_for_generate(batch_size, past_key_values, seq_length, config) + + def _check_generate_cache_sliding_window_too_small(self, cache_implementation): + """Test that reference sliding window cache works correctly when decoding more than sliding_window tokens at once.""" + for model_class in self.all_generative_model_classes: + config, inputs_dict = self.prepare_config_and_inputs_for_generate() + self.assertEqual(config.text_config.sliding_window, 4) + + model = model_class(config).to(torch_device).eval() + + # Resume from cache doesn't work with random attention mask. + inputs_dict["attention_mask"] = torch.ones_like(inputs_dict["attention_mask"]) + + out_reference = model.generate(**inputs_dict, max_new_tokens=10, do_sample=False) + + # Prefill the cache. We need at least max_new_tokens=2 to make the cache mark prefill + # as complete. Prefill is only marked as complete once single token is added to + # the cache (kv_length == 1). As the last decoded token isn't added to the cache + # we have to decode at least two tokens. + out_prefill = model.generate( + **inputs_dict, + max_new_tokens=2, + max_cache_len=100, + do_sample=False, + return_dict_in_generate=True, + use_cache=True, + cache_implementation=cache_implementation, + ) + + # Decode from cache by passing more than sliding_window unseen input ids. + out = model.generate( + input_ids=out_reference[:, :-3], + past_key_values=out_prefill.past_key_values, + max_new_tokens=3, + do_sample=False, + use_cache=True, + ) + self.assertEqual(out.tolist(), out_reference.tolist()) + + def test_generate_dynamic_cache_sliding_window_too_small(self): + self._check_generate_cache_sliding_window_too_small("dynamic") + + def test_generate_static_cache_sliding_window_too_small(self): + self._check_generate_cache_sliding_window_too_small("static") + @require_torch class UnlimitedOcrIntegrationTest(unittest.TestCase): From 7e14a2b3c4db60009f928ce7e030dbd0101ef98b Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 9 Jul 2026 14:11:44 +0000 Subject: [PATCH 121/195] fix get_max_length --- src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py | 2 +- src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 03d1a3e504cc..63fd846680a3 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1592,7 +1592,7 @@ def get_max_length(self) -> int: """Return the maximum cache shape of the cache""" if self.prefill_length is None: return self.max_cache_len - return self.prefill_length + self.sliding_window + return min(self.max_cache_len, self.prefill_length + self.sliding_window) def reset(self) -> None: super().reset() diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index a2455694736b..485a6bd9a354 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -972,7 +972,7 @@ def get_max_length(self) -> int: """Return the maximum cache shape of the cache""" if self.prefill_length is None: return self.max_cache_len - return self.prefill_length + self.sliding_window + return min(self.max_cache_len, self.prefill_length + self.sliding_window) def reset(self) -> None: super().reset() From cc7d26ed8755999fe0921f38a01b804a54332c60 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 9 Jul 2026 16:29:11 +0200 Subject: [PATCH 122/195] add comments --- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 6 ++++++ .../models/unlimited_ocr/modular_unlimited_ocr.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 63fd846680a3..7fedb140094a 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1314,6 +1314,9 @@ def update( full_value_states = torch.cat([self.values, value_states], dim=-2) # Prefill + # We assume kv_length > 1 during prefill and that the first decode step + # has kv_length == 1. This is required to differentiate between prefill + # and decode phases. Later decode steps can have any kv_length. if self.prefill_length is None and kv_length > 1: self.keys = full_key_states self.values = full_value_states @@ -1436,6 +1439,9 @@ def update( kv_length = key_states.shape[-2] # Prefill + # We assume kv_length > 1 during prefill and that the first decode step + # has kv_length == 1. This is required to differentiate between prefill + # and decode phases. Later decode steps can have any kv_length. if self.prefill_length is None and kv_length > 1: # Chunked prefill: the buffer was sized for the first chunk only, so grow it to hold all prefill seen so far plus this chunk plus the reserved window. required_length = min(self.max_cache_len, self.cumulative_length_int + kv_length + self.sliding_window) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 485a6bd9a354..1df72782d2f0 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -694,6 +694,9 @@ def update( full_value_states = torch.cat([self.values, value_states], dim=-2) # Prefill + # We assume kv_length > 1 during prefill and that the first decode step + # has kv_length == 1. This is required to differentiate between prefill + # and decode phases. Later decode steps can have any kv_length. if self.prefill_length is None and kv_length > 1: self.keys = full_key_states self.values = full_value_states @@ -816,6 +819,9 @@ def update( kv_length = key_states.shape[-2] # Prefill + # We assume kv_length > 1 during prefill and that the first decode step + # has kv_length == 1. This is required to differentiate between prefill + # and decode phases. Later decode steps can have any kv_length. if self.prefill_length is None and kv_length > 1: # Chunked prefill: the buffer was sized for the first chunk only, so grow it to hold all prefill seen so far plus this chunk plus the reserved window. required_length = min(self.max_cache_len, self.cumulative_length_int + kv_length + self.sliding_window) From 5f37dfabb6321ec70a523303bb1c89cf9e7e66e8 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 9 Jul 2026 18:06:48 +0200 Subject: [PATCH 123/195] update comment --- tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 3f92a6cf173f..0e82aed61e58 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -161,9 +161,9 @@ def _check_generate_cache_sliding_window_too_small(self, cache_implementation): out_reference = model.generate(**inputs_dict, max_new_tokens=10, do_sample=False) # Prefill the cache. We need at least max_new_tokens=2 to make the cache mark prefill - # as complete. Prefill is only marked as complete once single token is added to - # the cache (kv_length == 1). As the last decoded token isn't added to the cache - # we have to decode at least two tokens. + # as complete. Prefill is only marked as complete once a single token is added to + # the cache (kv_length == 1). As the last decoded token from a .generate call isn't + # added to the cache we have to set max_new_tokens=2. out_prefill = model.generate( **inputs_dict, max_new_tokens=2, From faf2f105447917ace7818e941ca764dac02bfab1 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 10 Jul 2026 08:57:15 +0200 Subject: [PATCH 124/195] update test --- .../unlimited_ocr/test_modeling_unlimited_ocr.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 0e82aed61e58..f1265edfb6f6 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -163,7 +163,7 @@ def _check_generate_cache_sliding_window_too_small(self, cache_implementation): # Prefill the cache. We need at least max_new_tokens=2 to make the cache mark prefill # as complete. Prefill is only marked as complete once a single token is added to # the cache (kv_length == 1). As the last decoded token from a .generate call isn't - # added to the cache we have to set max_new_tokens=2. + # added to the cache we have to decode at least 2. out_prefill = model.generate( **inputs_dict, max_new_tokens=2, @@ -174,9 +174,13 @@ def _check_generate_cache_sliding_window_too_small(self, cache_implementation): cache_implementation=cache_implementation, ) - # Decode from cache by passing more than sliding_window unseen input ids. + # Decode from cache and pass more than sliding_window unseen input ids. + input_ids_reference = out_reference[:, :-3] out = model.generate( - input_ids=out_reference[:, :-3], + input_ids=input_ids_reference, + # Pass attention mask explicitly as input_ids_reference sometimes contains randomly generated + # pad tokens which trips up the automatic attention mask generation. + attention_mask=torch.ones_like(input_ids_reference), past_key_values=out_prefill.past_key_values, max_new_tokens=3, do_sample=False, From d8accb9fb74a0e317414bbf6f5981dcdcb8bfad2 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 10 Jul 2026 09:20:55 +0200 Subject: [PATCH 125/195] add full cache test --- .../test_modeling_unlimited_ocr.py | 26 ++++++++++++------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index f1265edfb6f6..a92755d6f056 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -147,7 +147,7 @@ def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_l seq_length = min(seq_length, max_length) if max_length >= 0 else seq_length super()._check_past_key_values_for_generate(batch_size, past_key_values, seq_length, config) - def _check_generate_cache_sliding_window_too_small(self, cache_implementation): + def _check_generate_cache_sliding_window_too_small(self, cache_implementation: str, prefill_max_new_tokens: int): """Test that reference sliding window cache works correctly when decoding more than sliding_window tokens at once.""" for model_class in self.all_generative_model_classes: config, inputs_dict = self.prepare_config_and_inputs_for_generate() @@ -158,15 +158,15 @@ def _check_generate_cache_sliding_window_too_small(self, cache_implementation): # Resume from cache doesn't work with random attention mask. inputs_dict["attention_mask"] = torch.ones_like(inputs_dict["attention_mask"]) - out_reference = model.generate(**inputs_dict, max_new_tokens=10, do_sample=False) + out_reference = model.generate(**inputs_dict, max_new_tokens=15, do_sample=False) - # Prefill the cache. We need at least max_new_tokens=2 to make the cache mark prefill - # as complete. Prefill is only marked as complete once a single token is added to - # the cache (kv_length == 1). As the last decoded token from a .generate call isn't - # added to the cache we have to decode at least 2. + # Prefill the cache. + #Β As the last decoded token from a .generate call isn't added to the cache we have to + # decode at least 2 tokens for the cache update call to mark the prefill as complete. + #Β Prefill is only marked as complete once a single decoded token is added to the cache (kv_length == 1). out_prefill = model.generate( **inputs_dict, - max_new_tokens=2, + max_new_tokens=prefill_max_new_tokens, max_cache_len=100, do_sample=False, return_dict_in_generate=True, @@ -189,10 +189,18 @@ def _check_generate_cache_sliding_window_too_small(self, cache_implementation): self.assertEqual(out.tolist(), out_reference.tolist()) def test_generate_dynamic_cache_sliding_window_too_small(self): - self._check_generate_cache_sliding_window_too_small("dynamic") + self._check_generate_cache_sliding_window_too_small(cache_implementation="dynamic", prefill_max_new_tokens=2) def test_generate_static_cache_sliding_window_too_small(self): - self._check_generate_cache_sliding_window_too_small("static") + self._check_generate_cache_sliding_window_too_small(cache_implementation="static", prefill_max_new_tokens=2) + + def test_generate_dynamic_cache_sliding_window_too_small_cache_full(self): + """Continue from full cache""" + self._check_generate_cache_sliding_window_too_small(cache_implementation="dynamic", prefill_max_new_tokens=6) + + def test_generate_static_cache_sliding_window_too_small_cache_full(self): + """Continue from full cache""" + self._check_generate_cache_sliding_window_too_small(cache_implementation="static", prefill_max_new_tokens=6) @require_torch From 22c38d9d5c03a0029a00512dacec153c7194b6f9 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 10 Jul 2026 07:34:32 +0000 Subject: [PATCH 126/195] add detection docs --- docs/source/en/model_doc/unlimited_ocr.md | 58 +++++++++++++++++++++-- 1 file changed, 55 insertions(+), 3 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index d08d20e8b5b6..2cd7cc337a90 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -65,8 +65,12 @@ output = model.generate( ) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) # image [383, 87, 497, 171]\ntext [333, 201, 558, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... + +#Β All bounding boxes are in (x1, y1, x2, y2) format with coordinates normalized to [0, 999] ``` +### Batch processing + For batch processing, pass multiple images and prompts at once. Set `padding=True` for the processor if images have different sizes. @@ -93,13 +97,59 @@ output = model.generate( ) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) # image [383, 87, 497, 171]\ntext [333, 201, 558, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... -``` -Set `skip_special_tokens=False` to wrap all detections and region types in `<|det|>...<|/det|>` markers. This if useful for further post-processing of the output which will look like this: +#Β All bounding boxes are in (x1, y1, x2, y2) format with coordinates normalized to [0, 999] ``` -<|det|>image [383, 87, 497, 171]<|/det|>\n<|det|>text [333, 201, 558, 230]<|/det|>R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... + +### Region detections + +Set `skip_special_tokens=False` to wrap all detections and region types in `<|det|>...<|/det|>` markers. This is useful for further post-processing of the output, for example to plot the detected bounding boxes on the image. Each detection is wrapped as `<|det|>region_type [x1, y1, x2, y2]<|/det|>` with coordinates normalized to a `[0, 999]` range. Parse the markers with a regular expression and rescale the coordinates to the image size to plot the bounding boxes. + +```python +import re +from transformers import AutoProcessor, AutoModelForImageTextToText + +model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", device_map="auto") +processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") + +image = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" +inputs = processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) + +output = model.generate( + **inputs, + max_new_tokens=4096, + no_repeat_ngram_size=35, + no_repeat_ngram_window_size=128, +) +decoded = processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=False) +#Β <|det|>image [383, 87, 497, 171]<|/det|>\n<|det|>text [333, 201, 558, 230]<|/det|>R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... + +detections = re.findall(r"<\|det\|>(\S+) \[(\d+), (\d+), (\d+), (\d+)\]<\|/det\|>", decoded) + +# Visualization +import random +import matplotlib.pyplot as plt +import matplotlib.patches as patches +from transformers.image_utils import load_image + +image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg") +width, height = image.size + +figure, axis = plt.subplots(figsize=(10, 12)) +axis.imshow(image) +for region_type, x1, y1, x2, y2 in detections: + x1, y1, x2, y2 = int(x1) / 999 * width, int(y1) / 999 * height, int(x2) / 999 * width, int(y2) / 999 * height + color = (random.random(), random.random(), random.random()) + rectangle = patches.Rectangle((x1, y1), x2 - x1, y2 - y1, linewidth=1.5, edgecolor=color, facecolor="none") + axis.add_patch(rectangle) + axis.text(x1, y1, region_type, color="white", fontsize=8, backgroundcolor=color, verticalalignment="top") +axis.axis("off") +plt.show() ``` + + + Multi-page documents can be parsed jointly in a single forward pass by passing all page images together. Include one `` token per page in the text prompt so the model processes all pages as a continuous document. @@ -130,6 +180,8 @@ output = model.generate( ) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) # image [382, 87, 489, 174]\ntitle [333, 201, 556, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SCLUTION FORM... + +#Β All bounding boxes are in (x1, y1, x2, y2) format with coordinates normalized to [0, 999] ``` From cc6edbba6a0921948a3e5b8c805454a5abe9743e Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 10 Jul 2026 09:40:14 +0200 Subject: [PATCH 127/195] fix-repo --- tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index a92755d6f056..13238b865b22 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -161,9 +161,9 @@ def _check_generate_cache_sliding_window_too_small(self, cache_implementation: s out_reference = model.generate(**inputs_dict, max_new_tokens=15, do_sample=False) # Prefill the cache. - #Β As the last decoded token from a .generate call isn't added to the cache we have to + # As the last decoded token from a .generate call isn't added to the cache we have to # decode at least 2 tokens for the cache update call to mark the prefill as complete. - #Β Prefill is only marked as complete once a single decoded token is added to the cache (kv_length == 1). + # Prefill is only marked as complete once a single decoded token is added to the cache (kv_length == 1). out_prefill = model.generate( **inputs_dict, max_new_tokens=prefill_max_new_tokens, From d0636d6b446e517a4cfcb69d71a4cddf0f16a30c Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:28:07 +0100 Subject: [PATCH 128/195] move logits processor --- docs/source/en/model_doc/unlimited_ocr.md | 10 +-- .../generation/configuration_utils.py | 2 +- src/transformers/generation/utils.py | 4 +- .../models/unlimited_ocr/__init__.py | 1 - .../unlimited_ocr/generation_unlimited_ocr.py | 79 +++++++++++++++---- .../unlimited_ocr/modeling_unlimited_ocr.py | 39 +-------- .../unlimited_ocr/modular_unlimited_ocr.py | 74 ++++++++--------- .../test_modeling_unlimited_ocr.py | 12 ++- 8 files changed, 115 insertions(+), 106 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 2cd7cc337a90..3c67fc24ee9b 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -16,7 +16,7 @@ limitations under the License. ⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer. --> -*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-07-09.* +*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-07-27.* # UnlimitedOcr @@ -207,10 +207,6 @@ processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=T [[autodoc]] UnlimitedOcrSamVisionConfig -## UnlimitedOcrGenerationConfig - -[[autodoc]] UnlimitedOcrGenerationConfig - ## UnlimitedOcrImageProcessor [[autodoc]] UnlimitedOcrImageProcessor @@ -246,7 +242,3 @@ processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=T [[autodoc]] UnlimitedOcrForConditionalGeneration - forward - -## UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor - -[[autodoc]] UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor diff --git a/src/transformers/generation/configuration_utils.py b/src/transformers/generation/configuration_utils.py index d32c73a83109..1b9cebfe0ba3 100644 --- a/src/transformers/generation/configuration_utils.py +++ b/src/transformers/generation/configuration_utils.py @@ -1279,7 +1279,7 @@ def from_model_config(cls, model_config: Union["PreTrainedConfig", dict]) -> "Ge model_config = model_config.get_text_config(decoder=True) model_config = model_config.to_dict() - default_generation_config = cls() + default_generation_config = GenerationConfig() for attr in generation_config.to_dict(): is_unset = getattr(generation_config, attr) == getattr(default_generation_config, attr) if attr in model_config and is_unset: diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 0ffa5cb30679..de545f1b8aa1 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -399,7 +399,7 @@ def adjust_generation_fn( } # Load generation config try: - self.generation_config = self.generation_config_class.from_pretrained( + self.generation_config = GenerationConfig.from_pretrained( pretrained_model_name_or_path, _from_auto=from_auto_class, _from_pipeline=from_pipeline, @@ -412,7 +412,7 @@ def adjust_generation_fn( logger.info( "Generation config file not found, using a generation config created from the model config." ) - self.generation_config = self.generation_config_class.from_pretrained( + self.generation_config = GenerationConfig.from_pretrained( pretrained_model_name_or_path, config_file_name="config.json", _from_auto=from_auto_class, diff --git a/src/transformers/models/unlimited_ocr/__init__.py b/src/transformers/models/unlimited_ocr/__init__.py index a5fac7aee414..737519feca9a 100644 --- a/src/transformers/models/unlimited_ocr/__init__.py +++ b/src/transformers/models/unlimited_ocr/__init__.py @@ -20,7 +20,6 @@ if TYPE_CHECKING: from .configuration_unlimited_ocr import * - from .generation_unlimited_ocr import * from .image_processing_unlimited_ocr import * from .modeling_unlimited_ocr import * from .processing_unlimited_ocr import * diff --git a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py index 013703fc1346..b67046eec230 100644 --- a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py @@ -11,27 +11,21 @@ # 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. +from collections.abc import Callable +from typing import Any + import torch from ...generation.configuration_utils import GenerationConfig -from ...generation.logits_process import LOGITS_PROCESSOR_INPUTS_DOCSTRING, NoRepeatNGramLogitsProcessor +from ...generation.logits_process import ( + LOGITS_PROCESSOR_INPUTS_DOCSTRING, + LogitsProcessorList, + NoRepeatNGramLogitsProcessor, +) +from ...generation.utils import GenerationMixin from ...utils import add_start_docstrings -class UnlimitedOcrGenerationConfig(GenerationConfig): - r"""A GenerationConfig class with parameterization customized for UnlimitedOcr. - - Args: - no_repeat_ngram_window_size (`int`, *optional*): - If set together with `no_repeat_ngram_size`, n-gram repetitions are blocked only within this many - trailing tokens instead of over the whole sequence. - """ - - def __init__(self, no_repeat_ngram_window_size: int | None = None, **kwargs): - super().__init__(**kwargs) - self.no_repeat_ngram_window_size = no_repeat_ngram_window_size - - class UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor(NoRepeatNGramLogitsProcessor): r""" Identical to [`NoRepeatNGramLogitsProcessor`] but blocks n-gram repetitions only within the last @@ -55,4 +49,57 @@ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> to return super().__call__(input_ids[:, -self.window_size :], scores) -__all__ = ["UnlimitedOcrGenerationConfig", "UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor"] +class UnlimitedOcrGenerationMixin(GenerationMixin): + r""" + Adds support for the `no_repeat_ngram_window_size` generation option. If set together with `no_repeat_ngram_size`, + n-gram repetitions are blocked only within this many trailing tokens instead of over the whole sequence. + + `no_repeat_ngram_window_size` is a model specific option and must already be present on the generation config + before it can be passed to `generate`, either through the checkpoint's `generation_config.json` or by setting + `model.generation_config.no_repeat_ngram_window_size` after loading the model. + """ + + def _get_logits_processor( + self, + generation_config: GenerationConfig, + input_ids_seq_length: int | None = None, + encoder_input_ids: torch.LongTensor | None = None, + prefix_allowed_tokens_fn: Callable[[int, torch.Tensor], list[int]] | None = None, + logits_processor: LogitsProcessorList | None = None, + device: str | None = None, + model_kwargs: dict[str, Any] | None = None, + negative_prompt_ids: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + ) -> LogitsProcessorList: + no_repeat_ngram_size = generation_config.no_repeat_ngram_size + no_repeat_ngram_window_size = getattr(generation_config, "no_repeat_ngram_window_size", None) + use_sliding_window_processor = False + + if no_repeat_ngram_window_size is not None and no_repeat_ngram_size is not None and no_repeat_ngram_size > 0: + use_sliding_window_processor = True + logits_processor = LogitsProcessorList(logits_processor or []) + logits_processor.append( + UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor( + ngram_size=no_repeat_ngram_size, + window_size=no_repeat_ngram_window_size, + ) + ) + # Set to None to avoid also adding the default `NoRepeatNGramLogitsProcessor` + generation_config.no_repeat_ngram_size = None + + try: + processors = super()._get_logits_processor( + generation_config=generation_config, + input_ids_seq_length=input_ids_seq_length, + encoder_input_ids=encoder_input_ids, + prefix_allowed_tokens_fn=prefix_allowed_tokens_fn, + logits_processor=logits_processor, + device=device, + model_kwargs=model_kwargs, + negative_prompt_ids=negative_prompt_ids, + negative_prompt_attention_mask=negative_prompt_attention_mask, + ) + finally: + if use_sliding_window_processor: + generation_config.no_repeat_ngram_size = no_repeat_ngram_size + return processors diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 7fedb140094a..2ec738083da3 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -31,7 +31,6 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, DynamicLayer, DynamicSlidingWindowLayer, StaticSlidingWindowLayer from ...configuration_utils import PreTrainedConfig -from ...generation import GenerationMixin from ...integrations import ( use_experts_implementation, use_kernel_forward_from_hub, @@ -63,10 +62,7 @@ UnlimitedOcrVisionConfig, UnlimitedOcrVisionEncoderConfig, ) -from .generation_unlimited_ocr import ( - UnlimitedOcrGenerationConfig, - UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor, -) +from .generation_unlimited_ocr import UnlimitedOcrGenerationMixin logger = logging.get_logger(__name__) @@ -1981,9 +1977,8 @@ def forward( @auto_docstring -class UnlimitedOcrForConditionalGeneration(UnlimitedOcrPreTrainedModel, GenerationMixin): +class UnlimitedOcrForConditionalGeneration(UnlimitedOcrPreTrainedModel, UnlimitedOcrGenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} - generation_config_class = UnlimitedOcrGenerationConfig def __init__(self, config: UnlimitedOcrConfig): super().__init__(config) @@ -1994,9 +1989,6 @@ def __init__(self, config: UnlimitedOcrConfig): def get_output_embeddings(self) -> nn.Module: return self.lm_head - def pack_image_features(self): - raise NotImplementedError("UnlimitedOcr does not use pack_image_features") - @can_return_tuple @auto_docstring def get_image_features( @@ -2165,33 +2157,6 @@ def prepare_inputs_for_generation( return model_inputs - def _get_logits_processor(self, generation_config, logits_processor=None, **kwargs): - no_repeat_ngram_size = generation_config.no_repeat_ngram_size - no_repeat_ngram_window_size = generation_config.no_repeat_ngram_window_size - - use_sliding_window_processor = ( - no_repeat_ngram_window_size is not None and no_repeat_ngram_size is not None and no_repeat_ngram_size > 0 - ) - if use_sliding_window_processor: - logits_processor = list(logits_processor or []) + [ - UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor( - ngram_size=no_repeat_ngram_size, - window_size=no_repeat_ngram_window_size, - ) - ] - - # Set to None to avoid adding the default NoRepeatNgramLogitsProcessor - generation_config.no_repeat_ngram_size = None - - try: - processors = super()._get_logits_processor( - generation_config=generation_config, logits_processor=logits_processor, **kwargs - ) - finally: - if use_sliding_window_processor: - generation_config.no_repeat_ngram_size = no_repeat_ngram_size - return processors - __all__ = [ "UnlimitedOcrForConditionalGeneration", diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 1df72782d2f0..1fdddbdbe370 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -28,11 +28,10 @@ ) from ...configuration_utils import PreTrainedConfig from ...feature_extraction_utils import BatchFeature -from ...generation import GenerationMixin from ...image_transforms import group_images_by_shape, reorder_images from ...image_utils import PILImageResampling, SizeDict from ...masking_utils import and_masks, causal_mask_function, create_causal_mask, create_sliding_window_causal_mask -from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling from ...processing_utils import Unpack from ...tokenization_utils_base import TextInput from ...utils import ( @@ -59,7 +58,6 @@ ) from ..deepseek_ocr2.modeling_deepseek_ocr2 import ( DeepseekOcr2CausalLMOutputWithPast, - DeepseekOcr2ForConditionalGeneration, DeepseekOcr2Model, DeepseekOcr2ModelOutputWithPast, DeepseekOcr2ModelOutputWithPooling, @@ -71,10 +69,7 @@ ) from ..deepseek_ocr2.processing_deepseek_ocr2 import DeepseekOcr2Processor, DeepseekOcr2ProcessorKwargs from ..got_ocr2.configuration_got_ocr2 import GotOcr2VisionConfig -from .generation_unlimited_ocr import ( - UnlimitedOcrGenerationConfig, - UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor, -) +from .generation_unlimited_ocr import UnlimitedOcrGenerationMixin class UnlimitedOcrImageProcessorKwargs(DeepseekOcr2ImageProcessorKwargs): @@ -1236,8 +1231,42 @@ def forward( ) -class UnlimitedOcrForConditionalGeneration(DeepseekOcr2ForConditionalGeneration): - generation_config_class = UnlimitedOcrGenerationConfig +@auto_docstring +class UnlimitedOcrForConditionalGeneration(UnlimitedOcrPreTrainedModel, UnlimitedOcrGenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} + + def __init__(self, config: UnlimitedOcrConfig): + super().__init__(config) + self.model = UnlimitedOcrModel(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 + + @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 @@ -1383,33 +1412,6 @@ def prepare_inputs_for_generation( return model_inputs - def _get_logits_processor(self, generation_config, logits_processor=None, **kwargs): - no_repeat_ngram_size = generation_config.no_repeat_ngram_size - no_repeat_ngram_window_size = generation_config.no_repeat_ngram_window_size - - use_sliding_window_processor = ( - no_repeat_ngram_window_size is not None and no_repeat_ngram_size is not None and no_repeat_ngram_size > 0 - ) - if use_sliding_window_processor: - logits_processor = list(logits_processor or []) + [ - UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor( - ngram_size=no_repeat_ngram_size, - window_size=no_repeat_ngram_window_size, - ) - ] - - # Set to None to avoid adding the default NoRepeatNgramLogitsProcessor - generation_config.no_repeat_ngram_size = None - - try: - processors = GenerationMixin._get_logits_processor( - self, generation_config=generation_config, logits_processor=logits_processor, **kwargs - ) - finally: - if use_sliding_window_processor: - generation_config.no_repeat_ngram_size = no_repeat_ngram_size - return processors - __all__ = [ "UnlimitedOcrConfig", diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 13238b865b22..f7a66a41d472 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -221,6 +221,8 @@ def test_small_model_integration_test_document_parsing(self): model = UnlimitedOcrForConditionalGeneration.from_pretrained( self.model_id, revision=self.revision, device_map=torch_device ).eval() + # TODO: Pass as kwarg to generate one generation_config.json is on the hub + model.generation_config.no_repeat_ngram_window_size = 128 image = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -233,7 +235,6 @@ def test_small_model_integration_test_document_parsing(self): do_sample=False, max_new_tokens=20, no_repeat_ngram_size=35, - no_repeat_ngram_window_size=128, ) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True) EXPECTED_DECODED_TEXT = Expectations( @@ -250,6 +251,8 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se model = UnlimitedOcrForConditionalGeneration.from_pretrained( self.model_id, revision=self.revision, device_map=torch_device ).eval() + # TODO: Pass as kwarg to generate one generation_config.json is on the hub + model.generation_config.no_repeat_ngram_window_size = 128 image = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -262,7 +265,6 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se do_sample=False, max_new_tokens=20, no_repeat_ngram_size=35, - no_repeat_ngram_window_size=128, ) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=False) EXPECTED_DECODED_TEXT = Expectations( @@ -279,6 +281,8 @@ def test_small_model_integration_test_document_parsing_batched(self): model = UnlimitedOcrForConditionalGeneration.from_pretrained( self.model_id, revision=self.revision, device_map=torch_device ).eval() + # TODO: Pass as kwarg to generate one generation_config.json is on the hub + model.generation_config.no_repeat_ngram_window_size = 128 image1 = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -301,7 +305,6 @@ def test_small_model_integration_test_document_parsing_batched(self): do_sample=False, max_new_tokens=20, no_repeat_ngram_size=35, - no_repeat_ngram_window_size=128, ) decoded = self.processor.batch_decode( generate_ids[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True @@ -327,6 +330,8 @@ def test_small_model_integration_test_multi_page_document_parsing(self): model = UnlimitedOcrForConditionalGeneration.from_pretrained( self.model_id, revision=self.revision, device_map=torch_device ).eval() + # TODO: Pass as kwarg to generate one generation_config.json is on the hub + model.generation_config.no_repeat_ngram_window_size = 128 image1 = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -350,7 +355,6 @@ def test_small_model_integration_test_multi_page_document_parsing(self): do_sample=False, max_new_tokens=20, no_repeat_ngram_size=35, - no_repeat_ngram_window_size=128, ) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True) From f9ad01cdd76d20ac4a88bc8eb30c1ba8e2bcd9dc Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:15:39 +0000 Subject: [PATCH 129/195] update test --- .../unlimited_ocr/test_modeling_unlimited_ocr.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index f7a66a41d472..b09543804b39 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -221,8 +221,6 @@ def test_small_model_integration_test_document_parsing(self): model = UnlimitedOcrForConditionalGeneration.from_pretrained( self.model_id, revision=self.revision, device_map=torch_device ).eval() - # TODO: Pass as kwarg to generate one generation_config.json is on the hub - model.generation_config.no_repeat_ngram_window_size = 128 image = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -234,7 +232,6 @@ def test_small_model_integration_test_document_parsing(self): **inputs, do_sample=False, max_new_tokens=20, - no_repeat_ngram_size=35, ) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True) EXPECTED_DECODED_TEXT = Expectations( @@ -251,8 +248,6 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se model = UnlimitedOcrForConditionalGeneration.from_pretrained( self.model_id, revision=self.revision, device_map=torch_device ).eval() - # TODO: Pass as kwarg to generate one generation_config.json is on the hub - model.generation_config.no_repeat_ngram_window_size = 128 image = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -264,7 +259,6 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se **inputs, do_sample=False, max_new_tokens=20, - no_repeat_ngram_size=35, ) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=False) EXPECTED_DECODED_TEXT = Expectations( @@ -281,8 +275,6 @@ def test_small_model_integration_test_document_parsing_batched(self): model = UnlimitedOcrForConditionalGeneration.from_pretrained( self.model_id, revision=self.revision, device_map=torch_device ).eval() - # TODO: Pass as kwarg to generate one generation_config.json is on the hub - model.generation_config.no_repeat_ngram_window_size = 128 image1 = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -304,7 +296,6 @@ def test_small_model_integration_test_document_parsing_batched(self): **inputs, do_sample=False, max_new_tokens=20, - no_repeat_ngram_size=35, ) decoded = self.processor.batch_decode( generate_ids[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True @@ -330,8 +321,6 @@ def test_small_model_integration_test_multi_page_document_parsing(self): model = UnlimitedOcrForConditionalGeneration.from_pretrained( self.model_id, revision=self.revision, device_map=torch_device ).eval() - # TODO: Pass as kwarg to generate one generation_config.json is on the hub - model.generation_config.no_repeat_ngram_window_size = 128 image1 = load_image( url_to_local_path( "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" @@ -354,7 +343,6 @@ def test_small_model_integration_test_multi_page_document_parsing(self): **inputs, do_sample=False, max_new_tokens=20, - no_repeat_ngram_size=35, ) decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True) From 6ef423ec178c1c60582afd6d546c69cfccc6b7da Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:18:24 +0100 Subject: [PATCH 130/195] update docstring --- .../models/unlimited_ocr/generation_unlimited_ocr.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py index b67046eec230..d139b1ad28a0 100644 --- a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py @@ -53,10 +53,6 @@ class UnlimitedOcrGenerationMixin(GenerationMixin): r""" Adds support for the `no_repeat_ngram_window_size` generation option. If set together with `no_repeat_ngram_size`, n-gram repetitions are blocked only within this many trailing tokens instead of over the whole sequence. - - `no_repeat_ngram_window_size` is a model specific option and must already be present on the generation config - before it can be passed to `generate`, either through the checkpoint's `generation_config.json` or by setting - `model.generation_config.no_repeat_ngram_window_size` after loading the model. """ def _get_logits_processor( From f690cd6f90ff4270d1826d196c6d96378dd83300 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:25:29 +0100 Subject: [PATCH 131/195] remove import --- src/transformers/models/unlimited_ocr/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/transformers/models/unlimited_ocr/__init__.py b/src/transformers/models/unlimited_ocr/__init__.py index 737519feca9a..fbca2b95c1f0 100644 --- a/src/transformers/models/unlimited_ocr/__init__.py +++ b/src/transformers/models/unlimited_ocr/__init__.py @@ -19,7 +19,6 @@ if TYPE_CHECKING: - from .configuration_unlimited_ocr import * from .image_processing_unlimited_ocr import * from .modeling_unlimited_ocr import * from .processing_unlimited_ocr import * From 7196f3494d06308b759e0cf1ef985781126fd682 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:43:52 +0000 Subject: [PATCH 132/195] inherit --- .../image_processing_unlimited_ocr.py | 5 +-- .../unlimited_ocr/modular_unlimited_ocr.py | 44 ------------------- .../test_modeling_unlimited_ocr.py | 4 +- 3 files changed, 4 insertions(+), 49 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index 766da92b49e2..f8f7bf4e30a4 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -363,9 +363,8 @@ def pad_to_square( ) max_dim = max(height, width) - # Use round(x/2) instead of x//2 to match PIL.ImageOps.pad behavior - paste_x_left = round((max_dim - width) / 2) - paste_y_left = round((max_dim - height) / 2) + paste_x_left = (max_dim - width) // 2 + paste_y_left = (max_dim - height) // 2 paste_x_right = max_dim - width - paste_x_left paste_y_right = max_dim - height - paste_y_left padded_images = tvF.pad( diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 1fdddbdbe370..a52f495ff94f 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -16,7 +16,6 @@ import torch from huggingface_hub.dataclasses import strict from torch import nn -from torchvision.transforms.v2 import functional as tvF from ... import initialization as init from ...cache_utils import ( @@ -102,49 +101,6 @@ class UnlimitedOcrImageProcessor(DeepseekOcr2ImageProcessor): max_patches = 32 model_input_names = ["pixel_values", "num_local_patches", "patches_grid"] - def pad_to_square( - self, - images: "torch.Tensor", - background_color: int | tuple[int, int, int] = 0, - ) -> "torch.Tensor": - """ - Pads an image to a square based on the longest edge. - - Args: - images (`torch.Tensor`): - The images to pad. Shape: (batch_size, num_channels, height, width) or (num_channels, height, width). - background_color (`int` or `tuple[int, int, int]`, *optional*, defaults to 0): - The color to use for the padding. Can be an integer for single channel or a - tuple of integers representing for multi-channel images. If passed as integer - in multi-channel mode, it will default to `0` in subsequent channels. - Returns: - `torch.Tensor`: The padded images. - """ - height, width = images.shape[-2:] - - if height == width: - return images - - num_channels = images.shape[1] if len(images.shape) == 4 else images.shape[0] - if isinstance(background_color, int): - background_color = [background_color] + [0] * (num_channels - 1) - elif len(background_color) != num_channels: - raise ValueError( - f"background_color must have no more than {num_channels} elements to match the number of channels" - ) - - max_dim = max(height, width) - # Use round(x/2) instead of x//2 to match PIL.ImageOps.pad behavior - paste_x_left = round((max_dim - width) / 2) - paste_y_left = round((max_dim - height) / 2) - paste_x_right = max_dim - width - paste_x_left - paste_y_right = max_dim - height - paste_y_left - padded_images = tvF.pad( - images, padding=[paste_x_left, paste_y_left, paste_x_right, paste_y_right], fill=background_color - ) - - return padded_images - def _preprocess( self, images: list["torch.Tensor"], diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index b09543804b39..f21c6849240d 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -348,8 +348,8 @@ def test_small_model_integration_test_multi_page_document_parsing(self): EXPECTED_DECODED_TEXT = Expectations( { - ("cuda", None): "image [384, 87, 489, 171]\n", - ("cpu", None): "image [384, 87, 489, 171]\n", + ("cuda", None): "image [382, 87, 489, 174]\n", + ("cpu", None): "image [382, 87, 489, 174]\n", } ).get_expectation() self.assertEqual(decoded, EXPECTED_DECODED_TEXT) From a9c07724fc57322178013c526c9b1f0106628012 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:46:15 +0100 Subject: [PATCH 133/195] update docs --- docs/source/en/model_doc/unlimited_ocr.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 3c67fc24ee9b..ff37e74b0ad0 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -33,7 +33,7 @@ The abstract from the paper is the following: -Unlimited-OCR supports two inference configurations: the default "gundam" mode uses 640x640 tiles with dynamic cropping for high-resolution documents, and "base" mode uses a single 1024x1024 global view for standard-resolution inputs. To enable base mode set `crop_to_patches=False`. +Unlimited-OCR supports two inference configurations: the default "gundam" mode uses 640x640 tiles with dynamic cropping for high-resolution documents, and "base" mode uses a single 1024x1024 global view for standard-resolution inputs. To enable base mode set `crop_to_patches=False` in the processor. The vision tower follows the two-stage approach from [DeepSeek-OCR-2](./deepseek_ocr2): a SAM ViT-B encoder feeds into a CLIP ViT encoder. Unlike DeepSeek-OCR-2, the CLIP features are additionally concatenated with the SAM features to yield the final image tokens. Unlimited-OCR also omits the learnable patch queries from DeepSeek-OCR-2. From 5a96cc1e451e927d462d0956cc04f319e4285253 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 08:47:22 +0100 Subject: [PATCH 134/195] add _prefill --- .../unlimited_ocr/generation_unlimited_ocr.py | 24 ++++++++++ .../unlimited_ocr/modeling_unlimited_ocr.py | 47 ++++++++++--------- .../unlimited_ocr/modular_unlimited_ocr.py | 47 ++++++++++--------- .../test_modeling_unlimited_ocr.py | 20 +++++--- 4 files changed, 89 insertions(+), 49 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py index d139b1ad28a0..0c2474f2d125 100644 --- a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py @@ -55,6 +55,30 @@ class UnlimitedOcrGenerationMixin(GenerationMixin): n-gram repetitions are blocked only within this many trailing tokens instead of over the whole sequence. """ + def _prefill( + self, + input_ids: torch.LongTensor, + generation_config: GenerationConfig, + model_kwargs: dict, + is_first_iteration: bool = True, + ): + outputs = super()._prefill( + input_ids, + generation_config, + model_kwargs, + is_first_iteration=is_first_iteration, + ) + # The cache layers cannot tell prefill and decode apart on their own, so mark the prefill as complete once + # all prefill forward passes are done (a prefill can span several forward passes when chunked). + past_key_values = model_kwargs.get("past_key_values") + if past_key_values is None: + past_key_values = getattr(outputs, "past_key_values", None) + if past_key_values is not None: + for layer in past_key_values.layers: + if layer._layer_type == "reference_sliding_attention": + layer.end_prefill() + return outputs + def _get_logits_processor( self, generation_config: GenerationConfig, diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 2ec738083da3..ff73c3608824 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1282,11 +1282,13 @@ def _init_weights(self, module): class UnlimitedOcrDynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): - """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens before - the first decode step and applies a sliding window to all decoded tokens. + """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens and applies a + sliding window to all decoded tokens. Once ``sliding_window`` decode tokens have accumulated, the oldest decode tokens are evicted and replaced by the most recent ones. The prefill tokens always remain in the cache. + + The layer is in the prefill phase until ``end_prefill`` is called. """ _layer_type = "reference_sliding_attention" @@ -1310,18 +1312,11 @@ def update( full_value_states = torch.cat([self.values, value_states], dim=-2) # Prefill - # We assume kv_length > 1 during prefill and that the first decode step - # has kv_length == 1. This is required to differentiate between prefill - # and decode phases. Later decode steps can have any kv_length. - if self.prefill_length is None and kv_length > 1: + if self.prefill_length is None: self.keys = full_key_states self.values = full_value_states return self.keys, self.values - # First decode step (or an empty prompt that skipped prefill above): mark prefill as complete - if self.prefill_length is None: - self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 - # Cache growing if self.cumulative_length <= self.prefill_length + self.sliding_window - 1: self.keys = full_key_states @@ -1373,6 +1368,15 @@ def get_max_length(self) -> int: return -1 return self.prefill_length + self.sliding_window + def end_prefill(self) -> None: + """Mark the prefill phase as complete.""" + if self.prefill_length is None: + self.prefill_length = self.get_seq_length() + + def reset(self) -> None: + super().reset() + self.prefill_length = None + def crop(self, max_length: int) -> None: """ Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be @@ -1398,6 +1402,8 @@ class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): that are never eviced from the cache. Slots ``[prefill_length, prefill_length + sliding_window)`` hold the sliding window decode slots where the oldest entries are always replaced by the newest ones. + The layer is in the prefill phase until ``end_prefill`` is called. + Args: max_cache_len (`int`): Maximum number of tokens that can be stored, used for tensor preallocation. @@ -1435,10 +1441,7 @@ def update( kv_length = key_states.shape[-2] # Prefill - # We assume kv_length > 1 during prefill and that the first decode step - # has kv_length == 1. This is required to differentiate between prefill - # and decode phases. Later decode steps can have any kv_length. - if self.prefill_length is None and kv_length > 1: + if self.prefill_length is None: # Chunked prefill: the buffer was sized for the first chunk only, so grow it to hold all prefill seen so far plus this chunk plus the reserved window. required_length = min(self.max_cache_len, self.cumulative_length_int + kv_length + self.sliding_window) if self.keys.shape[-2] < required_length: @@ -1461,10 +1464,6 @@ def update( # Very important to return the `self` tensors here, as they have the static dynamo address return self.keys, self.values - # First decode step (or an empty prompt that skipped prefill above): mark prefill as complete - if self.prefill_length is None: - self.prefill_length = self.cumulative_length_int - # Everything below mirrors `StaticSlidingWindowLayer.update` with the difference that the sliding window # is in the `sliding_window` slots right after the prefill instead of spanning the whole buffer. window_start = self.prefill_length @@ -1559,8 +1558,8 @@ def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tens if not is_torchdynamo_compiling(): torch._dynamo.mark_static_address(self.cumulative_length) - prefill_seen = key_states.shape[-2] if key_states.shape[-2] > 1 else 0 - self._allocate_key_value_buffers(min(self.max_cache_len, prefill_seen + self.sliding_window)) + prefill_length = key_states.shape[-2] if self.prefill_length is None else self.prefill_length + self._allocate_key_value_buffers(min(self.max_cache_len, prefill_length + self.sliding_window)) self.is_initialized = True @@ -1572,7 +1571,7 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: kv_offset = 0 # Prefill - if self.prefill_length is None and query_length > 1: + if self.prefill_length is None: kv_length = min(self.max_cache_len, self.cumulative_length_int + query_length + self.sliding_window) # Decode: cache is already full elif is_full: @@ -1596,6 +1595,11 @@ def get_max_length(self) -> int: return self.max_cache_len return min(self.max_cache_len, self.prefill_length + self.sliding_window) + def end_prefill(self) -> None: + """Mark the prefill phase as complete.""" + if self.prefill_length is None: + self.prefill_length = self.cumulative_length_int + def reset(self) -> None: super().reset() self.prefill_length = None @@ -1697,6 +1701,7 @@ def create_reference_sliding_window_causal_mask(**kwargs): kv_offset = 0 else: layer = next(layer for layer in past_key_values.layers if layer._layer_type == "reference_sliding_attention") + # A layer that is still prefilling has no reference states yet, all its states are prefill states prefill_length = float("inf") if layer.prefill_length is None else layer.prefill_length _, kv_offset = layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 1fdddbdbe370..6f06e3f7fdd8 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -661,11 +661,13 @@ class UnlimitedOcrTextPreTrainedModel(DeepseekOcr2TextPreTrainedModel): class UnlimitedOcrDynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): - """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens before - the first decode step and applies a sliding window to all decoded tokens. + """Reference sliding-window attention (R-SWA) cache layer that keeps all prefill tokens and applies a + sliding window to all decoded tokens. Once ``sliding_window`` decode tokens have accumulated, the oldest decode tokens are evicted and replaced by the most recent ones. The prefill tokens always remain in the cache. + + The layer is in the prefill phase until ``end_prefill`` is called. """ _layer_type = "reference_sliding_attention" @@ -689,18 +691,11 @@ def update( full_value_states = torch.cat([self.values, value_states], dim=-2) # Prefill - # We assume kv_length > 1 during prefill and that the first decode step - # has kv_length == 1. This is required to differentiate between prefill - # and decode phases. Later decode steps can have any kv_length. - if self.prefill_length is None and kv_length > 1: + if self.prefill_length is None: self.keys = full_key_states self.values = full_value_states return self.keys, self.values - # First decode step (or an empty prompt that skipped prefill above): mark prefill as complete - if self.prefill_length is None: - self.prefill_length = self.keys.shape[-2] if self.keys.dim() > 1 else 0 - # Cache growing if self.cumulative_length <= self.prefill_length + self.sliding_window - 1: self.keys = full_key_states @@ -752,6 +747,15 @@ def get_max_length(self) -> int: return -1 return self.prefill_length + self.sliding_window + def end_prefill(self) -> None: + """Mark the prefill phase as complete.""" + if self.prefill_length is None: + self.prefill_length = self.get_seq_length() + + def reset(self) -> None: + super().reset() + self.prefill_length = None + def crop(self, max_length: int) -> None: """ Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be @@ -777,6 +781,8 @@ class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): that are never eviced from the cache. Slots ``[prefill_length, prefill_length + sliding_window)`` hold the sliding window decode slots where the oldest entries are always replaced by the newest ones. + The layer is in the prefill phase until ``end_prefill`` is called. + Args: max_cache_len (`int`): Maximum number of tokens that can be stored, used for tensor preallocation. @@ -814,10 +820,7 @@ def update( kv_length = key_states.shape[-2] # Prefill - # We assume kv_length > 1 during prefill and that the first decode step - # has kv_length == 1. This is required to differentiate between prefill - # and decode phases. Later decode steps can have any kv_length. - if self.prefill_length is None and kv_length > 1: + if self.prefill_length is None: # Chunked prefill: the buffer was sized for the first chunk only, so grow it to hold all prefill seen so far plus this chunk plus the reserved window. required_length = min(self.max_cache_len, self.cumulative_length_int + kv_length + self.sliding_window) if self.keys.shape[-2] < required_length: @@ -840,10 +843,6 @@ def update( # Very important to return the `self` tensors here, as they have the static dynamo address return self.keys, self.values - # First decode step (or an empty prompt that skipped prefill above): mark prefill as complete - if self.prefill_length is None: - self.prefill_length = self.cumulative_length_int - # Everything below mirrors `StaticSlidingWindowLayer.update` with the difference that the sliding window # is in the `sliding_window` slots right after the prefill instead of spanning the whole buffer. window_start = self.prefill_length @@ -938,8 +937,8 @@ def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tens if not is_torchdynamo_compiling(): torch._dynamo.mark_static_address(self.cumulative_length) - prefill_seen = key_states.shape[-2] if key_states.shape[-2] > 1 else 0 - self._allocate_key_value_buffers(min(self.max_cache_len, prefill_seen + self.sliding_window)) + prefill_length = key_states.shape[-2] if self.prefill_length is None else self.prefill_length + self._allocate_key_value_buffers(min(self.max_cache_len, prefill_length + self.sliding_window)) self.is_initialized = True @@ -951,7 +950,7 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: kv_offset = 0 # Prefill - if self.prefill_length is None and query_length > 1: + if self.prefill_length is None: kv_length = min(self.max_cache_len, self.cumulative_length_int + query_length + self.sliding_window) # Decode: cache is already full elif is_full: @@ -975,6 +974,11 @@ def get_max_length(self) -> int: return self.max_cache_len return min(self.max_cache_len, self.prefill_length + self.sliding_window) + def end_prefill(self) -> None: + """Mark the prefill phase as complete.""" + if self.prefill_length is None: + self.prefill_length = self.cumulative_length_int + def reset(self) -> None: super().reset() self.prefill_length = None @@ -1011,6 +1015,7 @@ def create_reference_sliding_window_causal_mask(**kwargs): kv_offset = 0 else: layer = next(layer for layer in past_key_values.layers if layer._layer_type == "reference_sliding_attention") + # A layer that is still prefilling has no reference states yet, all its states are prefill states prefill_length = float("inf") if layer.prefill_length is None else layer.prefill_length _, kv_offset = layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index b09543804b39..3dfa6919af4e 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -160,10 +160,7 @@ def _check_generate_cache_sliding_window_too_small(self, cache_implementation: s out_reference = model.generate(**inputs_dict, max_new_tokens=15, do_sample=False) - # Prefill the cache. - # As the last decoded token from a .generate call isn't added to the cache we have to - # decode at least 2 tokens for the cache update call to mark the prefill as complete. - # Prefill is only marked as complete once a single decoded token is added to the cache (kv_length == 1). + # Prefill the cache out_prefill = model.generate( **inputs_dict, max_new_tokens=prefill_max_new_tokens, @@ -184,15 +181,24 @@ def _check_generate_cache_sliding_window_too_small(self, cache_implementation: s past_key_values=out_prefill.past_key_values, max_new_tokens=3, do_sample=False, + return_dict_in_generate=True, use_cache=True, ) - self.assertEqual(out.tolist(), out_reference.tolist()) + self.assertEqual(out.sequences.tolist(), out_reference.tolist()) + + prompt_length = inputs_dict["input_ids"].shape[1] + reference_layers = [ + layer for layer in out.past_key_values.layers if layer._layer_type == "reference_sliding_attention" + ] + self.assertGreater(len(reference_layers), 0) + for layer in reference_layers: + self.assertEqual(layer.prefill_length, prompt_length) def test_generate_dynamic_cache_sliding_window_too_small(self): - self._check_generate_cache_sliding_window_too_small(cache_implementation="dynamic", prefill_max_new_tokens=2) + self._check_generate_cache_sliding_window_too_small(cache_implementation="dynamic", prefill_max_new_tokens=1) def test_generate_static_cache_sliding_window_too_small(self): - self._check_generate_cache_sliding_window_too_small(cache_implementation="static", prefill_max_new_tokens=2) + self._check_generate_cache_sliding_window_too_small(cache_implementation="static", prefill_max_new_tokens=1) def test_generate_dynamic_cache_sliding_window_too_small_cache_full(self): """Continue from full cache""" From 4e743808e65c193a5af9829cb18935375c50ac9a Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 09:00:59 +0100 Subject: [PATCH 135/195] update max_new_tokens --- docs/source/en/model_doc/unlimited_ocr.md | 6 +++--- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 2 +- .../models/unlimited_ocr/modular_unlimited_ocr.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index ff37e74b0ad0..c3b79ed281db 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -59,7 +59,7 @@ inputs = processor(images=image, text="document parsing.", return_tensors output = model.generate( **inputs, - max_new_tokens=4096, + max_new_tokens=32768, no_repeat_ngram_size=35, no_repeat_ngram_window_size=128, ) @@ -91,7 +91,7 @@ inputs = processor( output = model.generate( **inputs, - max_new_tokens=4096, + max_new_tokens=32768, no_repeat_ngram_size=35, no_repeat_ngram_window_size=128, ) @@ -117,7 +117,7 @@ inputs = processor(images=image, text="document parsing.", return_tensors output = model.generate( **inputs, - max_new_tokens=4096, + max_new_tokens=32768, no_repeat_ngram_size=35, no_repeat_ngram_window_size=128, ) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index ff73c3608824..1c32a1ecc28b 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -2057,7 +2057,7 @@ def forward( >>> output = model.generate( ... **inputs, - ... max_new_tokens=4096, + ... max_new_tokens=32768, ... no_repeat_ngram_size=35, ... no_repeat_ngram_window_size=128, ... ) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index d097aa0262c3..af259e50a291 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -1268,7 +1268,7 @@ def forward( >>> output = model.generate( ... **inputs, - ... max_new_tokens=4096, + ... max_new_tokens=32768, ... no_repeat_ngram_size=35, ... no_repeat_ngram_window_size=128, ... ) From ccf90daca59051bba6102f3be361b250eea985b3 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 09:54:56 +0100 Subject: [PATCH 136/195] Add return_detections --- docs/source/en/model_doc/unlimited_ocr.md | 10 +-- .../unlimited_ocr/modeling_unlimited_ocr.py | 2 +- .../unlimited_ocr/modular_unlimited_ocr.py | 62 +++++++++++++++++++ .../unlimited_ocr/processing_unlimited_ocr.py | 62 +++++++++++++++++++ .../test_modeling_unlimited_ocr.py | 12 +++- 5 files changed, 141 insertions(+), 7 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index c3b79ed281db..7230cdff1c02 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -103,7 +103,7 @@ processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=T ### Region detections -Set `skip_special_tokens=False` to wrap all detections and region types in `<|det|>...<|/det|>` markers. This is useful for further post-processing of the output, for example to plot the detected bounding boxes on the image. Each detection is wrapped as `<|det|>region_type [x1, y1, x2, y2]<|/det|>` with coordinates normalized to a `[0, 999]` range. Parse the markers with a regular expression and rescale the coordinates to the image size to plot the bounding boxes. +Set `skip_special_tokens=False` to wrap all detections and region types in `<|det|>...<|/det|>` markers. This is useful for further post-processing of the output, for example to plot the detected bounding boxes on the image. Each detection is wrapped as `<|det|>region_type [x1, y1, x2, y2]<|/det|>text...` with coordinates normalized to a `[0, 999]` range. Set `return_detections=True` to get an additional list of dictionaries with all detections parsed as `{"region_type": region_type, "box": [x1, y1, x2, y2], "text": "..."}`. ```python import re @@ -121,11 +121,9 @@ output = model.generate( no_repeat_ngram_size=35, no_repeat_ngram_window_size=128, ) -decoded = processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=False) +decoded, detections = processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=False, return_detections=True) # <|det|>image [383, 87, 497, 171]<|/det|>\n<|det|>text [333, 201, 558, 230]<|/det|>R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... -detections = re.findall(r"<\|det\|>(\S+) \[(\d+), (\d+), (\d+), (\d+)\]<\|/det\|>", decoded) - # Visualization import random import matplotlib.pyplot as plt @@ -137,7 +135,9 @@ width, height = image.size figure, axis = plt.subplots(figsize=(10, 12)) axis.imshow(image) -for region_type, x1, y1, x2, y2 in detections: +for detection in detections: + region_type = detection["region_type"] + x1, y1, x2, y2 = detection["box"] x1, y1, x2, y2 = int(x1) / 999 * width, int(y1) / 999 * height, int(x2) / 999 * width, int(y2) / 999 * height color = (random.random(), random.random(), random.random()) rectangle = patches.Rectangle((x1, y1), x2 - x1, y2 - y1, linewidth=1.5, edgecolor=color, facecolor="none") diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 1c32a1ecc28b..ff73c3608824 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -2057,7 +2057,7 @@ def forward( >>> output = model.generate( ... **inputs, - ... max_new_tokens=32768, + ... max_new_tokens=4096, ... no_repeat_ngram_size=35, ... no_repeat_ngram_window_size=128, ... ) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index af259e50a291..c19dca88a2de 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import math +import re import torch from huggingface_hub.dataclasses import strict @@ -224,6 +225,67 @@ def replace_image_token(self, image_inputs: dict, image_idx: int) -> TextInput: num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) return self.image_token * num_tokens + def batch_decode(self, *args, return_detections: bool = False, **kwargs): + """ + This method forwards all its arguments to PreTrainedTokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please refer to + the docstring of this method for more information. + + Args: + return_detections (`bool`, *optional*, defaults to `False`): + Whether or not to also return the layout detections parsed from the decoded text. + + Returns: + `str` or `tuple[str, list[dict]]`: The decoded text. If `return_detections` is `True`, a tuple of the + decoded text and a list of detections, where every detection is a dictionary with the keys `region_type`, + `box` and `text`. Boxes are in [x1, y1, x2, y2] format with coordinates normalized to [0, 999]. + """ + if not hasattr(self, "tokenizer"): + raise ValueError(f"Cannot batch decode text: {self.__class__.__name__} has no tokenizer.") + decoded = self.tokenizer.batch_decode(*args, **kwargs) + if return_detections: + detections = self._parse_detections(decoded) + return decoded, detections + return decoded + + def decode(self, *args, return_detections: bool = False, **kwargs): + """ + This method forwards all its arguments to PreTrainedTokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to + the docstring of this method for more information. + + Args: + return_detections (`bool`, *optional*, defaults to `False`): + Whether or not to also return the layout detections parsed from the decoded text. + + Returns: + `str` or `tuple[str, list[dict]]`: The decoded text. If `return_detections` is `True`, a tuple of the + decoded text and a list of detections, where every detection is a dictionary with the keys `region_type`, + `box` and `text`. Boxes are in [x1, y1, x2, y2] format with coordinates normalized to [0, 999]. + """ + if not hasattr(self, "tokenizer"): + raise ValueError(f"Cannot decode text: {self.__class__.__name__} has no tokenizer.") + decoded = self.tokenizer.decode(*args, **kwargs) + if return_detections: + detections = self._parse_detections(decoded) + return decoded, detections + return decoded + + def _parse_detections(self, decoded: str) -> list[dict]: + matches = re.findall( + r"<\|det\|>(\S+) \[(\d+), (\d+), (\d+), (\d+)\]<\|/det\|>(.*?)(?=<\|det\|>||\Z)", + decoded, + flags=re.DOTALL, + ) + detections = [] + for region_type, x1, y1, x2, y2, text in matches: + detections.append( + { + "region_type": region_type, + "box": [x1, y1, x2, y2], + "text": text, + } + ) + return detections + @auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index 1f8ca95c6404..a6d5551574ac 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -18,6 +18,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import math +import re from ...image_utils import make_flat_list_of_images from ...processing_utils import ProcessingKwargs, ProcessorMixin @@ -83,5 +84,66 @@ def replace_image_token(self, image_inputs: dict, image_idx: int) -> TextInput: num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) return self.image_token * num_tokens + def batch_decode(self, *args, return_detections: bool = False, **kwargs): + """ + This method forwards all its arguments to PreTrainedTokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please refer to + the docstring of this method for more information. + + Args: + return_detections (`bool`, *optional*, defaults to `False`): + Whether or not to also return the layout detections parsed from the decoded text. + + Returns: + `str` or `tuple[str, list[dict]]`: The decoded text. If `return_detections` is `True`, a tuple of the + decoded text and a list of detections, where every detection is a dictionary with the keys `region_type`, + `box` and `text`. Boxes are in [x1, y1, x2, y2] format with coordinates normalized to [0, 999]. + """ + if not hasattr(self, "tokenizer"): + raise ValueError(f"Cannot batch decode text: {self.__class__.__name__} has no tokenizer.") + decoded = self.tokenizer.batch_decode(*args, **kwargs) + if return_detections: + detections = self._parse_detections(decoded) + return decoded, detections + return decoded + + def decode(self, *args, return_detections: bool = False, **kwargs): + """ + This method forwards all its arguments to PreTrainedTokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to + the docstring of this method for more information. + + Args: + return_detections (`bool`, *optional*, defaults to `False`): + Whether or not to also return the layout detections parsed from the decoded text. + + Returns: + `str` or `tuple[str, list[dict]]`: The decoded text. If `return_detections` is `True`, a tuple of the + decoded text and a list of detections, where every detection is a dictionary with the keys `region_type`, + `box` and `text`. Boxes are in [x1, y1, x2, y2] format with coordinates normalized to [0, 999]. + """ + if not hasattr(self, "tokenizer"): + raise ValueError(f"Cannot decode text: {self.__class__.__name__} has no tokenizer.") + decoded = self.tokenizer.decode(*args, **kwargs) + if return_detections: + detections = self._parse_detections(decoded) + return decoded, detections + return decoded + + def _parse_detections(self, decoded: str) -> list[dict]: + matches = re.findall( + r"<\|det\|>(\S+) \[(\d+), (\d+), (\d+), (\d+)\]<\|/det\|>(.*?)(?=<\|det\|>||\Z)", + decoded, + flags=re.DOTALL, + ) + detections = [] + for region_type, x1, y1, x2, y2, text in matches: + detections.append( + { + "region_type": region_type, + "box": [x1, y1, x2, y2], + "text": text, + } + ) + return detections + __all__ = ["UnlimitedOcrProcessor"] diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 6c02253a9f41..e5e2d91e48bf 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -266,7 +266,9 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se do_sample=False, max_new_tokens=20, ) - decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=False) + decoded, detections = self.processor.decode( + generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=False, return_detections=True + ) EXPECTED_DECODED_TEXT = Expectations( { ("cuda", None): "<|det|>image [383, 87, 497, 171]<|/det|>\n<|det|>title [333", @@ -275,6 +277,14 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se ).get_expectation() # fmt: skip self.assertEqual(decoded, EXPECTED_DECODED_TEXT) + EXPECTED_DETECTIONS = Expectations( + { + ("cuda", None): [{"region_type": "image", "box": [383, 87, 497, 171], "text": "\n"}], + ("cpu", None): [{"region_type": "image", "box": [383, 87, 497, 171], "text": "\n"}], + } + ).get_expectation() # fmt: skip + self.assertEqual(detections, EXPECTED_DETECTIONS) + @slow @require_torch_accelerator def test_small_model_integration_test_document_parsing_batched(self): From 527a69d50efdf57451214d0fd5f967f0d3d0f371 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 09:54:56 +0100 Subject: [PATCH 137/195] Add return_detections --- docs/source/en/model_doc/unlimited_ocr.md | 10 +-- .../unlimited_ocr/modular_unlimited_ocr.py | 62 +++++++++++++++++++ .../unlimited_ocr/processing_unlimited_ocr.py | 62 +++++++++++++++++++ .../test_modeling_unlimited_ocr.py | 12 +++- 4 files changed, 140 insertions(+), 6 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index c3b79ed281db..7230cdff1c02 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -103,7 +103,7 @@ processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=T ### Region detections -Set `skip_special_tokens=False` to wrap all detections and region types in `<|det|>...<|/det|>` markers. This is useful for further post-processing of the output, for example to plot the detected bounding boxes on the image. Each detection is wrapped as `<|det|>region_type [x1, y1, x2, y2]<|/det|>` with coordinates normalized to a `[0, 999]` range. Parse the markers with a regular expression and rescale the coordinates to the image size to plot the bounding boxes. +Set `skip_special_tokens=False` to wrap all detections and region types in `<|det|>...<|/det|>` markers. This is useful for further post-processing of the output, for example to plot the detected bounding boxes on the image. Each detection is wrapped as `<|det|>region_type [x1, y1, x2, y2]<|/det|>text...` with coordinates normalized to a `[0, 999]` range. Set `return_detections=True` to get an additional list of dictionaries with all detections parsed as `{"region_type": region_type, "box": [x1, y1, x2, y2], "text": "..."}`. ```python import re @@ -121,11 +121,9 @@ output = model.generate( no_repeat_ngram_size=35, no_repeat_ngram_window_size=128, ) -decoded = processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=False) +decoded, detections = processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=False, return_detections=True) # <|det|>image [383, 87, 497, 171]<|/det|>\n<|det|>text [333, 201, 558, 230]<|/det|>R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... -detections = re.findall(r"<\|det\|>(\S+) \[(\d+), (\d+), (\d+), (\d+)\]<\|/det\|>", decoded) - # Visualization import random import matplotlib.pyplot as plt @@ -137,7 +135,9 @@ width, height = image.size figure, axis = plt.subplots(figsize=(10, 12)) axis.imshow(image) -for region_type, x1, y1, x2, y2 in detections: +for detection in detections: + region_type = detection["region_type"] + x1, y1, x2, y2 = detection["box"] x1, y1, x2, y2 = int(x1) / 999 * width, int(y1) / 999 * height, int(x2) / 999 * width, int(y2) / 999 * height color = (random.random(), random.random(), random.random()) rectangle = patches.Rectangle((x1, y1), x2 - x1, y2 - y1, linewidth=1.5, edgecolor=color, facecolor="none") diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index af259e50a291..c19dca88a2de 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import math +import re import torch from huggingface_hub.dataclasses import strict @@ -224,6 +225,67 @@ def replace_image_token(self, image_inputs: dict, image_idx: int) -> TextInput: num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) return self.image_token * num_tokens + def batch_decode(self, *args, return_detections: bool = False, **kwargs): + """ + This method forwards all its arguments to PreTrainedTokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please refer to + the docstring of this method for more information. + + Args: + return_detections (`bool`, *optional*, defaults to `False`): + Whether or not to also return the layout detections parsed from the decoded text. + + Returns: + `str` or `tuple[str, list[dict]]`: The decoded text. If `return_detections` is `True`, a tuple of the + decoded text and a list of detections, where every detection is a dictionary with the keys `region_type`, + `box` and `text`. Boxes are in [x1, y1, x2, y2] format with coordinates normalized to [0, 999]. + """ + if not hasattr(self, "tokenizer"): + raise ValueError(f"Cannot batch decode text: {self.__class__.__name__} has no tokenizer.") + decoded = self.tokenizer.batch_decode(*args, **kwargs) + if return_detections: + detections = self._parse_detections(decoded) + return decoded, detections + return decoded + + def decode(self, *args, return_detections: bool = False, **kwargs): + """ + This method forwards all its arguments to PreTrainedTokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to + the docstring of this method for more information. + + Args: + return_detections (`bool`, *optional*, defaults to `False`): + Whether or not to also return the layout detections parsed from the decoded text. + + Returns: + `str` or `tuple[str, list[dict]]`: The decoded text. If `return_detections` is `True`, a tuple of the + decoded text and a list of detections, where every detection is a dictionary with the keys `region_type`, + `box` and `text`. Boxes are in [x1, y1, x2, y2] format with coordinates normalized to [0, 999]. + """ + if not hasattr(self, "tokenizer"): + raise ValueError(f"Cannot decode text: {self.__class__.__name__} has no tokenizer.") + decoded = self.tokenizer.decode(*args, **kwargs) + if return_detections: + detections = self._parse_detections(decoded) + return decoded, detections + return decoded + + def _parse_detections(self, decoded: str) -> list[dict]: + matches = re.findall( + r"<\|det\|>(\S+) \[(\d+), (\d+), (\d+), (\d+)\]<\|/det\|>(.*?)(?=<\|det\|>||\Z)", + decoded, + flags=re.DOTALL, + ) + detections = [] + for region_type, x1, y1, x2, y2, text in matches: + detections.append( + { + "region_type": region_type, + "box": [x1, y1, x2, y2], + "text": text, + } + ) + return detections + @auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index 1f8ca95c6404..a6d5551574ac 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -18,6 +18,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import math +import re from ...image_utils import make_flat_list_of_images from ...processing_utils import ProcessingKwargs, ProcessorMixin @@ -83,5 +84,66 @@ def replace_image_token(self, image_inputs: dict, image_idx: int) -> TextInput: num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) return self.image_token * num_tokens + def batch_decode(self, *args, return_detections: bool = False, **kwargs): + """ + This method forwards all its arguments to PreTrainedTokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please refer to + the docstring of this method for more information. + + Args: + return_detections (`bool`, *optional*, defaults to `False`): + Whether or not to also return the layout detections parsed from the decoded text. + + Returns: + `str` or `tuple[str, list[dict]]`: The decoded text. If `return_detections` is `True`, a tuple of the + decoded text and a list of detections, where every detection is a dictionary with the keys `region_type`, + `box` and `text`. Boxes are in [x1, y1, x2, y2] format with coordinates normalized to [0, 999]. + """ + if not hasattr(self, "tokenizer"): + raise ValueError(f"Cannot batch decode text: {self.__class__.__name__} has no tokenizer.") + decoded = self.tokenizer.batch_decode(*args, **kwargs) + if return_detections: + detections = self._parse_detections(decoded) + return decoded, detections + return decoded + + def decode(self, *args, return_detections: bool = False, **kwargs): + """ + This method forwards all its arguments to PreTrainedTokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to + the docstring of this method for more information. + + Args: + return_detections (`bool`, *optional*, defaults to `False`): + Whether or not to also return the layout detections parsed from the decoded text. + + Returns: + `str` or `tuple[str, list[dict]]`: The decoded text. If `return_detections` is `True`, a tuple of the + decoded text and a list of detections, where every detection is a dictionary with the keys `region_type`, + `box` and `text`. Boxes are in [x1, y1, x2, y2] format with coordinates normalized to [0, 999]. + """ + if not hasattr(self, "tokenizer"): + raise ValueError(f"Cannot decode text: {self.__class__.__name__} has no tokenizer.") + decoded = self.tokenizer.decode(*args, **kwargs) + if return_detections: + detections = self._parse_detections(decoded) + return decoded, detections + return decoded + + def _parse_detections(self, decoded: str) -> list[dict]: + matches = re.findall( + r"<\|det\|>(\S+) \[(\d+), (\d+), (\d+), (\d+)\]<\|/det\|>(.*?)(?=<\|det\|>||\Z)", + decoded, + flags=re.DOTALL, + ) + detections = [] + for region_type, x1, y1, x2, y2, text in matches: + detections.append( + { + "region_type": region_type, + "box": [x1, y1, x2, y2], + "text": text, + } + ) + return detections + __all__ = ["UnlimitedOcrProcessor"] diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 6c02253a9f41..e5e2d91e48bf 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -266,7 +266,9 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se do_sample=False, max_new_tokens=20, ) - decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=False) + decoded, detections = self.processor.decode( + generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=False, return_detections=True + ) EXPECTED_DECODED_TEXT = Expectations( { ("cuda", None): "<|det|>image [383, 87, 497, 171]<|/det|>\n<|det|>title [333", @@ -275,6 +277,14 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se ).get_expectation() # fmt: skip self.assertEqual(decoded, EXPECTED_DECODED_TEXT) + EXPECTED_DETECTIONS = Expectations( + { + ("cuda", None): [{"region_type": "image", "box": [383, 87, 497, 171], "text": "\n"}], + ("cpu", None): [{"region_type": "image", "box": [383, 87, 497, 171], "text": "\n"}], + } + ).get_expectation() # fmt: skip + self.assertEqual(detections, EXPECTED_DETECTIONS) + @slow @require_torch_accelerator def test_small_model_integration_test_document_parsing_batched(self): From 060196563c04387b01f40b3659e46291738a212a Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 10:16:31 +0100 Subject: [PATCH 138/195] cleanup _preprocess --- .../image_processing_unlimited_ocr.py | 72 +++++++++---------- .../unlimited_ocr/modular_unlimited_ocr.py | 72 +++++++++---------- 2 files changed, 64 insertions(+), 80 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index f8f7bf4e30a4..7d78f7048f76 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -233,40 +233,48 @@ def _preprocess( ) -> BatchFeature: # --- Local patches (batched by shape group) --- local_patches_grouped = {} + patches_grid_grouped = {} + num_local_patches_grouped = {} - if crop_to_patches: - grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + for shape, stacked_images in grouped_images.items(): + height, width = shape[-2:] + num_images = stacked_images.shape[0] + if crop_to_patches and max(height, width) > tile_size: + num_columns, num_rows = get_optimal_tiled_canvas( + (height, width), (tile_size, tile_size), min_patches, max_patches + ) + stacked_patches, num_patches = self.crop_image_to_patches( + stacked_images, + min_patches=min_patches, + max_patches=max_patches, + tile_size=tile_size, + resample=resample, + ) + flat_patches = stacked_patches.reshape(-1, *stacked_patches.shape[2:]) + flat_patches = self.rescale_and_normalize( + flat_patches, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + local_patches_grouped[shape] = flat_patches.reshape(stacked_patches.shape) + else: + num_columns, num_rows, num_patches = 1, 1, 0 + local_patches_grouped[shape] = [None] * num_images + patches_grid_grouped[shape] = [[num_columns, num_rows]] * num_images + num_local_patches_grouped[shape] = [num_patches] * num_images - for shape, stacked_images in grouped_images.items(): - height, width = shape[-2:] - if max(height, width) > tile_size: - stacked_patches, _ = self.crop_image_to_patches( - stacked_images, - min_patches=min_patches, - max_patches=max_patches, - tile_size=tile_size, - resample=resample, - ) - flat_patches = stacked_patches.reshape(-1, *stacked_patches.shape[2:]) - flat_patches = self.rescale_and_normalize( - flat_patches, do_rescale, rescale_factor, do_normalize, image_mean, image_std - ) - local_patches_grouped[shape] = flat_patches.reshape(stacked_patches.shape) - else: - local_patches_grouped[shape] = [None] * stacked_images.shape[0] - - ordered_local = reorder_images(local_patches_grouped, grouped_images_index) - else: - ordered_local = [] + ordered_local = reorder_images(local_patches_grouped, grouped_images_index) + patches_grid = reorder_images(patches_grid_grouped, grouped_images_index) + num_local_patches = reorder_images(num_local_patches_grouped, grouped_images_index) flat_local_list = [patch for item in ordered_local if item is not None for patch in item] # --- Global view (batched by shape group) --- + # Different from DeepseekOcr2 which uses size.height or tile_size global_target_size = max(size.height, size.width) - grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) processed_global_grouped = {} for shape, stacked in grouped_images.items(): + # Different from DeepseekOcr2 which crops and pads all images if not crop_to_patches and global_target_size <= pad_if_larger_than: stacked = self.resize( stacked, SizeDict(height=global_target_size, width=global_target_size), resample=resample @@ -287,22 +295,6 @@ def _preprocess( data = {"pixel_values": all_pixel_values_global} if flat_local_list: data["pixel_values_local"] = flat_local_list - - # Compute per-image spatial crop grid and local-patch counts. - patches_grid = [] - num_local_patches = [] - for image in images: - height, width = image.shape[-2:] - if crop_to_patches and max(height, width) > tile_size: - num_columns, num_rows = get_optimal_tiled_canvas( - (height, width), (tile_size, tile_size), min_patches, max_patches - ) - num_local_patches.append(num_columns * num_rows) - else: - num_columns, num_rows = 1, 1 - num_local_patches.append(0) - patches_grid.append([num_columns, num_rows]) - data["num_local_patches"] = num_local_patches data["patches_grid"] = patches_grid diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index c19dca88a2de..fae2a05d91b1 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -123,40 +123,48 @@ def _preprocess( ) -> BatchFeature: # --- Local patches (batched by shape group) --- local_patches_grouped = {} + patches_grid_grouped = {} + num_local_patches_grouped = {} - if crop_to_patches: - grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + for shape, stacked_images in grouped_images.items(): + height, width = shape[-2:] + num_images = stacked_images.shape[0] + if crop_to_patches and max(height, width) > tile_size: + num_columns, num_rows = get_optimal_tiled_canvas( + (height, width), (tile_size, tile_size), min_patches, max_patches + ) + stacked_patches, num_patches = self.crop_image_to_patches( + stacked_images, + min_patches=min_patches, + max_patches=max_patches, + tile_size=tile_size, + resample=resample, + ) + flat_patches = stacked_patches.reshape(-1, *stacked_patches.shape[2:]) + flat_patches = self.rescale_and_normalize( + flat_patches, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + local_patches_grouped[shape] = flat_patches.reshape(stacked_patches.shape) + else: + num_columns, num_rows, num_patches = 1, 1, 0 + local_patches_grouped[shape] = [None] * num_images + patches_grid_grouped[shape] = [[num_columns, num_rows]] * num_images + num_local_patches_grouped[shape] = [num_patches] * num_images - for shape, stacked_images in grouped_images.items(): - height, width = shape[-2:] - if max(height, width) > tile_size: - stacked_patches, _ = self.crop_image_to_patches( - stacked_images, - min_patches=min_patches, - max_patches=max_patches, - tile_size=tile_size, - resample=resample, - ) - flat_patches = stacked_patches.reshape(-1, *stacked_patches.shape[2:]) - flat_patches = self.rescale_and_normalize( - flat_patches, do_rescale, rescale_factor, do_normalize, image_mean, image_std - ) - local_patches_grouped[shape] = flat_patches.reshape(stacked_patches.shape) - else: - local_patches_grouped[shape] = [None] * stacked_images.shape[0] - - ordered_local = reorder_images(local_patches_grouped, grouped_images_index) - else: - ordered_local = [] + ordered_local = reorder_images(local_patches_grouped, grouped_images_index) + patches_grid = reorder_images(patches_grid_grouped, grouped_images_index) + num_local_patches = reorder_images(num_local_patches_grouped, grouped_images_index) flat_local_list = [patch for item in ordered_local if item is not None for patch in item] # --- Global view (batched by shape group) --- + # Different from DeepseekOcr2 which uses size.height or tile_size global_target_size = max(size.height, size.width) - grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) processed_global_grouped = {} for shape, stacked in grouped_images.items(): + # Different from DeepseekOcr2 which crops and pads all images if not crop_to_patches and global_target_size <= pad_if_larger_than: stacked = self.resize( stacked, SizeDict(height=global_target_size, width=global_target_size), resample=resample @@ -177,22 +185,6 @@ def _preprocess( data = {"pixel_values": all_pixel_values_global} if flat_local_list: data["pixel_values_local"] = flat_local_list - - # Compute per-image spatial crop grid and local-patch counts. - patches_grid = [] - num_local_patches = [] - for image in images: - height, width = image.shape[-2:] - if crop_to_patches and max(height, width) > tile_size: - num_columns, num_rows = get_optimal_tiled_canvas( - (height, width), (tile_size, tile_size), min_patches, max_patches - ) - num_local_patches.append(num_columns * num_rows) - else: - num_columns, num_rows = 1, 1 - num_local_patches.append(0) - patches_grid.append([num_columns, num_rows]) - data["num_local_patches"] = num_local_patches data["patches_grid"] = patches_grid From 6a8189723d4e3b0bb3043765580d8ade61609fa1 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 10:19:23 +0100 Subject: [PATCH 139/195] remove total=False --- src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index fae2a05d91b1..dc35e88b18b5 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -194,7 +194,7 @@ def _preprocess( ) -class UnlimitedOcrProcessorKwargs(DeepseekOcr2ProcessorKwargs, total=False): +class UnlimitedOcrProcessorKwargs(DeepseekOcr2ProcessorKwargs): _defaults = { "text_kwargs": { "padding": False, From bb13eed29c3360178c2c7fd07f5132a2a8b26c8e Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 10:20:11 +0100 Subject: [PATCH 140/195] remane to maximum_pad_value --- .../unlimited_ocr/image_processing_unlimited_ocr.py | 10 +++++----- .../models/unlimited_ocr/modular_unlimited_ocr.py | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index 7d78f7048f76..bea2c8892fce 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -44,7 +44,7 @@ class UnlimitedOcrImageProcessorKwargs(ImagesKwargs, total=False): set to `True`. Can be overridden by the `max_patches` parameter in the `preprocess` method. tile_size (`int`, *optional*, defaults to `640`): The size of each local tile. Must match the model's query embedding size. - pad_if_larger_than (`int`, *optional*, defaults to `640`): + maximum_pad_value (`int`, *optional*, defaults to `640`): If `crop_to_patches` is `False` and `size.height/width` is larger than this value, the image will be resized directly to `size.height/width` without padding. Otherwise, images are resized and padded to `size.height/width` while preserving the aspect ratio. @@ -59,7 +59,7 @@ class UnlimitedOcrImageProcessorKwargs(ImagesKwargs, total=False): tile_size: int background_color: list[int] - pad_if_larger_than: int + maximum_pad_value: int @lru_cache(maxsize=10) @@ -153,7 +153,7 @@ class UnlimitedOcrImageProcessor(TorchvisionBackend): tile_size = 640 background_color = [127, 127, 127] model_input_names = ["pixel_values", "num_local_patches", "patches_grid"] - pad_if_larger_than = 640 + maximum_pad_value = 640 def __init__(self, **kwargs: Unpack[UnlimitedOcrImageProcessorKwargs]): super().__init__(**kwargs) @@ -220,7 +220,7 @@ def _preprocess( min_patches: int, max_patches: int, tile_size: int, - pad_if_larger_than: int, + maximum_pad_value: int, resample: "PILImageResampling | None", do_rescale: bool, rescale_factor: float, @@ -275,7 +275,7 @@ def _preprocess( processed_global_grouped = {} for shape, stacked in grouped_images.items(): # Different from DeepseekOcr2 which crops and pads all images - if not crop_to_patches and global_target_size <= pad_if_larger_than: + if not crop_to_patches and global_target_size <= maximum_pad_value: stacked = self.resize( stacked, SizeDict(height=global_target_size, width=global_target_size), resample=resample ) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index dc35e88b18b5..458dce929df4 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -85,7 +85,7 @@ class UnlimitedOcrImageProcessorKwargs(DeepseekOcr2ImageProcessorKwargs): set to `True`. Can be overridden by the `max_patches` parameter in the `preprocess` method. tile_size (`int`, *optional*, defaults to `640`): The size of each local tile. Must match the model's query embedding size. - pad_if_larger_than (`int`, *optional*, defaults to `640`): + maximum_pad_value (`int`, *optional*, defaults to `640`): If `crop_to_patches` is `False` and `size.height/width` is larger than this value, the image will be resized directly to `size.height/width` without padding. Otherwise, images are resized and padded to `size.height/width` while preserving the aspect ratio. @@ -93,12 +93,12 @@ class UnlimitedOcrImageProcessorKwargs(DeepseekOcr2ImageProcessorKwargs): The background color for padding. """ - pad_if_larger_than: int + maximum_pad_value: int class UnlimitedOcrImageProcessor(DeepseekOcr2ImageProcessor): tile_size = 640 - pad_if_larger_than = 640 + maximum_pad_value = 640 max_patches = 32 model_input_names = ["pixel_values", "num_local_patches", "patches_grid"] @@ -110,7 +110,7 @@ def _preprocess( min_patches: int, max_patches: int, tile_size: int, - pad_if_larger_than: int, + maximum_pad_value: int, resample: "PILImageResampling | None", do_rescale: bool, rescale_factor: float, @@ -165,7 +165,7 @@ def _preprocess( processed_global_grouped = {} for shape, stacked in grouped_images.items(): # Different from DeepseekOcr2 which crops and pads all images - if not crop_to_patches and global_target_size <= pad_if_larger_than: + if not crop_to_patches and global_target_size <= maximum_pad_value: stacked = self.resize( stacked, SizeDict(height=global_target_size, width=global_target_size), resample=resample ) From c47d816227410973bb48793d826a8d17aa0826d1 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 10:59:09 +0100 Subject: [PATCH 141/195] cleanup --- .../unlimited_ocr/modeling_unlimited_ocr.py | 22 +++++++-------- .../unlimited_ocr/modular_unlimited_ocr.py | 28 +++++++------------ 2 files changed, 20 insertions(+), 30 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 1c32a1ecc28b..1e1a279a4719 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -834,6 +834,7 @@ def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: class_embeds = self.class_embedding.expand(batch_size, 1, -1) embeddings = torch.cat([class_embeds, patch_embeds], dim=1) + # always interpolate embeddings = embeddings + self.interpolate_pos_encoding( embeddings, grid_height * self.patch_size, grid_width * self.patch_size ) @@ -947,14 +948,14 @@ def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwarg >>> patch_tokens = outputs.last_hidden_state ```""" - sam_encoder_outputs = self.sam_encoder(pixel_values, **kwargs) - sam_hidden_states = sam_encoder_outputs.last_hidden_state + sam_outputs = self.sam_encoder(pixel_values, **kwargs) + residual = sam_outputs.last_hidden_state - vision_encoder_outputs = self.vision_encoder(sam_hidden_states, **kwargs) + vision_encoder_outputs = self.vision_encoder(residual, **kwargs) hidden_states = vision_encoder_outputs.last_hidden_state - sam_hidden_states = sam_hidden_states.flatten(2).transpose(1, 2) - hidden_states = torch.cat([hidden_states[:, 1:], sam_hidden_states], dim=-1) + residual = residual.flatten(2).transpose(1, 2) + hidden_states = torch.cat([hidden_states[:, 1:], residual], dim=-1) return BaseModelOutput( last_hidden_state=hidden_states, @@ -1869,8 +1870,8 @@ def get_image_features( view_separator = self.view_separator[None, :] all_features = [] + num_queries_global = int(global_features.shape[1] ** 0.5) for idx in range(batch_size): - num_queries_global = int(global_features.shape[1] ** 0.5) global_grid = global_features[idx].reshape(num_queries_global, num_queries_global, hidden_size) global_grid = torch.cat([global_grid, newline.expand(num_queries_global, 1, hidden_size)], dim=1) global_flat = global_grid.reshape(-1, hidden_size) @@ -1879,15 +1880,12 @@ def get_image_features( if local_features is not None and local_features.shape[0] > 0: num_columns, num_rows = int(patches_grid[idx][0]), int(patches_grid[idx][1]) num_queries_local = int(local_features.shape[1] ** 0.5) + local_grid_shape = (num_rows * num_queries_local, -1, hidden_size) local_grid = local_features.reshape( num_rows, num_columns, num_queries_local, num_queries_local, hidden_size ) - local_grid = local_grid.permute(0, 2, 1, 3, 4).reshape( - num_rows * num_queries_local, num_columns * num_queries_local, hidden_size - ) - local_grid = torch.cat( - [local_grid, newline.expand(num_rows * num_queries_local, 1, hidden_size)], dim=1 - ) + local_grid = local_grid.transpose(1, 2).reshape(local_grid_shape) + local_grid = torch.cat([local_grid, newline.expand(local_grid_shape)], dim=1) local_flat = local_grid.reshape(-1, hidden_size) all_features.append(torch.cat([local_flat, global_flat, view_separator], dim=0)) else: diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 458dce929df4..9a7294f8fd01 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -42,7 +42,6 @@ is_torchdynamo_compiling, torch_int, ) -from ...utils.generic import merge_with_config_defaults from ...utils.output_capturing import capture_outputs from ..clip.configuration_clip import CLIPVisionConfig from ..clip.modeling_clip import CLIPAttention, CLIPEncoderLayer, CLIPVisionEmbeddings, CLIPVisionModel @@ -580,6 +579,7 @@ def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: class_embeds = self.class_embedding.expand(batch_size, 1, -1) embeddings = torch.cat([class_embeds, patch_embeds], dim=1) + # always interpolate embeddings = embeddings + self.interpolate_pos_encoding( embeddings, grid_height * self.patch_size, grid_width * self.patch_size ) @@ -624,8 +624,6 @@ def __init__(self, config: UnlimitedOcrVisionConfig): del self.query_768_resolution del self.query_1024_resolution - @can_return_tuple - @auto_docstring def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: r""" Example: @@ -650,14 +648,14 @@ def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwarg >>> patch_tokens = outputs.last_hidden_state ```""" - sam_encoder_outputs = self.sam_encoder(pixel_values, **kwargs) - sam_hidden_states = sam_encoder_outputs.last_hidden_state + sam_outputs = self.sam_encoder(pixel_values, **kwargs) + residual = sam_outputs.last_hidden_state - vision_encoder_outputs = self.vision_encoder(sam_hidden_states, **kwargs) + vision_encoder_outputs = self.vision_encoder(residual, **kwargs) hidden_states = vision_encoder_outputs.last_hidden_state - sam_hidden_states = sam_hidden_states.flatten(2).transpose(1, 2) - hidden_states = torch.cat([hidden_states[:, 1:], sam_hidden_states], dim=-1) + residual = residual.flatten(2).transpose(1, 2) + hidden_states = torch.cat([hidden_states[:, 1:], residual], dim=-1) return BaseModelOutput( last_hidden_state=hidden_states, @@ -1039,9 +1037,6 @@ def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): class UnlimitedOcrTextModel(DeepseekOcr2TextModel): - @merge_with_config_defaults - @capture_outputs - @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, @@ -1157,8 +1152,8 @@ def get_image_features( view_separator = self.view_separator[None, :] all_features = [] + num_queries_global = int(global_features.shape[1] ** 0.5) for idx in range(batch_size): - num_queries_global = int(global_features.shape[1] ** 0.5) global_grid = global_features[idx].reshape(num_queries_global, num_queries_global, hidden_size) global_grid = torch.cat([global_grid, newline.expand(num_queries_global, 1, hidden_size)], dim=1) global_flat = global_grid.reshape(-1, hidden_size) @@ -1167,15 +1162,12 @@ def get_image_features( if local_features is not None and local_features.shape[0] > 0: num_columns, num_rows = int(patches_grid[idx][0]), int(patches_grid[idx][1]) num_queries_local = int(local_features.shape[1] ** 0.5) + local_grid_shape = (num_rows * num_queries_local, -1, hidden_size) local_grid = local_features.reshape( num_rows, num_columns, num_queries_local, num_queries_local, hidden_size ) - local_grid = local_grid.permute(0, 2, 1, 3, 4).reshape( - num_rows * num_queries_local, num_columns * num_queries_local, hidden_size - ) - local_grid = torch.cat( - [local_grid, newline.expand(num_rows * num_queries_local, 1, hidden_size)], dim=1 - ) + local_grid = local_grid.transpose(1, 2).reshape(local_grid_shape) + local_grid = torch.cat([local_grid, newline.expand(local_grid_shape)], dim=1) local_flat = local_grid.reshape(-1, hidden_size) all_features.append(torch.cat([local_flat, global_flat, view_separator], dim=0)) else: From 20938584faaac174af53f7197f1fe101422cf510 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:22:11 +0100 Subject: [PATCH 142/195] add chat template --- docs/source/en/model_doc/unlimited_ocr.md | 67 ++++++++++++++----- .../unlimited_ocr/modeling_unlimited_ocr.py | 28 ++++++-- .../unlimited_ocr/modular_unlimited_ocr.py | 28 ++++++-- .../test_modeling_unlimited_ocr.py | 40 ++++++++--- .../test_processing_unlimited_ocr.py | 52 +++++++++++--- 5 files changed, 166 insertions(+), 49 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 7230cdff1c02..af8fa0001067 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -33,7 +33,7 @@ The abstract from the paper is the following: -Unlimited-OCR supports two inference configurations: the default "gundam" mode uses 640x640 tiles with dynamic cropping for high-resolution documents, and "base" mode uses a single 1024x1024 global view for standard-resolution inputs. To enable base mode set `crop_to_patches=False` in the processor. +Unlimited-OCR supports two inference configurations: the default "gundam" mode uses 640x640 tiles with dynamic cropping for high-resolution documents, and "base" mode uses a single 1024x1024 global view for standard-resolution inputs. To enable base mode set `crop_to_patches=False` in the processor, or pass it via `processor_kwargs` when using [`~ProcessorMixin.apply_chat_template`]. The vision tower follows the two-stage approach from [DeepSeek-OCR-2](./deepseek_ocr2): a SAM ViT-B encoder feeds into a CLIP ViT encoder. Unlike DeepSeek-OCR-2, the CLIP features are additionally concatenated with the SAM features to yield the final image tokens. Unlimited-OCR also omits the learnable patch queries from DeepSeek-OCR-2. @@ -42,8 +42,6 @@ The text model is identical to DeepSeek-OCR-2 with the additional Reference Slid This model was contributed by [guarin](https://huggingface.co/guarin). The original code can be found [here](https://github.com/baidu/Unlimited-OCR). -> [!TIP] -> For multi-page documents, pass all page images together with one `` token per page in the text prompt. The model processes all pages jointly within a single context window. @@ -55,7 +53,15 @@ model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", devic processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") image = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" -inputs = processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) +messages = [ + { + "role": "user", + "content": [{"type": "image", "url": image}, {"type": "text", "text": "document parsing."}], + } +] +inputs = processor.apply_chat_template( + messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" +).to(model.device) output = model.generate( **inputs, @@ -71,8 +77,7 @@ processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=T ### Batch processing -For batch processing, pass multiple images and prompts at once. Set `padding=True` for the processor -if images have different sizes. +For batch processing, pass a list of messages. Set `padding=True` for the processor if images have different sizes. ```python from transformers import AutoProcessor, AutoModelForImageTextToText @@ -82,11 +87,22 @@ processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") image1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" image2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_receipt.jpeg" -inputs = processor( - images=[image1, image2], - text=["document parsing.", "document parsing."], - padding=True, +messages = [ + [ + { + "role": "user", + "content": [{"type": "image", "url": image}, {"type": "text", "text": "document parsing."}], + } + ] + for image in [image1, image2] +] +inputs = processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, return_tensors="pt", + processor_kwargs={"padding": True}, ).to(model.device) output = model.generate( @@ -113,7 +129,15 @@ model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", devic processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") image = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" -inputs = processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) +messages = [ + { + "role": "user", + "content": [{"type": "image", "url": image}, {"type": "text", "text": "document parsing."}], + } +] +inputs = processor.apply_chat_template( + messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" +).to(model.device) output = model.generate( **inputs, @@ -152,10 +176,9 @@ plt.show() -Multi-page documents can be parsed jointly in a single forward pass by passing all page images together. Include one `` token per page in the text prompt so the model processes all pages as a continuous document. +Multi-page documents can be parsed jointly in a single forward pass by passing all page images together. Add one image block per page to the message so the model processes all pages as a continuous document. ```python -from PIL import Image from transformers import AutoProcessor, AutoModelForImageTextToText model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", device_map="auto") @@ -163,13 +186,21 @@ processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") page1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" page2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_receipt.jpeg" -num_pages = 2 -inputs = processor( - images=[page1, page2], - text="" * num_pages + "Multi page parsing.", - crop_to_patches=False, +messages = [ + { + "role": "user", + "content": [{"type": "image", "url": page} for page in [page1, page2]] + + [{"type": "text", "text": "Multi page parsing."}], + } +] +inputs = processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, return_tensors="pt", + processor_kwargs={"crop_to_patches": False}, ).to(model.device) output = model.generate( diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 1e1a279a4719..4ce6c2a456a9 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -2051,7 +2051,15 @@ def forward( >>> processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") >>> image = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" - >>> inputs = processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) + >>> messages = [ + ... { + ... "role": "user", + ... "content": [{"type": "image", "url": image}, {"type": "text", "text": "document parsing."}], + ... } + ... ] + >>> inputs = processor.apply_chat_template( + ... messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" + ... ).to(model.device) >>> output = model.generate( ... **inputs, @@ -2073,13 +2081,21 @@ def forward( >>> page1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" >>> page2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_receipt.jpeg" - >>> num_pages = 2 - >>> inputs = processor( - ... images=[page1, page2], - ... text="" * num_pages + "Multi page parsing.", - ... crop_to_patches=False, + >>> messages = [ + ... { + ... "role": "user", + ... "content": [{"type": "image", "url": page} for page in [page1, page2]] + ... + [{"type": "text", "text": "Multi page parsing."}], + ... } + ... ] + >>> inputs = processor.apply_chat_template( + ... messages, + ... add_generation_prompt=True, + ... tokenize=True, + ... return_dict=True, ... return_tensors="pt", + ... processor_kwargs={"crop_to_patches": False}, ... ).to(model.device) >>> output = model.generate( diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 9a7294f8fd01..89b908dbf187 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -1310,7 +1310,15 @@ def forward( >>> processor = AutoProcessor.from_pretrained("baidu/Unlimited-OCR") >>> image = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" - >>> inputs = processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) + >>> messages = [ + ... { + ... "role": "user", + ... "content": [{"type": "image", "url": image}, {"type": "text", "text": "document parsing."}], + ... } + ... ] + >>> inputs = processor.apply_chat_template( + ... messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" + ... ).to(model.device) >>> output = model.generate( ... **inputs, @@ -1332,13 +1340,21 @@ def forward( >>> page1 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg" >>> page2 = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_receipt.jpeg" - >>> num_pages = 2 - >>> inputs = processor( - ... images=[page1, page2], - ... text="" * num_pages + "Multi page parsing.", - ... crop_to_patches=False, + >>> messages = [ + ... { + ... "role": "user", + ... "content": [{"type": "image", "url": page} for page in [page1, page2]] + ... + [{"type": "text", "text": "Multi page parsing."}], + ... } + ... ] + >>> inputs = processor.apply_chat_template( + ... messages, + ... add_generation_prompt=True, + ... tokenize=True, + ... return_dict=True, ... return_tensors="pt", + ... processor_kwargs={"crop_to_patches": False}, ... ).to(model.device) >>> output = model.generate( diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index e5e2d91e48bf..681b66b4da4c 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -221,6 +221,16 @@ def setUp(self): def tearDown(self): cleanup(torch_device, gc_collect=True) + def get_messages(self, images, text): + if not isinstance(images, list): + images = [images] + return [ + { + "role": "user", + "content": [{"type": "image", "image": image} for image in images] + [{"type": "text", "text": text}], + } + ] + @slow @require_torch_accelerator def test_small_model_integration_test_document_parsing(self): @@ -232,7 +242,10 @@ def test_small_model_integration_test_document_parsing(self): "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" ) ) - inputs = self.processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) + messages = self.get_messages(image, "document parsing.") + inputs = self.processor.apply_chat_template( + messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" + ).to(model.device) with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): generate_ids = model.generate( **inputs, @@ -259,7 +272,10 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/image_ocr.jpg" ) ) - inputs = self.processor(images=image, text="document parsing.", return_tensors="pt").to(model.device) + messages = self.get_messages(image, "document parsing.") + inputs = self.processor.apply_chat_template( + messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" + ).to(model.device) with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): generate_ids = model.generate( **inputs, @@ -301,9 +317,12 @@ def test_small_model_integration_test_document_parsing_batched(self): "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png" ) ) - inputs = self.processor( - images=[image1, image2], - text=["document parsing.", "document parsing."], + messages = [self.get_messages(image, "document parsing.") for image in [image1, image2]] + inputs = self.processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, return_tensors="pt", ).to(model.device, dtype=torch.bfloat16) @@ -347,11 +366,14 @@ def test_small_model_integration_test_multi_page_document_parsing(self): "https://huggingface.co/datasets/hf-internal-testing/fixtures_got_ocr/resolve/main/multi_box.png" ) ) - inputs = self.processor( - images=[image1, image2], - text="Multi page parsing.", - crop_to_patches=False, + messages = self.get_messages([image1, image2], "Multi page parsing.") + inputs = self.processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, return_tensors="pt", + processor_kwargs={"crop_to_patches": False}, ).to(model.device) with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): diff --git a/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py b/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py index a1c426c6ead5..d95d43b0fcaa 100644 --- a/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py @@ -1,3 +1,16 @@ +# 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 unittest import torch @@ -24,13 +37,23 @@ class UnlimitedOcrProcessorTest(ProcessorTesterMixin, unittest.TestCase): image_text_kwargs_override_max_length = 310 image_unstructured_max_length = 300 + def get_messages(self, image): + return [ + { + "role": "user", + "content": [{"type": "image", "image": image}, {"type": "text", "text": "document parsing."}], + } + ] + def test_replace_image_tokens(self): processor = self.get_processor() - images = torch.randint(0, 256, (1, 3, 200, 300), dtype=torch.uint8) - prompt = "document parsing." + image = torch.randint(0, 256, (3, 200, 300), dtype=torch.uint8) + messages = self.get_messages(image) - inputs = processor(images=images, text=prompt, return_tensors="pt") + inputs = processor.apply_chat_template( + messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" + ) num_image_tokens = (inputs["input_ids"] == processor.image_token_id).sum().item() # image resized to 1024, followed by patch size 16 and 4x downsampling = 16 x 16 patches @@ -41,10 +64,12 @@ def test_replace_image_tokens(self): def test_replace_image_tokens_with_local(self): processor = self.get_processor() - images = torch.randint(0, 256, (1, 3, 500, 700), dtype=torch.uint8) - prompt = "document parsing." + image = torch.randint(0, 256, (3, 500, 700), dtype=torch.uint8) + messages = self.get_messages(image) - inputs = processor(images=images, text=prompt, return_tensors="pt") + inputs = processor.apply_chat_template( + messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" + ) num_image_tokens = (inputs["input_ids"] == processor.image_token_id).sum().item() # global is same as in test above @@ -57,10 +82,17 @@ def test_replace_image_tokens_with_local(self): def test_replace_image_tokens_no_crop(self): processor = self.get_processor() - images = torch.randint(0, 256, (1, 3, 500, 700), dtype=torch.uint8) - prompt = "document parsing." - - inputs = processor(images=images, text=prompt, return_tensors="pt", crop_to_patches=False) + image = torch.randint(0, 256, (3, 500, 700), dtype=torch.uint8) + messages = self.get_messages(image) + + inputs = processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + processor_kwargs={"crop_to_patches": False}, + ) num_image_tokens = (inputs["input_ids"] == processor.image_token_id).sum().item() # same as in test_replace_image_tokens From 9581dfcfb52d13c2962a3f801f19478e8ee909a5 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:53:13 +0100 Subject: [PATCH 143/195] add comment --- src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py | 2 +- src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 4ce6c2a456a9..687bb47104b9 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1804,7 +1804,7 @@ def forward( ) class UnlimitedOcrModel(UnlimitedOcrPreTrainedModel): base_model_prefix = "model" - _keys_to_ignore_on_load_unexpected = {"lm_head"} + _keys_to_ignore_on_load_unexpected = {"lm_head"} # unused and untied weight in original checkpoint def __init__(self, config: UnlimitedOcrConfig): super().__init__(config) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 89b908dbf187..609cb7381f8b 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -1098,7 +1098,7 @@ def forward( class UnlimitedOcrModel(DeepseekOcr2Model): - _keys_to_ignore_on_load_unexpected = {"lm_head"} + _keys_to_ignore_on_load_unexpected = {"lm_head"} # unused and untied weight in original checkpoint def __init__(self, config: UnlimitedOcrConfig): super().__init__(config) From 8fe09cb253a7f11bfe47d41d2606261b24599646 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:04:33 +0100 Subject: [PATCH 144/195] remove decorators --- src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 609cb7381f8b..e8e50ee90a65 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -494,7 +494,6 @@ class UnlimitedOcrPreTrainedModel(DeepseekOcr2PreTrainedModel): "UnlimitedOcrTextDecoderLayer", ] - @torch.no_grad() def _init_weights(self, module): super()._init_weights(module) if isinstance(module, UnlimitedOcrModel): @@ -1182,8 +1181,6 @@ def get_image_features( **local_outputs, ) - @can_return_tuple - @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, From 203e38cc571ba36a3070c2c74c31af1bdae10715 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:05:20 +0100 Subject: [PATCH 145/195] convert to int --- src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 2 +- .../models/unlimited_ocr/processing_unlimited_ocr.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index e8e50ee90a65..7fbf37c970b0 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -271,7 +271,7 @@ def _parse_detections(self, decoded: str) -> list[dict]: detections.append( { "region_type": region_type, - "box": [x1, y1, x2, y2], + "box": [int(x1), int(y1), int(x2), int(y2)], "text": text, } ) diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index a6d5551574ac..b31eca147055 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -139,7 +139,7 @@ def _parse_detections(self, decoded: str) -> list[dict]: detections.append( { "region_type": region_type, - "box": [x1, y1, x2, y2], + "box": [int(x1), int(y1), int(x2), int(y2)], "text": text, } ) From 5a30585f6a339bc5f09a51f609ffbe1cac6f221f Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:30:55 +0100 Subject: [PATCH 146/195] handle cache record_past --- .../unlimited_ocr/modeling_unlimited_ocr.py | 70 +++++++++++-------- .../unlimited_ocr/modular_unlimited_ocr.py | 70 +++++++++++-------- 2 files changed, 84 insertions(+), 56 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 687bb47104b9..8919fb3c97b7 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1313,7 +1313,7 @@ def update( full_value_states = torch.cat([self.values, value_states], dim=-2) # Prefill - if self.prefill_length is None: + if self.prefill_length is None or self.record_past: self.keys = full_key_states self.values = full_value_states return self.keys, self.values @@ -1383,13 +1383,27 @@ def crop(self, max_length: int) -> None: Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be negative to remove `max_length` tokens. """ + # If we are beyond the sliding window, we need to be more careful if self.prefill_length is not None and self.get_seq_length() >= self.prefill_length + self.sliding_window: - raise ValueError( - "Cannot `crop` a `UnlimitedOcrDynamicReferenceSlidingWindowLayer` after it has seen more tokens than its" - "prefill + sliding window (otherwise some states are lost)" - ) - DynamicLayer.crop(self, max_length) - self.cumulative_length = self.keys.shape[-2] + if not self.record_past: + raise RuntimeError( + "`crop` was called, but the current layer does not track past states, and the sliding window size was already " + "reached. Call `activate_past_recording` before `crop` to be able to rollback the cache." + ) + if max_length > 0: + raise RuntimeError( + "Once the sliding window size has been reached, `UnlimitedOcrDynamicReferenceSlidingWindowLayer` can only " + "be cropped by passing a negative int, to specify how many tokens to remove" + ) + tokens_to_remove = abs(max_length) + # We crop, and restrict the size back to the sliding window if still larger + self.keys = self.keys[:, :, -self.sliding_window + 1 - tokens_to_remove : -tokens_to_remove, :] + self.values = self.values[:, :, -self.sliding_window + 1 - tokens_to_remove : -tokens_to_remove, :] + self.cumulative_length = self.cumulative_length - tokens_to_remove + # If we did not reach the sliding window, we can do the same as for a full attention layer + else: + DynamicLayer.crop(self, max_length) + self.cumulative_length = self.keys.shape[-2] class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): @@ -1477,7 +1491,7 @@ def update( # In general, we should use a much simpler `cat` here as well, independently of the states size. However, # dynamo is currently bugged when doing it - see https://github.com/pytorch/pytorch/issues/159855 for more details if key_states.shape[-2] == 1: - # Roll the window region to the left by 1 position (the pinned prefill in front stays put) + # Roll the window region to the left by 1 position new_keys = self.keys[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) new_values = self.values[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) # Overwrite the last position with new states @@ -1544,26 +1558,6 @@ def update( ) return full_key_states, full_value_states - def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None: - self.dtype, self.device = key_states.dtype, key_states.device - self.batch_size, self.num_heads = key_states.shape[:2] - self.v_head_dim = value_states.shape[-1] - self.k_head_dim = key_states.shape[-1] - - self.cumulative_length = self.cumulative_length.to(self.device) - # Note: `mark_static_address` is used to tag the tensors as a fixed data pointer, preventing compiled graph - # breaks or cudagraph skips due to inplace mutations when updating the cache. However, it is not supported when - # tracing the graph, so we skip it in this case. As prefill should never be compiled, this is not an issue and it - # will still be run (except when users compile prefill explicitly, but this should be avoided!) - # Without this, we cannot use cudagraphs!! - if not is_torchdynamo_compiling(): - torch._dynamo.mark_static_address(self.cumulative_length) - - prefill_length = key_states.shape[-2] if self.prefill_length is None else self.prefill_length - self._allocate_key_value_buffers(min(self.max_cache_len, prefill_length + self.sliding_window)) - - self.is_initialized = True - def get_mask_sizes(self, query_length: int) -> tuple[int, int]: """Return the length and offset of the cache, used to generate the attention mask""" is_full = ( @@ -1605,6 +1599,26 @@ def reset(self) -> None: super().reset() self.prefill_length = None + def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None: + self.dtype, self.device = key_states.dtype, key_states.device + self.batch_size, self.num_heads = key_states.shape[:2] + self.v_head_dim = value_states.shape[-1] + self.k_head_dim = key_states.shape[-1] + + self.cumulative_length = self.cumulative_length.to(self.device) + # Note: `mark_static_address` is used to tag the tensors as a fixed data pointer, preventing compiled graph + # breaks or cudagraph skips due to inplace mutations when updating the cache. However, it is not supported when + # tracing the graph, so we skip it in this case. As prefill should never be compiled, this is not an issue and it + # will still be run (except when users compile prefill explicitly, but this should be avoided!) + # Without this, we cannot use cudagraphs!! + if not is_torchdynamo_compiling(): + torch._dynamo.mark_static_address(self.cumulative_length) + + prefill_length = key_states.shape[-2] if self.prefill_length is None else self.prefill_length + self._allocate_key_value_buffers(min(self.max_cache_len, prefill_length + self.sliding_window)) + + self.is_initialized = True + def _allocate_key_value_buffers(self, physical_length: int, copy_existing: bool = False) -> None: """(Re)allocate the static key/value buffers to `physical_length` slots and (re)tag the static address. diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 7fbf37c970b0..60c3c5d58d02 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -698,7 +698,7 @@ def update( full_value_states = torch.cat([self.values, value_states], dim=-2) # Prefill - if self.prefill_length is None: + if self.prefill_length is None or self.record_past: self.keys = full_key_states self.values = full_value_states return self.keys, self.values @@ -768,13 +768,27 @@ def crop(self, max_length: int) -> None: Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be negative to remove `max_length` tokens. """ + # If we are beyond the sliding window, we need to be more careful if self.prefill_length is not None and self.get_seq_length() >= self.prefill_length + self.sliding_window: - raise ValueError( - "Cannot `crop` a `UnlimitedOcrDynamicReferenceSlidingWindowLayer` after it has seen more tokens than its" - "prefill + sliding window (otherwise some states are lost)" - ) - DynamicLayer.crop(self, max_length) - self.cumulative_length = self.keys.shape[-2] + if not self.record_past: + raise RuntimeError( + "`crop` was called, but the current layer does not track past states, and the sliding window size was already " + "reached. Call `activate_past_recording` before `crop` to be able to rollback the cache." + ) + if max_length > 0: + raise RuntimeError( + "Once the sliding window size has been reached, `UnlimitedOcrDynamicReferenceSlidingWindowLayer` can only " + "be cropped by passing a negative int, to specify how many tokens to remove" + ) + tokens_to_remove = abs(max_length) + # We crop, and restrict the size back to the sliding window if still larger + self.keys = self.keys[:, :, -self.sliding_window + 1 - tokens_to_remove : -tokens_to_remove, :] + self.values = self.values[:, :, -self.sliding_window + 1 - tokens_to_remove : -tokens_to_remove, :] + self.cumulative_length = self.cumulative_length - tokens_to_remove + # If we did not reach the sliding window, we can do the same as for a full attention layer + else: + DynamicLayer.crop(self, max_length) + self.cumulative_length = self.keys.shape[-2] class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): @@ -862,7 +876,7 @@ def update( # In general, we should use a much simpler `cat` here as well, independently of the states size. However, # dynamo is currently bugged when doing it - see https://github.com/pytorch/pytorch/issues/159855 for more details if key_states.shape[-2] == 1: - # Roll the window region to the left by 1 position (the pinned prefill in front stays put) + # Roll the window region to the left by 1 position new_keys = self.keys[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) new_values = self.values[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) # Overwrite the last position with new states @@ -929,26 +943,6 @@ def update( ) return full_key_states, full_value_states - def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None: - self.dtype, self.device = key_states.dtype, key_states.device - self.batch_size, self.num_heads = key_states.shape[:2] - self.v_head_dim = value_states.shape[-1] - self.k_head_dim = key_states.shape[-1] - - self.cumulative_length = self.cumulative_length.to(self.device) - # Note: `mark_static_address` is used to tag the tensors as a fixed data pointer, preventing compiled graph - # breaks or cudagraph skips due to inplace mutations when updating the cache. However, it is not supported when - # tracing the graph, so we skip it in this case. As prefill should never be compiled, this is not an issue and it - # will still be run (except when users compile prefill explicitly, but this should be avoided!) - # Without this, we cannot use cudagraphs!! - if not is_torchdynamo_compiling(): - torch._dynamo.mark_static_address(self.cumulative_length) - - prefill_length = key_states.shape[-2] if self.prefill_length is None else self.prefill_length - self._allocate_key_value_buffers(min(self.max_cache_len, prefill_length + self.sliding_window)) - - self.is_initialized = True - def get_mask_sizes(self, query_length: int) -> tuple[int, int]: """Return the length and offset of the cache, used to generate the attention mask""" is_full = ( @@ -990,6 +984,26 @@ def reset(self) -> None: super().reset() self.prefill_length = None + def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None: + self.dtype, self.device = key_states.dtype, key_states.device + self.batch_size, self.num_heads = key_states.shape[:2] + self.v_head_dim = value_states.shape[-1] + self.k_head_dim = key_states.shape[-1] + + self.cumulative_length = self.cumulative_length.to(self.device) + # Note: `mark_static_address` is used to tag the tensors as a fixed data pointer, preventing compiled graph + # breaks or cudagraph skips due to inplace mutations when updating the cache. However, it is not supported when + # tracing the graph, so we skip it in this case. As prefill should never be compiled, this is not an issue and it + # will still be run (except when users compile prefill explicitly, but this should be avoided!) + # Without this, we cannot use cudagraphs!! + if not is_torchdynamo_compiling(): + torch._dynamo.mark_static_address(self.cumulative_length) + + prefill_length = key_states.shape[-2] if self.prefill_length is None else self.prefill_length + self._allocate_key_value_buffers(min(self.max_cache_len, prefill_length + self.sliding_window)) + + self.is_initialized = True + def _allocate_key_value_buffers(self, physical_length: int, copy_existing: bool = False) -> None: """(Re)allocate the static key/value buffers to `physical_length` slots and (re)tag the static address. From 6e31d534086d1dc13abac95b5946fc9aea7e033e Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:40:58 +0100 Subject: [PATCH 147/195] cleanup validate_inputs --- .../deepseek_ocr2/processing_deepseek_ocr2.py | 31 +++++++++-------- .../unlimited_ocr/modular_unlimited_ocr.py | 9 +++-- .../unlimited_ocr/processing_unlimited_ocr.py | 34 +++++++++++-------- 3 files changed, 42 insertions(+), 32 deletions(-) diff --git a/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py index f100b1b1be5f..8df57494cd4c 100644 --- a/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py @@ -41,6 +41,8 @@ class DeepseekOcr2ProcessorKwargs(ProcessingKwargs, total=False): @auto_docstring class DeepseekOcr2Processor(ProcessorMixin): + valid_processor_kwargs = DeepseekOcr2ProcessorKwargs + def __init__( self, image_processor=None, @@ -63,23 +65,24 @@ def __init__( super().__init__(image_processor, tokenizer, chat_template=chat_template, **kwargs) def validate_inputs(self, images=None, text=None, videos=None, audio=None, **kwargs): - super().validate_inputs(images=images, text=text, **kwargs) - if text is not None and images is not None: - if isinstance(text, str): - text = [text] - n_tokens = sum(sample.count(self.image_token) for sample in text) - n_images = len(make_flat_list_of_images(images)) - if n_tokens != n_images: - raise ValueError( - f"Number of {self.image_token} tokens in text ({n_tokens}) does not match " - f"number of images ({n_images})." - ) + super().validate_inputs(images=images, text=text, videos=videos, audio=audio, **kwargs) + if text is None: + raise ValueError("You have to specify text.") + + if isinstance(text, str): + text = [text] + total_placeholders = sum(prompt.count(self.image_token) for prompt in text) + num_images = len(make_flat_list_of_images(images)) if images is not None else 0 + if total_placeholders != num_images: + raise ValueError( + f"Found {total_placeholders} placeholders across the batch, but have {num_images} flattened images." + ) def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> TextInput: - size = self.image_processor.size["height"] - tile_size = self.image_processor.tile_size + size = kwargs.get("size") or self.image_processor.size + tile_size = kwargs.get("tile_size") or self.image_processor.tile_size - num_queries_global = math.ceil(size / self.patch_size / self.downsample_ratio) + num_queries_global = math.ceil(size["height"] / self.patch_size / self.downsample_ratio) global_tokens = num_queries_global * num_queries_global num_queries_local = math.ceil(tile_size / self.patch_size / self.downsample_ratio) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 60c3c5d58d02..42999e9e7556 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -202,9 +202,12 @@ class UnlimitedOcrProcessorKwargs(DeepseekOcr2ProcessorKwargs): class UnlimitedOcrProcessor(DeepseekOcr2Processor): - def replace_image_token(self, image_inputs: dict, image_idx: int) -> TextInput: - size = max(self.image_processor.size["height"], self.image_processor.size["width"]) - tile_size = self.image_processor.tile_size + valid_processor_kwargs = UnlimitedOcrProcessorKwargs + + def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> TextInput: + image_size = kwargs.get("size") or self.image_processor.size + tile_size = kwargs.get("tile_size") or self.image_processor.tile_size + size = max(image_size["height"], image_size["width"]) num_queries_global = math.ceil(size // self.patch_size / self.downsample_ratio) num_queries_local = math.ceil(tile_size // self.patch_size / self.downsample_ratio) diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index b31eca147055..82ec31a48141 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -36,6 +36,8 @@ class UnlimitedOcrProcessorKwargs(ProcessingKwargs, total=False): @auto_docstring class UnlimitedOcrProcessor(ProcessorMixin): + valid_processor_kwargs = UnlimitedOcrProcessorKwargs + def __init__( self, image_processor=None, @@ -58,21 +60,23 @@ def __init__( super().__init__(image_processor, tokenizer, chat_template=chat_template, **kwargs) def validate_inputs(self, images=None, text=None, videos=None, audio=None, **kwargs): - super().validate_inputs(images=images, text=text, **kwargs) - if text is not None and images is not None: - if isinstance(text, str): - text = [text] - n_tokens = sum(sample.count(self.image_token) for sample in text) - n_images = len(make_flat_list_of_images(images)) - if n_tokens != n_images: - raise ValueError( - f"Number of {self.image_token} tokens in text ({n_tokens}) does not match " - f"number of images ({n_images})." - ) - - def replace_image_token(self, image_inputs: dict, image_idx: int) -> TextInput: - size = max(self.image_processor.size["height"], self.image_processor.size["width"]) - tile_size = self.image_processor.tile_size + super().validate_inputs(images=images, text=text, videos=videos, audio=audio, **kwargs) + if text is None: + raise ValueError("You have to specify text.") + + if isinstance(text, str): + text = [text] + total_placeholders = sum(prompt.count(self.image_token) for prompt in text) + num_images = len(make_flat_list_of_images(images)) if images is not None else 0 + if total_placeholders != num_images: + raise ValueError( + f"Found {total_placeholders} placeholders across the batch, but have {num_images} flattened images." + ) + + def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> TextInput: + image_size = kwargs.get("size") or self.image_processor.size + tile_size = kwargs.get("tile_size") or self.image_processor.tile_size + size = max(image_size["height"], image_size["width"]) num_queries_global = math.ceil(size // self.patch_size / self.downsample_ratio) num_queries_local = math.ceil(tile_size // self.patch_size / self.downsample_ratio) From 7a2d9a90446e810f07f5d926f4451a5d209437d7 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:01:32 +0100 Subject: [PATCH 148/195] update pooler_output --- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 5 ++--- .../models/unlimited_ocr/modular_unlimited_ocr.py | 5 ++--- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 8919fb3c97b7..af9b545f8514 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1905,10 +1905,9 @@ def get_image_features( else: all_features.append(torch.cat([global_flat, view_separator], dim=0)) - image_features = torch.cat(all_features, dim=0) return UnlimitedOcrModelOutputWithPooling( last_hidden_state=global_vision_outputs.last_hidden_state, - pooler_output=image_features, + pooler_output=all_features, hidden_states=global_vision_outputs.hidden_states, attentions=global_vision_outputs.attentions, **local_outputs, @@ -1969,7 +1968,7 @@ def forward( image_features = self.get_image_features( pixel_values, pixel_values_local, num_local_patches, patches_grid, return_dict=True ).pooler_output - image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + image_features = torch.cat(image_features, dim=0).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) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 42999e9e7556..985cd47b1bdd 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -1189,10 +1189,9 @@ def get_image_features( else: all_features.append(torch.cat([global_flat, view_separator], dim=0)) - image_features = torch.cat(all_features, dim=0) return UnlimitedOcrModelOutputWithPooling( last_hidden_state=global_vision_outputs.last_hidden_state, - pooler_output=image_features, + pooler_output=all_features, hidden_states=global_vision_outputs.hidden_states, attentions=global_vision_outputs.attentions, **local_outputs, @@ -1228,7 +1227,7 @@ def forward( image_features = self.get_image_features( pixel_values, pixel_values_local, num_local_patches, patches_grid, return_dict=True ).pooler_output - image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) + image_features = torch.cat(image_features, dim=0).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) From 36adb9fff474361bab4d998c306bf27128f674a3 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:18:36 +0100 Subject: [PATCH 149/195] cleanup masking --- .../unlimited_ocr/modeling_unlimited_ocr.py | 38 +++++++++++++++--- .../unlimited_ocr/modular_unlimited_ocr.py | 39 ++++++++++++++++--- 2 files changed, 65 insertions(+), 12 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index af9b545f8514..f1cb8365c3d5 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -37,7 +37,14 @@ use_kernel_func_from_hub, use_kernelized_func, ) -from ...masking_utils import and_masks, causal_mask_function, create_causal_mask, create_sliding_window_causal_mask +from ...masking_utils import ( + BlockMask, + and_masks, + causal_mask_function, + create_causal_mask, + create_sliding_window_causal_mask, + or_masks, +) from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update @@ -1707,10 +1714,17 @@ def forward(self, x, position_ids): return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) -def create_reference_sliding_window_causal_mask(**kwargs): - past_key_values = kwargs["past_key_values"] - inputs_embeds = kwargs["inputs_embeds"] - +def create_reference_sliding_window_causal_mask( + config: PreTrainedConfig, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + past_key_values: Cache | None, + position_ids: torch.Tensor | None = None, + or_mask_function: Callable | None = None, + and_mask_function: Callable | None = None, + block_sequence_ids: torch.Tensor | None = None, + layer_idx: int | None = None, +) -> torch.Tensor | BlockMask | None: if past_key_values is None: prefill_length = float("inf") kv_offset = 0 @@ -1724,8 +1738,20 @@ def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): # Remove kv_offset to retrieve the kv_index with respect to prefill return kv_idx - kv_offset < prefill_length + prefill_mask_function = and_masks(causal_mask_function, prefill_overlay) + if or_mask_function is not None: + prefill_mask_function = or_masks(prefill_mask_function, or_mask_function) + return create_sliding_window_causal_mask( - or_mask_function=and_masks(causal_mask_function, prefill_overlay), **kwargs + config=config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + or_mask_function=prefill_mask_function, + and_mask_function=and_mask_function, + block_sequence_ids=block_sequence_ids, + layer_idx=layer_idx, ) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 985cd47b1bdd..4e9fbd87b8d9 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -13,6 +13,7 @@ # limitations under the License. import math import re +from collections.abc import Callable import torch from huggingface_hub.dataclasses import strict @@ -30,7 +31,14 @@ from ...feature_extraction_utils import BatchFeature from ...image_transforms import group_images_by_shape, reorder_images from ...image_utils import PILImageResampling, SizeDict -from ...masking_utils import and_masks, causal_mask_function, create_causal_mask, create_sliding_window_causal_mask +from ...masking_utils import ( + BlockMask, + and_masks, + causal_mask_function, + create_causal_mask, + create_sliding_window_causal_mask, + or_masks, +) from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling from ...processing_utils import Unpack from ...tokenization_utils_base import TextInput @@ -1030,10 +1038,17 @@ def _allocate_key_value_buffers(self, physical_length: int, copy_existing: bool torch._dynamo.mark_static_address(self.values) -def create_reference_sliding_window_causal_mask(**kwargs): - past_key_values = kwargs["past_key_values"] - inputs_embeds = kwargs["inputs_embeds"] - +def create_reference_sliding_window_causal_mask( + config: PreTrainedConfig, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + past_key_values: Cache | None, + position_ids: torch.Tensor | None = None, + or_mask_function: Callable | None = None, + and_mask_function: Callable | None = None, + block_sequence_ids: torch.Tensor | None = None, + layer_idx: int | None = None, +) -> torch.Tensor | BlockMask | None: if past_key_values is None: prefill_length = float("inf") kv_offset = 0 @@ -1047,8 +1062,20 @@ def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): # Remove kv_offset to retrieve the kv_index with respect to prefill return kv_idx - kv_offset < prefill_length + prefill_mask_function = and_masks(causal_mask_function, prefill_overlay) + if or_mask_function is not None: + prefill_mask_function = or_masks(prefill_mask_function, or_mask_function) + return create_sliding_window_causal_mask( - or_mask_function=and_masks(causal_mask_function, prefill_overlay), **kwargs + config=config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + or_mask_function=prefill_mask_function, + and_mask_function=and_mask_function, + block_sequence_ids=block_sequence_ids, + layer_idx=layer_idx, ) From 5b6487f41137dcd211e15232e7275ce4ec2922d9 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:07:19 +0100 Subject: [PATCH 150/195] flatten images early --- .../models/deepseek_ocr2/processing_deepseek_ocr2.py | 12 +++++++++--- .../models/unlimited_ocr/processing_unlimited_ocr.py | 12 +++++++++--- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py index 8df57494cd4c..61dfbea5aac0 100644 --- a/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py @@ -64,15 +64,21 @@ def __init__( self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token) super().__init__(image_processor, tokenizer, chat_template=chat_template, **kwargs) + def prepare_inputs_layout(self, images=None, text=None, videos=None, audio=None, **kwargs): + images, text, videos, audio = super().prepare_inputs_layout( + images=images, text=text, videos=videos, audio=audio, **kwargs + ) + if images is not None: + images = make_flat_list_of_images(images) + return images, text, videos, audio + def validate_inputs(self, images=None, text=None, videos=None, audio=None, **kwargs): super().validate_inputs(images=images, text=text, videos=videos, audio=audio, **kwargs) if text is None: raise ValueError("You have to specify text.") - if isinstance(text, str): - text = [text] total_placeholders = sum(prompt.count(self.image_token) for prompt in text) - num_images = len(make_flat_list_of_images(images)) if images is not None else 0 + num_images = len(images) if images is not None else 0 if total_placeholders != num_images: raise ValueError( f"Found {total_placeholders} placeholders across the batch, but have {num_images} flattened images." diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index 82ec31a48141..a1c771c680a5 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -59,15 +59,21 @@ def __init__( self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token) super().__init__(image_processor, tokenizer, chat_template=chat_template, **kwargs) + def prepare_inputs_layout(self, images=None, text=None, videos=None, audio=None, **kwargs): + images, text, videos, audio = super().prepare_inputs_layout( + images=images, text=text, videos=videos, audio=audio, **kwargs + ) + if images is not None: + images = make_flat_list_of_images(images) + return images, text, videos, audio + def validate_inputs(self, images=None, text=None, videos=None, audio=None, **kwargs): super().validate_inputs(images=images, text=text, videos=videos, audio=audio, **kwargs) if text is None: raise ValueError("You have to specify text.") - if isinstance(text, str): - text = [text] total_placeholders = sum(prompt.count(self.image_token) for prompt in text) - num_images = len(make_flat_list_of_images(images)) if images is not None else 0 + num_images = len(images) if images is not None else 0 if total_placeholders != num_images: raise ValueError( f"Found {total_placeholders} placeholders across the batch, but have {num_images} flattened images." From c8d8626ede4e417f1657c9cec61c303ea3189b56 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:58:08 +0100 Subject: [PATCH 151/195] move doc to multimodal section, fix hfoption block --- docs/source/en/_toctree.yml | 4 ++-- docs/source/en/model_doc/unlimited_ocr.md | 3 +-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 8aa3775f5c26..85ccc68b4b16 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -899,8 +899,6 @@ title: UL2 - local: model_doc/umt5 title: UMT5 - - local: model_doc/unlimited_ocr - title: UnlimitedOcr - local: model_doc/vaultgemma title: VaultGemma - local: model_doc/xmod @@ -1497,6 +1495,8 @@ title: TVP - local: model_doc/udop title: UDOP + - local: model_doc/unlimited_ocr + title: UnlimitedOcr - local: model_doc/video_llama_3 title: VideoLlama3 - local: model_doc/video_llava diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index af8fa0001067..87f1ab98d72b 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -122,7 +122,6 @@ processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=T Set `skip_special_tokens=False` to wrap all detections and region types in `<|det|>...<|/det|>` markers. This is useful for further post-processing of the output, for example to plot the detected bounding boxes on the image. Each detection is wrapped as `<|det|>region_type [x1, y1, x2, y2]<|/det|>text...` with coordinates normalized to a `[0, 999]` range. Set `return_detections=True` to get an additional list of dictionaries with all detections parsed as `{"region_type": region_type, "box": [x1, y1, x2, y2], "text": "..."}`. ```python -import re from transformers import AutoProcessor, AutoModelForImageTextToText model = AutoModelForImageTextToText.from_pretrained("baidu/Unlimited-OCR", device_map="auto") @@ -173,7 +172,7 @@ plt.show() - + Multi-page documents can be parsed jointly in a single forward pass by passing all page images together. Add one image block per page to the message so the model processes all pages as a continuous document. From 977bc7055176d9454258cc120ab07204bab1f4c1 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:58:08 +0100 Subject: [PATCH 152/195] remove no-op image_newline renaming --- src/transformers/conversion_mapping.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 259a80f1ecd6..957aa000802c 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -1624,7 +1624,6 @@ def _build_checkpoint_conversion_mapping(): WeightRenaming(r"(^|model\.)sam_model", r"\1vision_tower.sam_encoder"), WeightRenaming(r"(^|model\.)vision_model", r"\1vision_tower.vision_encoder"), WeightRenaming(r"(^|model\.)view_seperator", r"\1view_separator"), - WeightRenaming(r"(^|model\.)image_newline", r"\1image_newline"), WeightRenaming(r"(^|model\.)projector\.layers\.weight", r"\1multi_modal_projector.weight"), WeightRenaming(r"(^|model\.)projector\.layers\.bias", r"\1multi_modal_projector.bias"), WeightRenaming(r"(^|model\.)embed_tokens\.", r"\1language_model.embed_tokens."), From ea4d55a6f95d4361aec348a9a71ff928713afbee Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:58:08 +0100 Subject: [PATCH 153/195] add license header --- .../unlimited_ocr/test_generation_unlimited_ocr.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/models/unlimited_ocr/test_generation_unlimited_ocr.py b/tests/models/unlimited_ocr/test_generation_unlimited_ocr.py index 4be03717a7b4..4f7e3a81491b 100644 --- a/tests/models/unlimited_ocr/test_generation_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_generation_unlimited_ocr.py @@ -1,3 +1,17 @@ +# 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 unittest import torch From fa90758874140db5b16265093cc33c0edaab92fc Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 17:07:14 +0100 Subject: [PATCH 154/195] fix config and image processor docstrings --- .../unlimited_ocr/configuration_unlimited_ocr.py | 4 ++-- .../unlimited_ocr/image_processing_unlimited_ocr.py | 7 ++++--- .../models/unlimited_ocr/modular_unlimited_ocr.py | 11 ++++++----- 3 files changed, 12 insertions(+), 10 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index 166809f4b540..a2d44347eacd 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -166,9 +166,9 @@ class UnlimitedOcrTextConfig(PreTrainedConfig): MLP type (`"dense"` or `"sparse"`) for each decoder layer, e.g. `["dense", "sparse", "sparse", ...]`. layer_types (`list[str]`, *optional*): Attention type for each decoder layer. Defaults to `"reference_sliding_attention"` on every layer. - use_sliding_window (`int`, defaults to `True`): + use_sliding_window (`bool`, *optional*, defaults to `True`): Whether to use reference sliding window attention. - sliding_window (`int`, *optional*, defaults to 128): + sliding_window (`int`, *optional*, defaults to `128`): Sliding window size for reference sliding window attention. If set, every token attends to the last `sliding_window` and all image and prompt tokens. diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index bea2c8892fce..ea66ee4df057 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -45,9 +45,10 @@ class UnlimitedOcrImageProcessorKwargs(ImagesKwargs, total=False): tile_size (`int`, *optional*, defaults to `640`): The size of each local tile. Must match the model's query embedding size. maximum_pad_value (`int`, *optional*, defaults to `640`): - If `crop_to_patches` is `False` and `size.height/width` is larger than this value, - the image will be resized directly to `size.height/width` without padding. Otherwise, - images are resized and padded to `size.height/width` while preserving the aspect ratio. + If `crop_to_patches` is `False` and `max(size.height, size.width)` is smaller than or equal to this + value, the image is resized directly to a square of `max(size.height, size.width)` without preserving + the aspect ratio. Otherwise, the image is resized while preserving the aspect ratio and then padded to + a square with `background_color`. background_color (`list[int]`, *optional*, defaults to `[127, 127, 127]`): The background color for padding. """ diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 4e9fbd87b8d9..659720b4e9f6 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -93,9 +93,10 @@ class UnlimitedOcrImageProcessorKwargs(DeepseekOcr2ImageProcessorKwargs): tile_size (`int`, *optional*, defaults to `640`): The size of each local tile. Must match the model's query embedding size. maximum_pad_value (`int`, *optional*, defaults to `640`): - If `crop_to_patches` is `False` and `size.height/width` is larger than this value, - the image will be resized directly to `size.height/width` without padding. Otherwise, - images are resized and padded to `size.height/width` while preserving the aspect ratio. + If `crop_to_patches` is `False` and `max(size.height, size.width)` is smaller than or equal to this + value, the image is resized directly to a square of `max(size.height, size.width)` without preserving + the aspect ratio. Otherwise, the image is resized while preserving the aspect ratio and then padded to + a square with `background_color`. background_color (`list[int]`, *optional*, defaults to `[127, 127, 127]`): The background color for padding. """ @@ -396,9 +397,9 @@ class UnlimitedOcrTextConfig(DeepseekOcr2TextConfig): MLP type (`"dense"` or `"sparse"`) for each decoder layer, e.g. `["dense", "sparse", "sparse", ...]`. layer_types (`list[str]`, *optional*): Attention type for each decoder layer. Defaults to `"reference_sliding_attention"` on every layer. - use_sliding_window (`int`, defaults to `True`): + use_sliding_window (`bool`, *optional*, defaults to `True`): Whether to use reference sliding window attention. - sliding_window (`int`, *optional*, defaults to 128): + sliding_window (`int`, *optional*, defaults to `128`): Sliding window size for reference sliding window attention. If set, every token attends to the last `sliding_window` and all image and prompt tokens. From bf24d276a23d234612f36d4f1c90b48f86559376 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 28 Jul 2026 17:17:46 +0100 Subject: [PATCH 155/195] remove dead _input_embed_layer from vision encoder --- src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py | 1 - src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index f1cb8365c3d5..0bbf4d3d4c4a 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -891,7 +891,6 @@ class UnlimitedOcrVisionEncoder(UnlimitedOcrPreTrainedModel): config: UnlimitedOcrVisionConfig main_input_name = "patch_embeds" input_modalities = ("image",) - _input_embed_layer = "patch_embedding" _can_record_outputs = { "hidden_states": UnlimitedOcrEncoderLayer, "attentions": UnlimitedOcrAttention, diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 659720b4e9f6..ea1b78adf03a 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -604,6 +604,7 @@ def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: ) class UnlimitedOcrVisionEncoder(CLIPVisionModel): main_input_name = "patch_embeds" + _input_embed_layer = AttributeError() _can_record_outputs = { "hidden_states": UnlimitedOcrEncoderLayer, "attentions": UnlimitedOcrAttention, From 922231ce04ef291fd401eaa365dc72b0f354a178 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 29 Jul 2026 09:07:45 +0100 Subject: [PATCH 156/195] reuse base decode implementations and fix batch detections --- .../unlimited_ocr/modular_unlimited_ocr.py | 19 ++++++++----------- .../unlimited_ocr/processing_unlimited_ocr.py | 17 +++++++---------- 2 files changed, 15 insertions(+), 21 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index ea1b78adf03a..c0423bcb0db4 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -40,7 +40,7 @@ or_masks, ) from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling -from ...processing_utils import Unpack +from ...processing_utils import ProcessorMixin, Unpack from ...tokenization_utils_base import TextInput from ...utils import ( TensorType, @@ -238,15 +238,14 @@ def batch_decode(self, *args, return_detections: bool = False, **kwargs): Whether or not to also return the layout detections parsed from the decoded text. Returns: - `str` or `tuple[str, list[dict]]`: The decoded text. If `return_detections` is `True`, a tuple of the - decoded text and a list of detections, where every detection is a dictionary with the keys `region_type`, - `box` and `text`. Boxes are in [x1, y1, x2, y2] format with coordinates normalized to [0, 999]. + `list[str]` or `tuple[list[str], list[list[dict]]]`: The decoded text. If `return_detections` is `True`, + a tuple of the decoded text and the detections of every sequence, where every detection is a dictionary + with the keys `region_type`, `box` and `text`. Boxes are in [x1, y1, x2, y2] format with coordinates + normalized to [0, 999]. """ - if not hasattr(self, "tokenizer"): - raise ValueError(f"Cannot batch decode text: {self.__class__.__name__} has no tokenizer.") - decoded = self.tokenizer.batch_decode(*args, **kwargs) + decoded = ProcessorMixin.batch_decode(self, *args, **kwargs) if return_detections: - detections = self._parse_detections(decoded) + detections = [self._parse_detections(text) for text in decoded] return decoded, detections return decoded @@ -264,9 +263,7 @@ def decode(self, *args, return_detections: bool = False, **kwargs): decoded text and a list of detections, where every detection is a dictionary with the keys `region_type`, `box` and `text`. Boxes are in [x1, y1, x2, y2] format with coordinates normalized to [0, 999]. """ - if not hasattr(self, "tokenizer"): - raise ValueError(f"Cannot decode text: {self.__class__.__name__} has no tokenizer.") - decoded = self.tokenizer.decode(*args, **kwargs) + decoded = ProcessorMixin.decode(self, *args, **kwargs) if return_detections: detections = self._parse_detections(decoded) return decoded, detections diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index a1c771c680a5..ba8b2e750220 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -104,15 +104,14 @@ def batch_decode(self, *args, return_detections: bool = False, **kwargs): Whether or not to also return the layout detections parsed from the decoded text. Returns: - `str` or `tuple[str, list[dict]]`: The decoded text. If `return_detections` is `True`, a tuple of the - decoded text and a list of detections, where every detection is a dictionary with the keys `region_type`, - `box` and `text`. Boxes are in [x1, y1, x2, y2] format with coordinates normalized to [0, 999]. + `list[str]` or `tuple[list[str], list[list[dict]]]`: The decoded text. If `return_detections` is `True`, + a tuple of the decoded text and the detections of every sequence, where every detection is a dictionary + with the keys `region_type`, `box` and `text`. Boxes are in [x1, y1, x2, y2] format with coordinates + normalized to [0, 999]. """ - if not hasattr(self, "tokenizer"): - raise ValueError(f"Cannot batch decode text: {self.__class__.__name__} has no tokenizer.") - decoded = self.tokenizer.batch_decode(*args, **kwargs) + decoded = super().batch_decode(*args, **kwargs) if return_detections: - detections = self._parse_detections(decoded) + detections = [self._parse_detections(text) for text in decoded] return decoded, detections return decoded @@ -130,9 +129,7 @@ def decode(self, *args, return_detections: bool = False, **kwargs): decoded text and a list of detections, where every detection is a dictionary with the keys `region_type`, `box` and `text`. Boxes are in [x1, y1, x2, y2] format with coordinates normalized to [0, 999]. """ - if not hasattr(self, "tokenizer"): - raise ValueError(f"Cannot decode text: {self.__class__.__name__} has no tokenizer.") - decoded = self.tokenizer.decode(*args, **kwargs) + decoded = super().decode(*args, **kwargs) if return_detections: detections = self._parse_detections(decoded) return decoded, detections From 51afb7fd246298707f35ee58ba71f86386ff491e Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 29 Jul 2026 09:11:59 +0100 Subject: [PATCH 157/195] only check image placeholders when images are passed --- .../models/deepseek_ocr2/processing_deepseek_ocr2.py | 12 ++++++------ .../models/unlimited_ocr/processing_unlimited_ocr.py | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py index 61dfbea5aac0..369ebd9ce7e8 100644 --- a/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/processing_deepseek_ocr2.py @@ -77,12 +77,12 @@ def validate_inputs(self, images=None, text=None, videos=None, audio=None, **kwa if text is None: raise ValueError("You have to specify text.") - total_placeholders = sum(prompt.count(self.image_token) for prompt in text) - num_images = len(images) if images is not None else 0 - if total_placeholders != num_images: - raise ValueError( - f"Found {total_placeholders} placeholders across the batch, but have {num_images} flattened images." - ) + if images is not None: + total_placeholders = sum(prompt.count(self.image_token) for prompt in text) + if total_placeholders != len(images): + raise ValueError( + f"Found {total_placeholders} placeholders across the batch, but have {len(images)} flattened images." + ) def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> TextInput: size = kwargs.get("size") or self.image_processor.size diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index ba8b2e750220..8852c265f457 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -72,12 +72,12 @@ def validate_inputs(self, images=None, text=None, videos=None, audio=None, **kwa if text is None: raise ValueError("You have to specify text.") - total_placeholders = sum(prompt.count(self.image_token) for prompt in text) - num_images = len(images) if images is not None else 0 - if total_placeholders != num_images: - raise ValueError( - f"Found {total_placeholders} placeholders across the batch, but have {num_images} flattened images." - ) + if images is not None: + total_placeholders = sum(prompt.count(self.image_token) for prompt in text) + if total_placeholders != len(images): + raise ValueError( + f"Found {total_placeholders} placeholders across the batch, but have {len(images)} flattened images." + ) def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> TextInput: image_size = kwargs.get("size") or self.image_processor.size From bb9341d38ab6523db223a63b9af78403769874ff Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 29 Jul 2026 09:16:10 +0100 Subject: [PATCH 158/195] fix-repo --- docs/source/en/model_doc/unlimited_ocr.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 87f1ab98d72b..7a6e7b7606ac 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -16,7 +16,7 @@ limitations under the License. ⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer. --> -*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-07-27.* +*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-07-29.* # UnlimitedOcr From d4f5a8f453d28e035c5114362f7da70fcd584508 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 29 Jul 2026 09:36:20 +0100 Subject: [PATCH 159/195] improve mask handling --- .../unlimited_ocr/modeling_unlimited_ocr.py | 15 ++++++++++++--- .../models/unlimited_ocr/modular_unlimited_ocr.py | 15 ++++++++++++--- .../unlimited_ocr/test_modeling_unlimited_ocr.py | 11 +++++++++++ 3 files changed, 35 insertions(+), 6 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 0bbf4d3d4c4a..ee41e80750bc 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1724,11 +1724,16 @@ def create_reference_sliding_window_causal_mask( block_sequence_ids: torch.Tensor | None = None, layer_idx: int | None = None, ) -> torch.Tensor | BlockMask | None: - if past_key_values is None: + layer = None + if past_key_values is not None: + layer = next( + (layer for layer in past_key_values.layers if layer._layer_type == "reference_sliding_attention"), None + ) + + if layer is None: prefill_length = float("inf") kv_offset = 0 else: - layer = next(layer for layer in past_key_values.layers if layer._layer_type == "reference_sliding_attention") # A layer that is still prefilling has no reference states yet, all its states are prefill states prefill_length = float("inf") if layer.prefill_length is None else layer.prefill_length _, kv_offset = layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) @@ -1812,8 +1817,12 @@ def forward( causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), - "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), } + # The reference sliding window layers are not always activated depending on the config + if "reference_sliding_attention" in self.config.layer_types: + causal_mask_mapping["reference_sliding_attention"] = create_reference_sliding_window_causal_mask( + **mask_kwargs + ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index c0423bcb0db4..411900e05158 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -1048,11 +1048,16 @@ def create_reference_sliding_window_causal_mask( block_sequence_ids: torch.Tensor | None = None, layer_idx: int | None = None, ) -> torch.Tensor | BlockMask | None: - if past_key_values is None: + layer = None + if past_key_values is not None: + layer = next( + (layer for layer in past_key_values.layers if layer._layer_type == "reference_sliding_attention"), None + ) + + if layer is None: prefill_length = float("inf") kv_offset = 0 else: - layer = next(layer for layer in past_key_values.layers if layer._layer_type == "reference_sliding_attention") # A layer that is still prefilling has no reference states yet, all its states are prefill states prefill_length = float("inf") if layer.prefill_length is None else layer.prefill_length _, kv_offset = layer.get_mask_sizes(query_length=inputs_embeds.shape[1]) @@ -1115,8 +1120,12 @@ def forward( causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), - "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), } + # The reference sliding window layers are not always activated depending on the config + if "reference_sliding_attention" in self.config.layer_types: + causal_mask_mapping["reference_sliding_attention"] = create_reference_sliding_window_causal_mask( + **mask_kwargs + ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 681b66b4da4c..e11d95fbd318 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -208,6 +208,17 @@ def test_generate_static_cache_sliding_window_too_small_cache_full(self): """Continue from full cache""" self._check_generate_cache_sliding_window_too_small(cache_implementation="static", prefill_max_new_tokens=6) + def test_generate_without_sliding_window(self): + """With `use_sliding_window=False` every layer is a full attention layer.""" + model_tester = self.model_tester_class(self, use_sliding_window=False) + config, inputs_dict = model_tester.prepare_config_and_inputs_for_common() + self.assertEqual(config.text_config.layer_types, ["full_attention"] * config.text_config.num_hidden_layers) + + for model_class in self.all_generative_model_classes: + model = model_class(config).to(torch_device).eval() + out = model.generate(**inputs_dict, max_new_tokens=3, do_sample=False, return_dict_in_generate=True) + self.assertTrue(all(not layer.is_sliding for layer in out.past_key_values.layers)) + @require_torch class UnlimitedOcrIntegrationTest(unittest.TestCase): From 664f30b81edc30b10aed842467e0af0442adc40f Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 29 Jul 2026 11:13:17 +0100 Subject: [PATCH 160/195] replace end_prefill with set_prefill_length --- .../unlimited_ocr/generation_unlimited_ocr.py | 28 +++++---- .../unlimited_ocr/modeling_unlimited_ocr.py | 41 ++++++++---- .../unlimited_ocr/modular_unlimited_ocr.py | 41 ++++++++---- .../test_modeling_unlimited_ocr.py | 62 +++++++++++++++++++ 4 files changed, 135 insertions(+), 37 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py index 0c2474f2d125..9b4ef5c8849d 100644 --- a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py @@ -62,22 +62,28 @@ def _prefill( model_kwargs: dict, is_first_iteration: bool = True, ): - outputs = super()._prefill( + # A prefill can span several forward passes when chunked, which the cache layers cannot tell apart from + # decode steps, so declare its length before the first one. + past_key_values = model_kwargs.get("past_key_values") + if past_key_values is not None and past_key_values.get_seq_length() == 0: + prefill_length = self._get_prefill_length(input_ids, model_kwargs) + for layer in past_key_values.layers: + if layer._layer_type == "reference_sliding_attention": + layer.set_prefill_length(prefill_length) + return super()._prefill( input_ids, generation_config, model_kwargs, is_first_iteration=is_first_iteration, ) - # The cache layers cannot tell prefill and decode apart on their own, so mark the prefill as complete once - # all prefill forward passes are done (a prefill can span several forward passes when chunked). - past_key_values = model_kwargs.get("past_key_values") - if past_key_values is None: - past_key_values = getattr(outputs, "past_key_values", None) - if past_key_values is not None: - for layer in past_key_values.layers: - if layer._layer_type == "reference_sliding_attention": - layer.end_prefill() - return outputs + + def _get_prefill_length(self, input_ids: torch.LongTensor, model_kwargs: dict) -> int: + """Number of tokens the prefill caches, including padding tokens.""" + inputs_embeds = model_kwargs.get("inputs_embeds") + # `input_ids` is empty when generating from embeddings + if inputs_embeds is not None: + return inputs_embeds.shape[1] + return input_ids.shape[1] def _get_logits_processor( self, diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index ee41e80750bc..fa5dd96c0471 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1295,7 +1295,9 @@ class UnlimitedOcrDynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): Once ``sliding_window`` decode tokens have accumulated, the oldest decode tokens are evicted and replaced by the most recent ones. The prefill tokens always remain in the cache. - The layer is in the prefill phase until ``end_prefill`` is called. + The states of the first ``update`` are the prefill states, unless their length is declared upfront with + ``set_prefill_length``, which is required when the prefill spans several ``update`` calls, e.g. for chunked + prefill. """ _layer_type = "reference_sliding_attention" @@ -1319,7 +1321,10 @@ def update( full_value_states = torch.cat([self.values, value_states], dim=-2) # Prefill - if self.prefill_length is None or self.record_past: + if self.record_past or self.prefill_length is None: + # While recording the past, all states are kept and the layer stays in the prefill phase + if not self.record_past: + self.prefill_length = self.cumulative_length self.keys = full_key_states self.values = full_value_states return self.keys, self.values @@ -1375,10 +1380,10 @@ def get_max_length(self) -> int: return -1 return self.prefill_length + self.sliding_window - def end_prefill(self) -> None: - """Mark the prefill phase as complete.""" - if self.prefill_length is None: - self.prefill_length = self.get_seq_length() + def set_prefill_length(self, prefill_length: int) -> None: + """Declare how many leading tokens are prefill states, before they are cached.""" + if self.prefill_length is None and not self.record_past: + self.prefill_length = prefill_length def reset(self) -> None: super().reset() @@ -1423,7 +1428,9 @@ class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): that are never eviced from the cache. Slots ``[prefill_length, prefill_length + sliding_window)`` hold the sliding window decode slots where the oldest entries are always replaced by the newest ones. - The layer is in the prefill phase until ``end_prefill`` is called. + The states of the first ``update`` are the prefill states, unless their length is declared upfront with + ``set_prefill_length``, which is required when the prefill spans several ``update`` calls, e.g. for chunked + prefill. Args: max_cache_len (`int`): @@ -1462,7 +1469,7 @@ def update( kv_length = key_states.shape[-2] # Prefill - if self.prefill_length is None: + if self.is_prefilling: # Chunked prefill: the buffer was sized for the first chunk only, so grow it to hold all prefill seen so far plus this chunk plus the reserved window. required_length = min(self.max_cache_len, self.cumulative_length_int + kv_length + self.sliding_window) if self.keys.shape[-2] < required_length: @@ -1482,6 +1489,9 @@ def update( self.cumulative_length_int += kv_length self.cumulative_length.add_(kv_length) + if self.prefill_length is None: + self.prefill_length = self.cumulative_length_int + # Very important to return the `self` tensors here, as they have the static dynamo address return self.keys, self.values @@ -1571,8 +1581,8 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: ) kv_offset = 0 - # Prefill - if self.prefill_length is None: + # Prefill: the buffer is not necessarily allocated yet, so its size cannot be used + if self.is_prefilling: kv_length = min(self.max_cache_len, self.cumulative_length_int + query_length + self.sliding_window) # Decode: cache is already full elif is_full: @@ -1596,10 +1606,15 @@ def get_max_length(self) -> int: return self.max_cache_len return min(self.max_cache_len, self.prefill_length + self.sliding_window) - def end_prefill(self) -> None: - """Mark the prefill phase as complete.""" + @property + def is_prefilling(self) -> bool: + """Whether the prefill states are still being cached.""" + return self.prefill_length is None or self.cumulative_length_int < self.prefill_length + + def set_prefill_length(self, prefill_length: int) -> None: + """Declare how many leading tokens are prefill states, before they are cached.""" if self.prefill_length is None: - self.prefill_length = self.cumulative_length_int + self.prefill_length = prefill_length def reset(self) -> None: super().reset() diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 411900e05158..d94b12057938 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -684,7 +684,9 @@ class UnlimitedOcrDynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): Once ``sliding_window`` decode tokens have accumulated, the oldest decode tokens are evicted and replaced by the most recent ones. The prefill tokens always remain in the cache. - The layer is in the prefill phase until ``end_prefill`` is called. + The states of the first ``update`` are the prefill states, unless their length is declared upfront with + ``set_prefill_length``, which is required when the prefill spans several ``update`` calls, e.g. for chunked + prefill. """ _layer_type = "reference_sliding_attention" @@ -708,7 +710,10 @@ def update( full_value_states = torch.cat([self.values, value_states], dim=-2) # Prefill - if self.prefill_length is None or self.record_past: + if self.record_past or self.prefill_length is None: + # While recording the past, all states are kept and the layer stays in the prefill phase + if not self.record_past: + self.prefill_length = self.cumulative_length self.keys = full_key_states self.values = full_value_states return self.keys, self.values @@ -764,10 +769,10 @@ def get_max_length(self) -> int: return -1 return self.prefill_length + self.sliding_window - def end_prefill(self) -> None: - """Mark the prefill phase as complete.""" - if self.prefill_length is None: - self.prefill_length = self.get_seq_length() + def set_prefill_length(self, prefill_length: int) -> None: + """Declare how many leading tokens are prefill states, before they are cached.""" + if self.prefill_length is None and not self.record_past: + self.prefill_length = prefill_length def reset(self) -> None: super().reset() @@ -812,7 +817,9 @@ class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): that are never eviced from the cache. Slots ``[prefill_length, prefill_length + sliding_window)`` hold the sliding window decode slots where the oldest entries are always replaced by the newest ones. - The layer is in the prefill phase until ``end_prefill`` is called. + The states of the first ``update`` are the prefill states, unless their length is declared upfront with + ``set_prefill_length``, which is required when the prefill spans several ``update`` calls, e.g. for chunked + prefill. Args: max_cache_len (`int`): @@ -851,7 +858,7 @@ def update( kv_length = key_states.shape[-2] # Prefill - if self.prefill_length is None: + if self.is_prefilling: # Chunked prefill: the buffer was sized for the first chunk only, so grow it to hold all prefill seen so far plus this chunk plus the reserved window. required_length = min(self.max_cache_len, self.cumulative_length_int + kv_length + self.sliding_window) if self.keys.shape[-2] < required_length: @@ -871,6 +878,9 @@ def update( self.cumulative_length_int += kv_length self.cumulative_length.add_(kv_length) + if self.prefill_length is None: + self.prefill_length = self.cumulative_length_int + # Very important to return the `self` tensors here, as they have the static dynamo address return self.keys, self.values @@ -960,8 +970,8 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: ) kv_offset = 0 - # Prefill - if self.prefill_length is None: + # Prefill: the buffer is not necessarily allocated yet, so its size cannot be used + if self.is_prefilling: kv_length = min(self.max_cache_len, self.cumulative_length_int + query_length + self.sliding_window) # Decode: cache is already full elif is_full: @@ -985,10 +995,15 @@ def get_max_length(self) -> int: return self.max_cache_len return min(self.max_cache_len, self.prefill_length + self.sliding_window) - def end_prefill(self) -> None: - """Mark the prefill phase as complete.""" + @property + def is_prefilling(self) -> bool: + """Whether the prefill states are still being cached.""" + return self.prefill_length is None or self.cumulative_length_int < self.prefill_length + + def set_prefill_length(self, prefill_length: int) -> None: + """Declare how many leading tokens are prefill states, before they are cached.""" if self.prefill_length is None: - self.prefill_length = self.cumulative_length_int + self.prefill_length = prefill_length def reset(self) -> None: super().reset() diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index e11d95fbd318..1a6b9cf7df9e 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -38,6 +38,8 @@ import torch from transformers import ( + DynamicCache, + StaticCache, UnlimitedOcrForConditionalGeneration, UnlimitedOcrModel, ) @@ -208,6 +210,66 @@ def test_generate_static_cache_sliding_window_too_small_cache_full(self): """Continue from full cache""" self._check_generate_cache_sliding_window_too_small(cache_implementation="static", prefill_max_new_tokens=6) + def _check_manual_forward_cache(self, cache_implementation: str): + """Test that the states of the first forward pass are the prefill states if set_prefill_length is not explicitly called.""" + for model_class in self.all_generative_model_classes: + config, inputs_dict = self.prepare_config_and_inputs_for_generate() + sliding_window = config.text_config.sliding_window + num_decode_steps = 2 * sliding_window + + model = model_class(config).to(torch_device).eval() + input_ids = inputs_dict["input_ids"] + prompt_length = input_ids.shape[1] + past_key_values = ( + StaticCache(config=config.get_text_config(), max_cache_len=prompt_length + num_decode_steps) + if cache_implementation == "static" + else DynamicCache(config=config.get_text_config()) + ) + + next_tokens = input_ids + with torch.no_grad(): + for _ in range(num_decode_steps + 1): + outputs = model(input_ids=next_tokens, past_key_values=past_key_values, use_cache=True) + next_tokens = outputs.logits[:, -1:].argmax(dim=-1) + + reference_layers = [ + layer for layer in past_key_values.layers if layer._layer_type == "reference_sliding_attention" + ] + self.assertGreater(len(reference_layers), 0) + for layer in reference_layers: + self.assertEqual(layer.prefill_length, prompt_length) + + def test_manual_forward_dynamic_cache(self): + self._check_manual_forward_cache(cache_implementation="dynamic") + + def test_manual_forward_loop_static_cache(self): + self._check_manual_forward_cache(cache_implementation="static") + + def test_generate_cache_chunked_prefill(self): + """Test that a chunked prefill keeps all prompt tokens as prefill, not only the first chunk.""" + for model_class in self.all_generative_model_classes: + config, inputs_dict = self.prepare_config_and_inputs_for_generate() + model = model_class(config).to(torch_device).eval() + + input_ids = inputs_dict["input_ids"] + prompt_length = input_ids.shape[1] + out = model.generate( + input_ids=input_ids, + attention_mask=torch.ones_like(input_ids), + max_new_tokens=3, + min_new_tokens=3, + do_sample=False, + prefill_chunk_size=prompt_length // 2, + return_dict_in_generate=True, + ) + + reference_layers = [ + layer for layer in out.past_key_values.layers if layer._layer_type == "reference_sliding_attention" + ] + self.assertGreater(len(reference_layers), 0) + for layer in reference_layers: + self.assertEqual(layer.prefill_length, prompt_length) + def test_generate_without_sliding_window(self): """With `use_sliding_window=False` every layer is a full attention layer.""" model_tester = self.model_tester_class(self, use_sliding_window=False) From 3774c0ecd0cdc2cf00f6eea6aef21ac5f4e81cf7 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:43:33 +0100 Subject: [PATCH 161/195] simplify --- .../unlimited_ocr/modeling_unlimited_ocr.py | 48 +++++++++---------- .../unlimited_ocr/modular_unlimited_ocr.py | 48 +++++++++---------- 2 files changed, 44 insertions(+), 52 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index fa5dd96c0471..739b85786638 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1320,19 +1320,16 @@ def update( full_key_states = torch.cat([self.keys, key_states], dim=-2) full_value_states = torch.cat([self.values, value_states], dim=-2) - # Prefill - if self.record_past or self.prefill_length is None: - # While recording the past, all states are kept and the layer stays in the prefill phase - if not self.record_past: - self.prefill_length = self.cumulative_length - self.keys = full_key_states - self.values = full_value_states - return self.keys, self.values - # Cache growing - if self.cumulative_length <= self.prefill_length + self.sliding_window - 1: + if ( + self.record_past + or self.prefill_length is None + or self.cumulative_length <= self.prefill_length + self.sliding_window - 1 + ): self.keys = full_key_states self.values = full_value_states + if self.prefill_length is None: + self.prefill_length = self.cumulative_length # Cache full elif self.keys.shape[-2] == self.prefill_length + self.sliding_window - 1: self.keys[:, :, -self.sliding_window + 1 :].copy_(full_key_states[:, :, -self.sliding_window + 1 :, :]) @@ -1382,7 +1379,7 @@ def get_max_length(self) -> int: def set_prefill_length(self, prefill_length: int) -> None: """Declare how many leading tokens are prefill states, before they are cached.""" - if self.prefill_length is None and not self.record_past: + if self.prefill_length is None: self.prefill_length = prefill_length def reset(self) -> None: @@ -1441,8 +1438,8 @@ class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): _layer_type = "reference_sliding_attention" - def __init__(self, max_cache_len: int, sliding_window: int): - super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window) + def __init__(self, max_cache_len: int, sliding_window: int, **kwargs): + super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window, **kwargs) # Keep `max_cache_len` as max value for length bookkeeping. # The physical buffer doesn't exceed `prefill_length + sliding_window`. self.max_cache_len = max_cache_len @@ -1469,11 +1466,15 @@ def update( kv_length = key_states.shape[-2] # Prefill - if self.is_prefilling: - # Chunked prefill: the buffer was sized for the first chunk only, so grow it to hold all prefill seen so far plus this chunk plus the reserved window. - required_length = min(self.max_cache_len, self.cumulative_length_int + kv_length + self.sliding_window) + if self.prefill_length is None or self.cumulative_length_int < self.prefill_length: + if self.prefill_length is None: + self.prefill_length = self.cumulative_length_int + kv_length + + # Resize buffer if necessary + required_length = min(self.max_cache_len, self.prefill_length + self.sliding_window) if self.keys.shape[-2] < required_length: self._allocate_key_value_buffers(required_length, copy_existing=True) + # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length @@ -1489,9 +1490,6 @@ def update( self.cumulative_length_int += kv_length self.cumulative_length.add_(kv_length) - if self.prefill_length is None: - self.prefill_length = self.cumulative_length_int - # Very important to return the `self` tensors here, as they have the static dynamo address return self.keys, self.values @@ -1582,8 +1580,11 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: kv_offset = 0 # Prefill: the buffer is not necessarily allocated yet, so its size cannot be used - if self.is_prefilling: - kv_length = min(self.max_cache_len, self.cumulative_length_int + query_length + self.sliding_window) + if self.prefill_length is None or self.cumulative_length_int < self.prefill_length: + prefill_length = ( + self.cumulative_length_int + query_length if self.prefill_length is None else self.prefill_length + ) + kv_length = min(self.max_cache_len, prefill_length + self.sliding_window) # Decode: cache is already full elif is_full: kv_offset = max(self.cumulative_length_int - self.prefill_length - self.sliding_window + 1, 0) @@ -1606,11 +1607,6 @@ def get_max_length(self) -> int: return self.max_cache_len return min(self.max_cache_len, self.prefill_length + self.sliding_window) - @property - def is_prefilling(self) -> bool: - """Whether the prefill states are still being cached.""" - return self.prefill_length is None or self.cumulative_length_int < self.prefill_length - def set_prefill_length(self, prefill_length: int) -> None: """Declare how many leading tokens are prefill states, before they are cached.""" if self.prefill_length is None: diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index d94b12057938..6cc3a98ca17c 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -709,19 +709,16 @@ def update( full_key_states = torch.cat([self.keys, key_states], dim=-2) full_value_states = torch.cat([self.values, value_states], dim=-2) - # Prefill - if self.record_past or self.prefill_length is None: - # While recording the past, all states are kept and the layer stays in the prefill phase - if not self.record_past: - self.prefill_length = self.cumulative_length - self.keys = full_key_states - self.values = full_value_states - return self.keys, self.values - # Cache growing - if self.cumulative_length <= self.prefill_length + self.sliding_window - 1: + if ( + self.record_past + or self.prefill_length is None + or self.cumulative_length <= self.prefill_length + self.sliding_window - 1 + ): self.keys = full_key_states self.values = full_value_states + if self.prefill_length is None: + self.prefill_length = self.cumulative_length # Cache full elif self.keys.shape[-2] == self.prefill_length + self.sliding_window - 1: self.keys[:, :, -self.sliding_window + 1 :].copy_(full_key_states[:, :, -self.sliding_window + 1 :, :]) @@ -771,7 +768,7 @@ def get_max_length(self) -> int: def set_prefill_length(self, prefill_length: int) -> None: """Declare how many leading tokens are prefill states, before they are cached.""" - if self.prefill_length is None and not self.record_past: + if self.prefill_length is None: self.prefill_length = prefill_length def reset(self) -> None: @@ -830,8 +827,8 @@ class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): _layer_type = "reference_sliding_attention" - def __init__(self, max_cache_len: int, sliding_window: int): - super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window) + def __init__(self, max_cache_len: int, sliding_window: int, **kwargs): + super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window, **kwargs) # Keep `max_cache_len` as max value for length bookkeeping. # The physical buffer doesn't exceed `prefill_length + sliding_window`. self.max_cache_len = max_cache_len @@ -858,11 +855,15 @@ def update( kv_length = key_states.shape[-2] # Prefill - if self.is_prefilling: - # Chunked prefill: the buffer was sized for the first chunk only, so grow it to hold all prefill seen so far plus this chunk plus the reserved window. - required_length = min(self.max_cache_len, self.cumulative_length_int + kv_length + self.sliding_window) + if self.prefill_length is None or self.cumulative_length_int < self.prefill_length: + if self.prefill_length is None: + self.prefill_length = self.cumulative_length_int + kv_length + + # Resize buffer if necessary + required_length = min(self.max_cache_len, self.prefill_length + self.sliding_window) if self.keys.shape[-2] < required_length: self._allocate_key_value_buffers(required_length, copy_existing=True) + # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length @@ -878,9 +879,6 @@ def update( self.cumulative_length_int += kv_length self.cumulative_length.add_(kv_length) - if self.prefill_length is None: - self.prefill_length = self.cumulative_length_int - # Very important to return the `self` tensors here, as they have the static dynamo address return self.keys, self.values @@ -971,8 +969,11 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: kv_offset = 0 # Prefill: the buffer is not necessarily allocated yet, so its size cannot be used - if self.is_prefilling: - kv_length = min(self.max_cache_len, self.cumulative_length_int + query_length + self.sliding_window) + if self.prefill_length is None or self.cumulative_length_int < self.prefill_length: + prefill_length = ( + self.cumulative_length_int + query_length if self.prefill_length is None else self.prefill_length + ) + kv_length = min(self.max_cache_len, prefill_length + self.sliding_window) # Decode: cache is already full elif is_full: kv_offset = max(self.cumulative_length_int - self.prefill_length - self.sliding_window + 1, 0) @@ -995,11 +996,6 @@ def get_max_length(self) -> int: return self.max_cache_len return min(self.max_cache_len, self.prefill_length + self.sliding_window) - @property - def is_prefilling(self) -> bool: - """Whether the prefill states are still being cached.""" - return self.prefill_length is None or self.cumulative_length_int < self.prefill_length - def set_prefill_length(self, prefill_length: int) -> None: """Declare how many leading tokens are prefill states, before they are cached.""" if self.prefill_length is None: From 96d4a29403f419301e77c3cfbcdc9dfeb2143a2b Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:57:44 +0100 Subject: [PATCH 162/195] try super.update --- .../unlimited_ocr/modeling_unlimited_ocr.py | 100 +++++------------- .../unlimited_ocr/modular_unlimited_ocr.py | 100 +++++------------- 2 files changed, 48 insertions(+), 152 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 739b85786638..89c314945a0e 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1493,84 +1493,32 @@ def update( # Very important to return the `self` tensors here, as they have the static dynamo address return self.keys, self.values - # Everything below mirrors `StaticSlidingWindowLayer.update` with the difference that the sliding window - # is in the `sliding_window` slots right after the prefill instead of spanning the whole buffer. - window_start = self.prefill_length - current_length = self.cumulative_length_int - self.prefill_length - is_full = current_length >= self.sliding_window - # Update it now that we saved the value above - self.cumulative_length_int += kv_length - - if is_full: - # In general, we should use a much simpler `cat` here as well, independently of the states size. However, - # dynamo is currently bugged when doing it - see https://github.com/pytorch/pytorch/issues/159855 for more details - if key_states.shape[-2] == 1: - # Roll the window region to the left by 1 position - new_keys = self.keys[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) - new_values = self.values[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) - # Overwrite the last position with new states - # (note: very important to use a tensor to index here, see https://github.com/pytorch/pytorch/issues/159855) - index = torch.tensor([-1], dtype=int, device=self.device) - new_keys[:, :, index] = key_states - new_values[:, :, index] = value_states - - # Copy back into `self` (do not just assign again) in order to keep the static dynamo address - self.keys[:, :, window_start : window_start + self.sliding_window, :].copy_(new_keys) - self.values[:, :, window_start : window_start + self.sliding_window, :].copy_(new_values) - - # Very important to return the `self` tensors here, as they have the static dynamo address - return self.keys, self.values - # Already full but using more than 1 new token (e.g. chat continuation, etc...) - else: - full_key_states = torch.cat( - ( - self.keys[:, :, :window_start, :], - self.keys[:, :, window_start + 1 : window_start + self.sliding_window, :], - key_states, - ), - dim=-2, - ) - full_value_states = torch.cat( - ( - self.values[:, :, :window_start, :], - self.values[:, :, window_start + 1 : window_start + self.sliding_window, :], - value_states, - ), - dim=-2, - ) - # Not yet full, but becoming full on this update - elif current_length + kv_length > self.sliding_window: - full_key_states = torch.cat((self.keys[:, :, : window_start + current_length, :], key_states), dim=-2) - full_value_states = torch.cat( - (self.values[:, :, : window_start + current_length, :], value_states), dim=-2 - ) - else: - # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs - # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten. - cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length - try: - self.keys.index_copy_(2, cache_position, key_states) - self.values.index_copy_(2, cache_position, value_states) - except NotImplementedError: - # Fallback for devices like MPS where index_copy_ might not be supported. - self.keys[:, :, cache_position] = key_states - self.values[:, :, cache_position] = value_states - - # Update the tensor version of the length in-place (we don't need to update it if we are already outside - # of this branch, as we don't need the tensor anymore) - self.cumulative_length.add_(kv_length) - - # Very important to return the `self` tensors here, as they have the static dynamo address + # Decode + # Call `StaticSlidingWindowLayer.update` to reduce code duplication. This requires setting temporary + # attributes to match expectations of parent class. + keys, values = self.keys, self.values + window = slice(self.prefill_length, self.prefill_length + self.sliding_window) + sliding_keys_buffer, sliding_values_buffer = keys[:, :, window, :], values[:, :, window, :] + self.keys, self.values = sliding_keys_buffer, sliding_values_buffer + self.cumulative_length -= self.prefill_length + self.cumulative_length_int -= self.prefill_length + max_cache_len, self.max_cache_len = self.max_cache_len, self.sliding_window + try: + sliding_keys, sliding_values = super().update(key_states, value_states, *args, **kwargs) + finally: + self.keys, self.values = keys, values + self.cumulative_length += self.prefill_length + self.cumulative_length_int += self.prefill_length + self.max_cache_len = max_cache_len + + # The parent returned the original buffers + if sliding_keys is sliding_keys_buffer: return self.keys, self.values - - # We only cache the last `sliding_window` tokens - self.keys[:, :, window_start : window_start + self.sliding_window, :].copy_( - full_key_states[:, :, -self.sliding_window :, :] + # The parent returned concatenated states + return ( + torch.cat((self.keys[:, :, : self.prefill_length, :], sliding_keys), dim=-2), + torch.cat((self.values[:, :, : self.prefill_length, :], sliding_values), dim=-2), ) - self.values[:, :, window_start : window_start + self.sliding_window, :].copy_( - full_value_states[:, :, -self.sliding_window :, :] - ) - return full_key_states, full_value_states def get_mask_sizes(self, query_length: int) -> tuple[int, int]: """Return the length and offset of the cache, used to generate the attention mask""" diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 6cc3a98ca17c..5aa07e96a4f3 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -882,84 +882,32 @@ def update( # Very important to return the `self` tensors here, as they have the static dynamo address return self.keys, self.values - # Everything below mirrors `StaticSlidingWindowLayer.update` with the difference that the sliding window - # is in the `sliding_window` slots right after the prefill instead of spanning the whole buffer. - window_start = self.prefill_length - current_length = self.cumulative_length_int - self.prefill_length - is_full = current_length >= self.sliding_window - # Update it now that we saved the value above - self.cumulative_length_int += kv_length - - if is_full: - # In general, we should use a much simpler `cat` here as well, independently of the states size. However, - # dynamo is currently bugged when doing it - see https://github.com/pytorch/pytorch/issues/159855 for more details - if key_states.shape[-2] == 1: - # Roll the window region to the left by 1 position - new_keys = self.keys[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) - new_values = self.values[:, :, window_start : window_start + self.sliding_window, :].roll(-1, dims=-2) - # Overwrite the last position with new states - # (note: very important to use a tensor to index here, see https://github.com/pytorch/pytorch/issues/159855) - index = torch.tensor([-1], dtype=int, device=self.device) - new_keys[:, :, index] = key_states - new_values[:, :, index] = value_states - - # Copy back into `self` (do not just assign again) in order to keep the static dynamo address - self.keys[:, :, window_start : window_start + self.sliding_window, :].copy_(new_keys) - self.values[:, :, window_start : window_start + self.sliding_window, :].copy_(new_values) - - # Very important to return the `self` tensors here, as they have the static dynamo address - return self.keys, self.values - # Already full but using more than 1 new token (e.g. chat continuation, etc...) - else: - full_key_states = torch.cat( - ( - self.keys[:, :, :window_start, :], - self.keys[:, :, window_start + 1 : window_start + self.sliding_window, :], - key_states, - ), - dim=-2, - ) - full_value_states = torch.cat( - ( - self.values[:, :, :window_start, :], - self.values[:, :, window_start + 1 : window_start + self.sliding_window, :], - value_states, - ), - dim=-2, - ) - # Not yet full, but becoming full on this update - elif current_length + kv_length > self.sliding_window: - full_key_states = torch.cat((self.keys[:, :, : window_start + current_length, :], key_states), dim=-2) - full_value_states = torch.cat( - (self.values[:, :, : window_start + current_length, :], value_states), dim=-2 - ) - else: - # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs - # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten. - cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length - try: - self.keys.index_copy_(2, cache_position, key_states) - self.values.index_copy_(2, cache_position, value_states) - except NotImplementedError: - # Fallback for devices like MPS where index_copy_ might not be supported. - self.keys[:, :, cache_position] = key_states - self.values[:, :, cache_position] = value_states - - # Update the tensor version of the length in-place (we don't need to update it if we are already outside - # of this branch, as we don't need the tensor anymore) - self.cumulative_length.add_(kv_length) - - # Very important to return the `self` tensors here, as they have the static dynamo address + # Decode + # Call `StaticSlidingWindowLayer.update` to reduce code duplication. This requires setting temporary + # attributes to match expectations of parent class. + keys, values = self.keys, self.values + window = slice(self.prefill_length, self.prefill_length + self.sliding_window) + sliding_keys_buffer, sliding_values_buffer = keys[:, :, window, :], values[:, :, window, :] + self.keys, self.values = sliding_keys_buffer, sliding_values_buffer + self.cumulative_length -= self.prefill_length + self.cumulative_length_int -= self.prefill_length + max_cache_len, self.max_cache_len = self.max_cache_len, self.sliding_window + try: + sliding_keys, sliding_values = super().update(key_states, value_states, *args, **kwargs) + finally: + self.keys, self.values = keys, values + self.cumulative_length += self.prefill_length + self.cumulative_length_int += self.prefill_length + self.max_cache_len = max_cache_len + + # The parent returned the original buffers + if sliding_keys is sliding_keys_buffer: return self.keys, self.values - - # We only cache the last `sliding_window` tokens - self.keys[:, :, window_start : window_start + self.sliding_window, :].copy_( - full_key_states[:, :, -self.sliding_window :, :] + # The parent returned concatenated states + return ( + torch.cat((self.keys[:, :, : self.prefill_length, :], sliding_keys), dim=-2), + torch.cat((self.values[:, :, : self.prefill_length, :], sliding_values), dim=-2), ) - self.values[:, :, window_start : window_start + self.sliding_window, :].copy_( - full_value_states[:, :, -self.sliding_window :, :] - ) - return full_key_states, full_value_states def get_mask_sizes(self, query_length: int) -> tuple[int, int]: """Return the length and offset of the cache, used to generate the attention mask""" From 3fdba03cf9e4982c23a8b490379c2021cd77f305 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:06:30 +0100 Subject: [PATCH 163/195] cleanup --- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 5 ----- .../models/unlimited_ocr/modular_unlimited_ocr.py | 5 ----- 2 files changed, 10 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 89c314945a0e..c7e0fc3a5fba 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1571,11 +1571,6 @@ def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tens self.k_head_dim = key_states.shape[-1] self.cumulative_length = self.cumulative_length.to(self.device) - # Note: `mark_static_address` is used to tag the tensors as a fixed data pointer, preventing compiled graph - # breaks or cudagraph skips due to inplace mutations when updating the cache. However, it is not supported when - # tracing the graph, so we skip it in this case. As prefill should never be compiled, this is not an issue and it - # will still be run (except when users compile prefill explicitly, but this should be avoided!) - # Without this, we cannot use cudagraphs!! if not is_torchdynamo_compiling(): torch._dynamo.mark_static_address(self.cumulative_length) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 5aa07e96a4f3..6ca9062caf40 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -960,11 +960,6 @@ def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tens self.k_head_dim = key_states.shape[-1] self.cumulative_length = self.cumulative_length.to(self.device) - # Note: `mark_static_address` is used to tag the tensors as a fixed data pointer, preventing compiled graph - # breaks or cudagraph skips due to inplace mutations when updating the cache. However, it is not supported when - # tracing the graph, so we skip it in this case. As prefill should never be compiled, this is not an issue and it - # will still be run (except when users compile prefill explicitly, but this should be avoided!) - # Without this, we cannot use cudagraphs!! if not is_torchdynamo_compiling(): torch._dynamo.mark_static_address(self.cumulative_length) From 745a3e270d1142dba9751ca5186c0d9178b16b76 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:32:04 +0100 Subject: [PATCH 164/195] fix kwargs --- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 4 ++-- .../models/unlimited_ocr/modular_unlimited_ocr.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index c7e0fc3a5fba..081e58cb6f71 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1302,8 +1302,8 @@ class UnlimitedOcrDynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): _layer_type = "reference_sliding_attention" - def __init__(self, config: PreTrainedConfig | None = None, sliding_window: int | None = None): - super().__init__(config=config, sliding_window=sliding_window) + def __init__(self, sliding_window: int | None = None, **kwargs): + super().__init__(sliding_window=sliding_window, **kwargs) self.prefill_length: int | None = None def update( diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 6ca9062caf40..922e3082c0c8 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -691,8 +691,8 @@ class UnlimitedOcrDynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): _layer_type = "reference_sliding_attention" - def __init__(self, config: PreTrainedConfig | None = None, sliding_window: int | None = None): - super().__init__(config=config, sliding_window=sliding_window) + def __init__(self, sliding_window: int | None = None, **kwargs): + super().__init__(sliding_window=sliding_window, **kwargs) self.prefill_length: int | None = None def update( From 127656316757b265d4a0d37669e9643cdcfca5a1 Mon Sep 17 00:00:00 2001 From: guarin Date: Wed, 29 Jul 2026 17:44:02 +0100 Subject: [PATCH 165/195] move logits processor test to test_logits_process.py Co-Authored-By: Claude Opus 5 (1M context) --- tests/generation/test_logits_process.py | 16 ++++++++ .../test_generation_unlimited_ocr.py | 38 ------------------- 2 files changed, 16 insertions(+), 38 deletions(-) delete mode 100644 tests/models/unlimited_ocr/test_generation_unlimited_ocr.py diff --git a/tests/generation/test_logits_process.py b/tests/generation/test_logits_process.py index 83f170a4d555..12ea0f8b880f 100644 --- a/tests/generation/test_logits_process.py +++ b/tests/generation/test_logits_process.py @@ -61,6 +61,9 @@ DiaEOSChannelFilterLogitsProcessor, DiaEOSDelayPatternLogitsProcessor, ) + from transformers.models.unlimited_ocr.generation_unlimited_ocr import ( + UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor, + ) @require_torch @@ -649,6 +652,19 @@ def test_no_repeat_ngram_dist_processor(self): self.assertFalse(torch.all(scores == filtered_scores_2_gram)) self.assertFalse(torch.all(scores == filtered_scores_3_gram)) + def test_sliding_window_no_repeat_ngram_dist_processor(self): + vocab_size = 3 + # The (0, 1) bigram appears at the start, so a full-sequence processor would forbid token 1 + # after the trailing 0. A small window should not see that early bigram. + input_ids = torch.tensor([[0, 1, 2, 0]], device=torch_device, dtype=torch.long) + scores = torch.zeros((1, vocab_size), device=torch_device, dtype=torch.float) + + small_window = UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor(ngram_size=2, window_size=2) + full_window = UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor(ngram_size=2, window_size=4) + + self.assertListEqual(torch.isinf(small_window(input_ids, scores.clone())).tolist(), [[False, False, False]]) + self.assertListEqual(torch.isinf(full_window(input_ids, scores.clone())).tolist(), [[False, True, False]]) + def test_encoder_no_repeat_ngram_dist_processor(self): vocab_size = 3 num_beams = 2 diff --git a/tests/models/unlimited_ocr/test_generation_unlimited_ocr.py b/tests/models/unlimited_ocr/test_generation_unlimited_ocr.py deleted file mode 100644 index 4f7e3a81491b..000000000000 --- a/tests/models/unlimited_ocr/test_generation_unlimited_ocr.py +++ /dev/null @@ -1,38 +0,0 @@ -# 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 unittest - -import torch - -from transformers.models.unlimited_ocr.generation_unlimited_ocr import ( - UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor, -) -from transformers.testing_utils import require_torch, torch_device - - -@require_torch -class UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessorTest(unittest.TestCase): - def test_window_limits_ngram_lookup(self): - vocab_size = 3 - # The (0, 1) bigram appears at the start, so a full-sequence processor would forbid token 1 - # after the trailing 0. A small window should not see that early bigram. - input_ids = torch.tensor([[0, 1, 2, 0]], device=torch_device, dtype=torch.long) - scores = torch.zeros((1, vocab_size), device=torch_device, dtype=torch.float) - - small_window = UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor(ngram_size=2, window_size=2) - full_window = UnlimitedOcrSlidingWindowNoRepeatNgramLogitsProcessor(ngram_size=2, window_size=4) - - self.assertListEqual(torch.isinf(small_window(input_ids, scores.clone())).tolist(), [[False, False, False]]) - self.assertListEqual(torch.isinf(full_window(input_ids, scores.clone())).tolist(), [[False, True, False]]) From 3e1e53746e11030bb2783e099adb1ceb99ddcb0c Mon Sep 17 00:00:00 2001 From: guarin Date: Wed, 29 Jul 2026 17:46:05 +0100 Subject: [PATCH 166/195] don't filter text config kwargs Co-Authored-By: Claude Opus 5 (1M context) --- .../models/unlimited_ocr/configuration_unlimited_ocr.py | 6 ++---- .../models/unlimited_ocr/modular_unlimited_ocr.py | 6 ++---- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index a2d44347eacd..b90c9df584bd 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -312,10 +312,8 @@ def __post_init__(self, **kwargs): self.vision_config = self.sub_configs["vision_config"](**self.vision_config) if self.text_config is None: - text_config_cls = self.sub_configs["text_config"] - text_keys = text_config_cls().to_dict().keys() - text_kwargs = {key: kwargs[key] for key in text_keys if key in kwargs} - self.text_config = text_config_cls(**text_kwargs) + # The reference implementation defines the text config values on the main config + self.text_config = self.sub_configs["text_config"](**kwargs) elif isinstance(self.text_config, dict): self.text_config = self.sub_configs["text_config"](**self.text_config) super().__post_init__(**kwargs) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 922e3082c0c8..a88d7756ed18 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -475,10 +475,8 @@ def __post_init__(self, **kwargs): self.vision_config = self.sub_configs["vision_config"](**self.vision_config) if self.text_config is None: - text_config_cls = self.sub_configs["text_config"] - text_keys = text_config_cls().to_dict().keys() - text_kwargs = {key: kwargs[key] for key in text_keys if key in kwargs} - self.text_config = text_config_cls(**text_kwargs) + # The reference implementation defines the text config values on the main config + self.text_config = self.sub_configs["text_config"](**kwargs) elif isinstance(self.text_config, dict): self.text_config = self.sub_configs["text_config"](**self.text_config) PreTrainedConfig.__post_init__(self, **kwargs) From 94e7e406ae4283374a7a516ec02df004f15bac56 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:52:00 +0100 Subject: [PATCH 167/195] remove use_sliding_window --- .../models/unlimited_ocr/configuration_unlimited_ocr.py | 6 +----- .../models/unlimited_ocr/modular_unlimited_ocr.py | 6 +----- tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py | 5 ++--- 3 files changed, 4 insertions(+), 13 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index b90c9df584bd..b266591fc02d 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -166,11 +166,9 @@ class UnlimitedOcrTextConfig(PreTrainedConfig): MLP type (`"dense"` or `"sparse"`) for each decoder layer, e.g. `["dense", "sparse", "sparse", ...]`. layer_types (`list[str]`, *optional*): Attention type for each decoder layer. Defaults to `"reference_sliding_attention"` on every layer. - use_sliding_window (`bool`, *optional*, defaults to `True`): - Whether to use reference sliding window attention. sliding_window (`int`, *optional*, defaults to `128`): Sliding window size for reference sliding window attention. If set, every token attends to the last - `sliding_window` and all image and prompt tokens. + `sliding_window` and all image and prompt tokens. Set to `None` to use full attention on every layer. Example: @@ -247,11 +245,9 @@ class UnlimitedOcrTextConfig(PreTrainedConfig): base_config_key = "text_config" mlp_layer_types: list[str] | None = None layer_types: list[str] | None = None - use_sliding_window: bool = True sliding_window: int | None = 128 def __post_init__(self, **kwargs): - self.sliding_window = self.sliding_window if self.use_sliding_window else None if self.layer_types is None: self.layer_types = [ "reference_sliding_attention" if self.sliding_window is not None else "full_attention" diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index a88d7756ed18..bb2107fffb2a 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -394,11 +394,9 @@ class UnlimitedOcrTextConfig(DeepseekOcr2TextConfig): MLP type (`"dense"` or `"sparse"`) for each decoder layer, e.g. `["dense", "sparse", "sparse", ...]`. layer_types (`list[str]`, *optional*): Attention type for each decoder layer. Defaults to `"reference_sliding_attention"` on every layer. - use_sliding_window (`bool`, *optional*, defaults to `True`): - Whether to use reference sliding window attention. sliding_window (`int`, *optional*, defaults to `128`): Sliding window size for reference sliding window attention. If set, every token attends to the last - `sliding_window` and all image and prompt tokens. + `sliding_window` and all image and prompt tokens. Set to `None` to use full attention on every layer. Example: @@ -426,11 +424,9 @@ class UnlimitedOcrTextConfig(DeepseekOcr2TextConfig): topk_group: int | None = 1 num_experts_per_tok: int | None = 6 layer_types: list[str] | None = None - use_sliding_window: bool = True sliding_window: int | None = 128 def __post_init__(self, **kwargs): - self.sliding_window = self.sliding_window if self.use_sliding_window else None if self.layer_types is None: self.layer_types = [ "reference_sliding_attention" if self.sliding_window is not None else "full_attention" diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 1a6b9cf7df9e..6f0574430af6 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -80,7 +80,6 @@ def __init__(self, parent, **kwargs): kwargs.setdefault("moe_intermediate_size", 64) kwargs.setdefault("num_experts_per_tok", 2) kwargs.setdefault("sliding_window", 4) - kwargs.setdefault("use_sliding_window", True) super().__init__(parent, **kwargs) self.sam_config = { @@ -271,8 +270,8 @@ def test_generate_cache_chunked_prefill(self): self.assertEqual(layer.prefill_length, prompt_length) def test_generate_without_sliding_window(self): - """With `use_sliding_window=False` every layer is a full attention layer.""" - model_tester = self.model_tester_class(self, use_sliding_window=False) + """With `sliding_window=None` every layer is a full attention layer.""" + model_tester = self.model_tester_class(self, sliding_window=None) config, inputs_dict = model_tester.prepare_config_and_inputs_for_common() self.assertEqual(config.text_config.layer_types, ["full_attention"] * config.text_config.num_hidden_layers) From 417946d368894f4330800b48ed0ef2368f17b6dd Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:31:32 +0100 Subject: [PATCH 168/195] handle patches grid --- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 6 +++++- .../models/unlimited_ocr/modular_unlimited_ocr.py | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 081e58cb6f71..2d279486cb24 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1842,7 +1842,7 @@ def get_image_features( num_local_patches (`list[int]` or `torch.Tensor`, *optional*): Number of local patches per image, e.g. `[6, 0, 4]`. patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): - The patches grid `(num_columns, num_rows)` per image. + The patches grid `(num_columns, num_rows)` per image. Required if `pixel_values_local` is passed. """ if isinstance(num_local_patches, torch.Tensor): num_local_patches = num_local_patches.tolist() @@ -2000,6 +2000,7 @@ def get_image_features( pixel_values: torch.FloatTensor, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, + patches_grid: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: r""" @@ -2009,11 +2010,14 @@ def get_image_features( 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]`. + patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + The patches grid `(num_columns, num_rows)` per image. Required if `pixel_values_local` is passed. """ return self.model.get_image_features( pixel_values=pixel_values, pixel_values_local=pixel_values_local, num_local_patches=num_local_patches, + patches_grid=patches_grid, **kwargs, ) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index bb2107fffb2a..e9d023835edc 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -1121,7 +1121,7 @@ def get_image_features( num_local_patches (`list[int]` or `torch.Tensor`, *optional*): Number of local patches per image, e.g. `[6, 0, 4]`. patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): - The patches grid `(num_columns, num_rows)` per image. + The patches grid `(num_columns, num_rows)` per image. Required if `pixel_values_local` is passed. """ if isinstance(num_local_patches, torch.Tensor): num_local_patches = num_local_patches.tolist() @@ -1254,6 +1254,7 @@ def get_image_features( pixel_values: torch.FloatTensor, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, + patches_grid: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: r""" @@ -1263,11 +1264,14 @@ def get_image_features( 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]`. + patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + The patches grid `(num_columns, num_rows)` per image. Required if `pixel_values_local` is passed. """ return self.model.get_image_features( pixel_values=pixel_values, pixel_values_local=pixel_values_local, num_local_patches=num_local_patches, + patches_grid=patches_grid, **kwargs, ) From 6b5f625deeb0d05b9d40011401920d29afee561e Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:50:57 +0000 Subject: [PATCH 169/195] update expectations --- .../unlimited_ocr/test_modeling_unlimited_ocr.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 1a6b9cf7df9e..568cca54de52 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -328,7 +328,7 @@ def test_small_model_integration_test_document_parsing(self): decoded = self.processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True) EXPECTED_DECODED_TEXT = Expectations( { - ("cuda", None): "image [383, 87, 497, 171]\ntitle [333", + ("cuda", 9): "image [383, 88, 497, 175]\ntitle [333", ("cpu", None): "image [383, 87, 497, 171]\ntitle [333", } ).get_expectation() # fmt: skip @@ -360,7 +360,7 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se ) EXPECTED_DECODED_TEXT = Expectations( { - ("cuda", None): "<|det|>image [383, 87, 497, 171]<|/det|>\n<|det|>title [333", + ("cuda", 9): "<|det|>image [383, 88, 497, 175]<|/det|>\n<|det|>title [333", ("cpu", None): "<|det|>image [383, 87, 497, 171]<|/det|>\n<|det|>title [333", } ).get_expectation() # fmt: skip @@ -368,7 +368,7 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se EXPECTED_DETECTIONS = Expectations( { - ("cuda", None): [{"region_type": "image", "box": [383, 87, 497, 171], "text": "\n"}], + ("cuda", 9): [{"region_type": "image", "box": [383, 88, 497, 175], "text": "\n"}], ("cpu", None): [{"region_type": "image", "box": [383, 87, 497, 171], "text": "\n"}], } ).get_expectation() # fmt: skip @@ -411,8 +411,8 @@ def test_small_model_integration_test_document_parsing_batched(self): EXPECTED_DECODED_TEXT = Expectations( { - ("cuda", None): [ - "image [383, 88, 497, 171]\ntitle [333", + ("cuda", 9): [ + "image [383, 88, 497, 174]\ntitle [333", "header [53, 23, 365, 41]Advanced Template and Styl", ], ("cpu", None): [ @@ -459,7 +459,7 @@ def test_small_model_integration_test_multi_page_document_parsing(self): EXPECTED_DECODED_TEXT = Expectations( { - ("cuda", None): "image [382, 87, 489, 174]\n", + ("cuda", 9): "image [382, 87, 489, 180]\n", ("cpu", None): "image [382, 87, 489, 174]\n", } ).get_expectation() From df084607ddc268e23ab995cc176dabab604b6ac9 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Wed, 29 Jul 2026 19:07:57 +0100 Subject: [PATCH 170/195] remove test_all_params_have_gradient --- tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 061164cb27bb..69064c5e075d 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -125,7 +125,6 @@ def get_config(self): @require_torch class UnlimitedOcrModelTest(VLMModelTest, unittest.TestCase): model_tester_class = UnlimitedOcrVisionText2TextModelTester - test_all_params_have_gradient = False @unittest.skip( reason="UnlimitedOcrVisionModel builds a hybrid bidirectional+causal mask internally, so SDPA is always called with a non-null `attn_mask`." From f835baf28e43308cd8aea72cee0c2a9903b1231a Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:02:19 +0100 Subject: [PATCH 171/195] refactor dynamic cache --- .../unlimited_ocr/modeling_unlimited_ocr.py | 119 ++++++++---------- .../unlimited_ocr/modular_unlimited_ocr.py | 118 ++++++++--------- 2 files changed, 103 insertions(+), 134 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 2d279486cb24..cdd9b720bcc5 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -29,7 +29,7 @@ from ... import initialization as init from ...activations import ACT2FN -from ...cache_utils import Cache, DynamicCache, DynamicLayer, DynamicSlidingWindowLayer, StaticSlidingWindowLayer +from ...cache_utils import Cache, DynamicCache, DynamicSlidingWindowLayer, StaticSlidingWindowLayer from ...configuration_utils import PreTrainedConfig from ...integrations import ( use_experts_implementation, @@ -1305,6 +1305,9 @@ class UnlimitedOcrDynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): def __init__(self, sliding_window: int | None = None, **kwargs): super().__init__(sliding_window=sliding_window, **kwargs) self.prefill_length: int | None = None + self.prefill_cumulative_length: int = 0 + self.prefill_keys: torch.Tensor | None = None + self.prefill_values: torch.Tensor | None = None def update( self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs @@ -1313,64 +1316,47 @@ def update( if not self.is_initialized: self.lazy_initialization(key_states, value_states) - kv_length = key_states.shape[-2] - self.cumulative_length += kv_length - - # Compute the full states - full_key_states = torch.cat([self.keys, key_states], dim=-2) - full_value_states = torch.cat([self.values, value_states], dim=-2) + # Prefill + if self.prefill_length is None or self.prefill_cumulative_length < self.prefill_length: + kv_length = key_states.shape[-2] + self.prefill_cumulative_length += kv_length - # Cache growing - if ( - self.record_past - or self.prefill_length is None - or self.cumulative_length <= self.prefill_length + self.sliding_window - 1 - ): - self.keys = full_key_states - self.values = full_value_states if self.prefill_length is None: - self.prefill_length = self.cumulative_length - # Cache full - elif self.keys.shape[-2] == self.prefill_length + self.sliding_window - 1: - self.keys[:, :, -self.sliding_window + 1 :].copy_(full_key_states[:, :, -self.sliding_window + 1 :, :]) - self.values[:, :, -self.sliding_window + 1 :].copy_(full_value_states[:, :, -self.sliding_window + 1 :, :]) - # Cache full after this update and full_key_states > cache size - else: - self.keys = torch.cat( - [ - full_key_states[:, :, : self.prefill_length, :], - full_key_states[:, :, -self.sliding_window + 1 :, :], - ], - dim=-2, - ) - self.values = torch.cat( - [ - full_value_states[:, :, : self.prefill_length, :], - full_value_states[:, :, -self.sliding_window + 1 :, :], - ], - dim=-2, - ) + self.prefill_length = self.prefill_cumulative_length + + self.prefill_keys = torch.cat([self.prefill_keys, key_states], dim=-2) + self.prefill_values = torch.cat([self.prefill_values, value_states], dim=-2) + + return self.prefill_keys, self.prefill_values - # Return full states to avoid losing context in case we added multiple tokens at once + sliding_key_states, sliding_value_states = super().update(key_states=key_states, value_states=value_states) + full_key_states = torch.cat([self.prefill_keys, sliding_key_states], dim=-2) + full_value_states = torch.cat([self.prefill_values, sliding_value_states], dim=-2) return full_key_states, full_value_states + def lazy_initialization(self, key_states, value_states): + super().lazy_initialization(key_states, value_states) + self.prefill_keys = self.keys.clone() + self.prefill_values = self.values.clone() + def get_mask_sizes(self, query_length: int) -> tuple[int, int]: """Return the length and offset of the cache, used to generate the attention mask""" - is_full = ( - self.prefill_length is not None and self.cumulative_length >= self.prefill_length + self.sliding_window - ) - - kv_offset = 0 - if is_full: - kv_offset = max(self.cumulative_length - self.prefill_length - self.sliding_window + 1, 0) - kv_length = self.prefill_length + self.sliding_window - 1 + query_length + # Prefill + if self.prefill_length is None or self.prefill_cumulative_length < self.prefill_length: + kv_offset = 0 + kv_length = self.prefill_cumulative_length + query_length else: - kv_length = self.cumulative_length + query_length + kv_length, kv_offset = super().get_mask_sizes(query_length) + kv_length += self.prefill_length + kv_offset += self.prefill_length # Returned kv_offset is with respect to sliding window keys. # Remove kv_offset from kv_idx to retrieve the prefill indices. return kv_length, kv_offset + def get_seq_length(self): + return self.prefill_cumulative_length + super().get_seq_length() + def get_max_length(self) -> int: """Return the maximum cache shape of the cache""" if self.prefill_length is None: @@ -1384,34 +1370,33 @@ def set_prefill_length(self, prefill_length: int) -> None: def reset(self) -> None: super().reset() + if self.is_initialized: + self.prefill_keys.zero_() + self.prefill_values.zero_() self.prefill_length = None + self.prefill_cumulative_length = 0 def crop(self, max_length: int) -> None: """ Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be negative to remove `max_length` tokens. """ - # If we are beyond the sliding window, we need to be more careful - if self.prefill_length is not None and self.get_seq_length() >= self.prefill_length + self.sliding_window: - if not self.record_past: - raise RuntimeError( - "`crop` was called, but the current layer does not track past states, and the sliding window size was already " - "reached. Call `activate_past_recording` before `crop` to be able to rollback the cache." - ) - if max_length > 0: - raise RuntimeError( - "Once the sliding window size has been reached, `UnlimitedOcrDynamicReferenceSlidingWindowLayer` can only " - "be cropped by passing a negative int, to specify how many tokens to remove" - ) - tokens_to_remove = abs(max_length) - # We crop, and restrict the size back to the sliding window if still larger - self.keys = self.keys[:, :, -self.sliding_window + 1 - tokens_to_remove : -tokens_to_remove, :] - self.values = self.values[:, :, -self.sliding_window + 1 - tokens_to_remove : -tokens_to_remove, :] - self.cumulative_length = self.cumulative_length - tokens_to_remove - # If we did not reach the sliding window, we can do the same as for a full attention layer - else: - DynamicLayer.crop(self, max_length) - self.cumulative_length = self.keys.shape[-2] + # Sliding window + if self.prefill_length is not None and self.get_seq_length() >= self.prefill_length: + sliding_max_length = max(max_length, -self.cumulative_length) + max_length = abs(max_length - sliding_max_length) + super().crop(sliding_max_length) + + # Prefill + if max_length <= 0: + max_length = self.prefill_cumulative_length - abs(max_length) + + if self.prefill_cumulative_length <= max_length: + return + + self.prefill_keys = self.prefill_keys[..., :max_length, :] + self.prefill_values = self.prefill_values[..., :max_length, :] + self.prefill_cumulative_length = max_length class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index e9d023835edc..c8bb12857a99 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -23,7 +23,6 @@ from ...cache_utils import ( Cache, DynamicCache, - DynamicLayer, DynamicSlidingWindowLayer, StaticSlidingWindowLayer, ) @@ -688,6 +687,9 @@ class UnlimitedOcrDynamicReferenceSlidingWindowLayer(DynamicSlidingWindowLayer): def __init__(self, sliding_window: int | None = None, **kwargs): super().__init__(sliding_window=sliding_window, **kwargs) self.prefill_length: int | None = None + self.prefill_cumulative_length: int = 0 + self.prefill_keys: torch.Tensor | None = None + self.prefill_values: torch.Tensor | None = None def update( self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs @@ -696,64 +698,47 @@ def update( if not self.is_initialized: self.lazy_initialization(key_states, value_states) - kv_length = key_states.shape[-2] - self.cumulative_length += kv_length - - # Compute the full states - full_key_states = torch.cat([self.keys, key_states], dim=-2) - full_value_states = torch.cat([self.values, value_states], dim=-2) + # Prefill + if self.prefill_length is None or self.prefill_cumulative_length < self.prefill_length: + kv_length = key_states.shape[-2] + self.prefill_cumulative_length += kv_length - # Cache growing - if ( - self.record_past - or self.prefill_length is None - or self.cumulative_length <= self.prefill_length + self.sliding_window - 1 - ): - self.keys = full_key_states - self.values = full_value_states if self.prefill_length is None: - self.prefill_length = self.cumulative_length - # Cache full - elif self.keys.shape[-2] == self.prefill_length + self.sliding_window - 1: - self.keys[:, :, -self.sliding_window + 1 :].copy_(full_key_states[:, :, -self.sliding_window + 1 :, :]) - self.values[:, :, -self.sliding_window + 1 :].copy_(full_value_states[:, :, -self.sliding_window + 1 :, :]) - # Cache full after this update and full_key_states > cache size - else: - self.keys = torch.cat( - [ - full_key_states[:, :, : self.prefill_length, :], - full_key_states[:, :, -self.sliding_window + 1 :, :], - ], - dim=-2, - ) - self.values = torch.cat( - [ - full_value_states[:, :, : self.prefill_length, :], - full_value_states[:, :, -self.sliding_window + 1 :, :], - ], - dim=-2, - ) + self.prefill_length = self.prefill_cumulative_length + + self.prefill_keys = torch.cat([self.prefill_keys, key_states], dim=-2) + self.prefill_values = torch.cat([self.prefill_values, value_states], dim=-2) + + return self.prefill_keys, self.prefill_values - # Return full states to avoid losing context in case we added multiple tokens at once + sliding_key_states, sliding_value_states = super().update(key_states=key_states, value_states=value_states) + full_key_states = torch.cat([self.prefill_keys, sliding_key_states], dim=-2) + full_value_states = torch.cat([self.prefill_values, sliding_value_states], dim=-2) return full_key_states, full_value_states + def lazy_initialization(self, key_states, value_states): + super().lazy_initialization(key_states, value_states) + self.prefill_keys = self.keys.clone() + self.prefill_values = self.values.clone() + def get_mask_sizes(self, query_length: int) -> tuple[int, int]: """Return the length and offset of the cache, used to generate the attention mask""" - is_full = ( - self.prefill_length is not None and self.cumulative_length >= self.prefill_length + self.sliding_window - ) - - kv_offset = 0 - if is_full: - kv_offset = max(self.cumulative_length - self.prefill_length - self.sliding_window + 1, 0) - kv_length = self.prefill_length + self.sliding_window - 1 + query_length + # Prefill + if self.prefill_length is None or self.prefill_cumulative_length < self.prefill_length: + kv_offset = 0 + kv_length = self.prefill_cumulative_length + query_length else: - kv_length = self.cumulative_length + query_length + kv_length, kv_offset = super().get_mask_sizes(query_length) + kv_length += self.prefill_length + kv_offset += self.prefill_length # Returned kv_offset is with respect to sliding window keys. # Remove kv_offset from kv_idx to retrieve the prefill indices. return kv_length, kv_offset + def get_seq_length(self): + return self.prefill_cumulative_length + super().get_seq_length() + def get_max_length(self) -> int: """Return the maximum cache shape of the cache""" if self.prefill_length is None: @@ -767,34 +752,33 @@ def set_prefill_length(self, prefill_length: int) -> None: def reset(self) -> None: super().reset() + if self.is_initialized: + self.prefill_keys.zero_() + self.prefill_values.zero_() self.prefill_length = None + self.prefill_cumulative_length = 0 def crop(self, max_length: int) -> None: """ Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be negative to remove `max_length` tokens. """ - # If we are beyond the sliding window, we need to be more careful - if self.prefill_length is not None and self.get_seq_length() >= self.prefill_length + self.sliding_window: - if not self.record_past: - raise RuntimeError( - "`crop` was called, but the current layer does not track past states, and the sliding window size was already " - "reached. Call `activate_past_recording` before `crop` to be able to rollback the cache." - ) - if max_length > 0: - raise RuntimeError( - "Once the sliding window size has been reached, `UnlimitedOcrDynamicReferenceSlidingWindowLayer` can only " - "be cropped by passing a negative int, to specify how many tokens to remove" - ) - tokens_to_remove = abs(max_length) - # We crop, and restrict the size back to the sliding window if still larger - self.keys = self.keys[:, :, -self.sliding_window + 1 - tokens_to_remove : -tokens_to_remove, :] - self.values = self.values[:, :, -self.sliding_window + 1 - tokens_to_remove : -tokens_to_remove, :] - self.cumulative_length = self.cumulative_length - tokens_to_remove - # If we did not reach the sliding window, we can do the same as for a full attention layer - else: - DynamicLayer.crop(self, max_length) - self.cumulative_length = self.keys.shape[-2] + # Sliding window + if self.prefill_length is not None and self.get_seq_length() >= self.prefill_length: + sliding_max_length = max(max_length, -self.cumulative_length) + max_length = abs(max_length - sliding_max_length) + super().crop(sliding_max_length) + + # Prefill + if max_length <= 0: + max_length = self.prefill_cumulative_length - abs(max_length) + + if self.prefill_cumulative_length <= max_length: + return + + self.prefill_keys = self.prefill_keys[..., :max_length, :] + self.prefill_values = self.prefill_values[..., :max_length, :] + self.prefill_cumulative_length = max_length class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): From 1657694ca3e7d4dd7ad1a5f70216fe6fd1da4979 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:12:20 +0100 Subject: [PATCH 172/195] update check --- src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py | 2 +- src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index cdd9b720bcc5..856d60e4ffff 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1382,7 +1382,7 @@ def crop(self, max_length: int) -> None: negative to remove `max_length` tokens. """ # Sliding window - if self.prefill_length is not None and self.get_seq_length() >= self.prefill_length: + if self.cumulative_length > 0: sliding_max_length = max(max_length, -self.cumulative_length) max_length = abs(max_length - sliding_max_length) super().crop(sliding_max_length) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index c8bb12857a99..9a65cc3fd658 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -764,7 +764,7 @@ def crop(self, max_length: int) -> None: negative to remove `max_length` tokens. """ # Sliding window - if self.prefill_length is not None and self.get_seq_length() >= self.prefill_length: + if self.cumulative_length > 0: sliding_max_length = max(max_length, -self.cumulative_length) max_length = abs(max_length - sliding_max_length) super().crop(sliding_max_length) From d7a8a182b5c223a377e224e32d2f70ad4340e199 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:24:59 +0100 Subject: [PATCH 173/195] fix offset --- src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py | 1 - src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 1 - 2 files changed, 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 856d60e4ffff..f06ac3103de2 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1348,7 +1348,6 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: else: kv_length, kv_offset = super().get_mask_sizes(query_length) kv_length += self.prefill_length - kv_offset += self.prefill_length # Returned kv_offset is with respect to sliding window keys. # Remove kv_offset from kv_idx to retrieve the prefill indices. diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 9a65cc3fd658..bf3d4ae88c31 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -730,7 +730,6 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: else: kv_length, kv_offset = super().get_mask_sizes(query_length) kv_length += self.prefill_length - kv_offset += self.prefill_length # Returned kv_offset is with respect to sliding window keys. # Remove kv_offset from kv_idx to retrieve the prefill indices. From 0eefd0c4dec38e031c4af99c3f825788c26fa23f Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:44:29 +0100 Subject: [PATCH 174/195] add reorder_cache --- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 11 +++++++++++ .../models/unlimited_ocr/modular_unlimited_ocr.py | 11 +++++++++++ 2 files changed, 22 insertions(+) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index f06ac3103de2..61803be9d77b 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1375,6 +1375,17 @@ def reset(self) -> None: self.prefill_length = None self.prefill_cumulative_length = 0 + def reorder_cache(self, beam_idx: torch.LongTensor) -> None: + """Reorders this layer's cache for beam search.""" + # Sliding window + if self.cumulative_length > 0: + super().reorder_cache(beam_idx) + + # Prefill + if self.prefill_cumulative_length > 0: + self.prefill_keys = self.prefill_keys.index_select(0, beam_idx.to(self.prefill_keys.device)) + self.prefill_values = self.prefill_values.index_select(0, beam_idx.to(self.prefill_values.device)) + def crop(self, max_length: int) -> None: """ Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index bf3d4ae88c31..8c4275f91924 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -757,6 +757,17 @@ def reset(self) -> None: self.prefill_length = None self.prefill_cumulative_length = 0 + def reorder_cache(self, beam_idx: torch.LongTensor) -> None: + """Reorders this layer's cache for beam search.""" + # Sliding window + if self.cumulative_length > 0: + super().reorder_cache(beam_idx) + + # Prefill + if self.prefill_cumulative_length > 0: + self.prefill_keys = self.prefill_keys.index_select(0, beam_idx.to(self.prefill_keys.device)) + self.prefill_values = self.prefill_values.index_select(0, beam_idx.to(self.prefill_values.device)) + def crop(self, max_length: int) -> None: """ Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be From 1eeb9ec94bc101136f1b8f193385539d42f6ddf8 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 30 Jul 2026 11:05:40 +0100 Subject: [PATCH 175/195] handle overlapping prefill/window overlap --- .../unlimited_ocr/modeling_unlimited_ocr.py | 16 +++++++++++----- .../unlimited_ocr/modular_unlimited_ocr.py | 16 +++++++++++----- 2 files changed, 22 insertions(+), 10 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 61803be9d77b..e0d807688f50 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1319,15 +1319,21 @@ def update( # Prefill if self.prefill_length is None or self.prefill_cumulative_length < self.prefill_length: kv_length = key_states.shape[-2] - self.prefill_cumulative_length += kv_length if self.prefill_length is None: - self.prefill_length = self.prefill_cumulative_length + self.prefill_length = self.prefill_cumulative_length + kv_length - self.prefill_keys = torch.cat([self.prefill_keys, key_states], dim=-2) - self.prefill_values = torch.cat([self.prefill_values, value_states], dim=-2) + prefill_kv_length = min(kv_length, self.prefill_length - self.prefill_cumulative_length) + self.prefill_cumulative_length += prefill_kv_length - return self.prefill_keys, self.prefill_values + self.prefill_keys = torch.cat([self.prefill_keys, key_states[..., :prefill_kv_length, :]], dim=-2) + self.prefill_values = torch.cat([self.prefill_values, value_states[..., :prefill_kv_length, :]], dim=-2) + + if prefill_kv_length == kv_length: + return self.prefill_keys, self.prefill_values + + key_states = key_states[..., prefill_kv_length:, :] + value_states = value_states[..., prefill_kv_length:, :] sliding_key_states, sliding_value_states = super().update(key_states=key_states, value_states=value_states) full_key_states = torch.cat([self.prefill_keys, sliding_key_states], dim=-2) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 8c4275f91924..1c3159da9598 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -701,15 +701,21 @@ def update( # Prefill if self.prefill_length is None or self.prefill_cumulative_length < self.prefill_length: kv_length = key_states.shape[-2] - self.prefill_cumulative_length += kv_length if self.prefill_length is None: - self.prefill_length = self.prefill_cumulative_length + self.prefill_length = self.prefill_cumulative_length + kv_length - self.prefill_keys = torch.cat([self.prefill_keys, key_states], dim=-2) - self.prefill_values = torch.cat([self.prefill_values, value_states], dim=-2) + prefill_kv_length = min(kv_length, self.prefill_length - self.prefill_cumulative_length) + self.prefill_cumulative_length += prefill_kv_length - return self.prefill_keys, self.prefill_values + self.prefill_keys = torch.cat([self.prefill_keys, key_states[..., :prefill_kv_length, :]], dim=-2) + self.prefill_values = torch.cat([self.prefill_values, value_states[..., :prefill_kv_length, :]], dim=-2) + + if prefill_kv_length == kv_length: + return self.prefill_keys, self.prefill_values + + key_states = key_states[..., prefill_kv_length:, :] + value_states = value_states[..., prefill_kv_length:, :] sliding_key_states, sliding_value_states = super().update(key_states=key_states, value_states=value_states) full_key_states = torch.cat([self.prefill_keys, sliding_key_states], dim=-2) From 774d4894ab8b1fddfaaa8f3f9a867cde0609d596 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:21:47 +0000 Subject: [PATCH 176/195] fix crop --- .../unlimited_ocr/modeling_unlimited_ocr.py | 27 ++++++++++--------- .../unlimited_ocr/modular_unlimited_ocr.py | 27 ++++++++++--------- 2 files changed, 30 insertions(+), 24 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 61803be9d77b..3764dc9e71bb 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1391,22 +1391,25 @@ def crop(self, max_length: int) -> None: Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be negative to remove `max_length` tokens. """ - # Sliding window - if self.cumulative_length > 0: - sliding_max_length = max(max_length, -self.cumulative_length) - max_length = abs(max_length - sliding_max_length) - super().crop(sliding_max_length) - - # Prefill + total_length = self.get_seq_length() if max_length <= 0: - max_length = self.prefill_cumulative_length - abs(max_length) + max_length = total_length - abs(max_length) - if self.prefill_cumulative_length <= max_length: + if total_length <= max_length: return - self.prefill_keys = self.prefill_keys[..., :max_length, :] - self.prefill_values = self.prefill_values[..., :max_length, :] - self.prefill_cumulative_length = max_length + # Prefill tokens are only evicted once the target length drops below how many are currently cached; + # whatever remains of the budget after keeping them is what the sliding window buffer gets to keep. + new_prefill_length = min(self.prefill_cumulative_length, max_length) + new_sliding_length = max_length - new_prefill_length + + if self.cumulative_length > new_sliding_length: + super().crop(-(self.cumulative_length - new_sliding_length)) + + if new_prefill_length < self.prefill_cumulative_length: + self.prefill_keys = self.prefill_keys[..., :new_prefill_length, :] + self.prefill_values = self.prefill_values[..., :new_prefill_length, :] + self.prefill_cumulative_length = new_prefill_length class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 8c4275f91924..1e6aa2b2a654 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -773,22 +773,25 @@ def crop(self, max_length: int) -> None: Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be negative to remove `max_length` tokens. """ - # Sliding window - if self.cumulative_length > 0: - sliding_max_length = max(max_length, -self.cumulative_length) - max_length = abs(max_length - sliding_max_length) - super().crop(sliding_max_length) - - # Prefill + total_length = self.get_seq_length() if max_length <= 0: - max_length = self.prefill_cumulative_length - abs(max_length) + max_length = total_length - abs(max_length) - if self.prefill_cumulative_length <= max_length: + if total_length <= max_length: return - self.prefill_keys = self.prefill_keys[..., :max_length, :] - self.prefill_values = self.prefill_values[..., :max_length, :] - self.prefill_cumulative_length = max_length + # Prefill tokens are only evicted once the target length drops below how many are currently cached; + # whatever remains of the budget after keeping them is what the sliding window buffer gets to keep. + new_prefill_length = min(self.prefill_cumulative_length, max_length) + new_sliding_length = max_length - new_prefill_length + + if self.cumulative_length > new_sliding_length: + super().crop(-(self.cumulative_length - new_sliding_length)) + + if new_prefill_length < self.prefill_cumulative_length: + self.prefill_keys = self.prefill_keys[..., :new_prefill_length, :] + self.prefill_values = self.prefill_values[..., :new_prefill_length, :] + self.prefill_cumulative_length = new_prefill_length class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): From 0ad2158c534a5a978c007afdba64d09259306c31 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:56:25 +0000 Subject: [PATCH 177/195] update tests --- .../test_modeling_unlimited_ocr.py | 26 ++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 69064c5e075d..3f161b14e455 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -145,7 +145,31 @@ def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_l # max cache length can be smaller than sequence length max_length = past_key_values.get_max_length() seq_length = min(seq_length, max_length) if max_length >= 0 else seq_length - super()._check_past_key_values_for_generate(batch_size, past_key_values, seq_length, config) + + if not any(hasattr(layer, "prefill_keys") for layer in past_key_values.layers): + super()._check_past_key_values_for_generate(batch_size, past_key_values, seq_length, config) + return + + # Super method assumes that there is only layer.keys/values but reference sliding window layers + # can have layer.prefill_keys/prefill_values as well which results in wrong assertions in the + # super method. + def seq_len(tensor): + return tensor.shape[-2] if tensor.dim() > 1 else 0 + + config = config.get_text_config(decoder=True) + self.assertEqual(config.num_hidden_layers, len(past_key_values)) + for layer_idx, layer in enumerate(past_key_values.layers): + layer_config = config.per_layer_config[layer_idx] + attention_shape = self._get_attention_shape(batch_size, seq_length, layer_config) + combined_keys_length = seq_len(layer.prefill_keys) + seq_len(layer.keys) + combined_values_length = seq_len(layer.prefill_values) + seq_len(layer.values) + self.assertEqual( + (*layer.prefill_keys.shape[:-2], combined_keys_length, layer.prefill_keys.shape[-1]), attention_shape + ) + self.assertEqual( + (*layer.prefill_values.shape[:-2], combined_values_length, layer.prefill_values.shape[-1]), + attention_shape, + ) def _check_generate_cache_sliding_window_too_small(self, cache_implementation: str, prefill_max_new_tokens: int): """Test that reference sliding window cache works correctly when decoding more than sliding_window tokens at once.""" From 1cf342aa8f41d302a0a27d5114140c5c35f81159 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Thu, 30 Jul 2026 11:59:41 +0100 Subject: [PATCH 178/195] fix-repo --- tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 3f161b14e455..ef15eed7c75d 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -150,9 +150,9 @@ def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_l super()._check_past_key_values_for_generate(batch_size, past_key_values, seq_length, config) return - # Super method assumes that there is only layer.keys/values but reference sliding window layers - # can have layer.prefill_keys/prefill_values as well which results in wrong assertions in the - # super method. + # Super method assumes that there is only layer.keys/values but reference sliding window layers + # can have layer.prefill_keys/prefill_values as well which results in wrong assertions in the + # super method. def seq_len(tensor): return tensor.shape[-2] if tensor.dim() > 1 else 0 From 76076d7832af55c0c93fc8a00be611a5601db156 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 31 Jul 2026 08:32:18 +0100 Subject: [PATCH 179/195] cleanup --- docs/source/en/model_doc/unlimited_ocr.md | 2 +- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 4 ++-- .../models/unlimited_ocr/modular_unlimited_ocr.py | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 7a6e7b7606ac..290aa0cc91a7 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -16,7 +16,7 @@ limitations under the License. ⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer. --> -*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-07-29.* +*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-07-31.* # UnlimitedOcr diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index e21798188e31..08bfe9206b76 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1342,8 +1342,8 @@ def update( def lazy_initialization(self, key_states, value_states): super().lazy_initialization(key_states, value_states) - self.prefill_keys = self.keys.clone() - self.prefill_values = self.values.clone() + self.prefill_keys = torch.empty_like(self.keys) + self.prefill_values = torch.empty_like(self.values) def get_mask_sizes(self, query_length: int) -> tuple[int, int]: """Return the length and offset of the cache, used to generate the attention mask""" diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index eb9dd170e81c..54c4128d86a7 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -724,8 +724,8 @@ def update( def lazy_initialization(self, key_states, value_states): super().lazy_initialization(key_states, value_states) - self.prefill_keys = self.keys.clone() - self.prefill_values = self.values.clone() + self.prefill_keys = torch.empty_like(self.keys) + self.prefill_values = torch.empty_like(self.values) def get_mask_sizes(self, query_length: int) -> tuple[int, int]: """Return the length and offset of the cache, used to generate the attention mask""" From 5e2134445853a1a74acf7cdbbdbec91410152089 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 31 Jul 2026 12:52:20 +0100 Subject: [PATCH 180/195] match static cache impl with dynamic one --- .../unlimited_ocr/generation_unlimited_ocr.py | 1 + .../unlimited_ocr/modeling_unlimited_ocr.py | 182 +++++++++--------- .../unlimited_ocr/modular_unlimited_ocr.py | 182 +++++++++--------- 3 files changed, 193 insertions(+), 172 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py index 9b4ef5c8849d..67f095fc2e7d 100644 --- a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py @@ -70,6 +70,7 @@ def _prefill( for layer in past_key_values.layers: if layer._layer_type == "reference_sliding_attention": layer.set_prefill_length(prefill_length) + return super()._prefill( input_ids, generation_config, diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 08bfe9206b76..9622af93ba74 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1351,6 +1351,7 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: if self.prefill_length is None or self.prefill_cumulative_length < self.prefill_length: kv_offset = 0 kv_length = self.prefill_cumulative_length + query_length + # Decode else: kv_length, kv_offset = super().get_mask_sizes(query_length) kv_length += self.prefill_length @@ -1359,7 +1360,7 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: # Remove kv_offset from kv_idx to retrieve the prefill indices. return kv_length, kv_offset - def get_seq_length(self): + def get_seq_length(self) -> int: return self.prefill_cumulative_length + super().get_seq_length() def get_max_length(self) -> int: @@ -1421,7 +1422,7 @@ def crop(self, max_length: int) -> None: class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): """ A static cache layer that stores the key and value states as static tensors of shape - `[batch_size, num_heads, min(max_cache_len, prefill_size + sliding_window), head_dim]`. + `[batch_size, num_heads, min(max_cache_len, prefill_length + sliding_window), head_dim]`. It lazily allocates its full backing tensors, and then mutates them in-place. Built for `torch.compile` support. @@ -1444,11 +1445,20 @@ class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): def __init__(self, max_cache_len: int, sliding_window: int, **kwargs): super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window, **kwargs) - # Keep `max_cache_len` as max value for length bookkeeping. - # The physical buffer doesn't exceed `prefill_length + sliding_window`. - self.max_cache_len = max_cache_len - self.sliding_window = sliding_window + # Effective sliding window length is determined by max_cache_len in parent + self.sliding_window = self.max_cache_len + self.prefill_max_cache_len = max_cache_len - self.sliding_window self.prefill_length: int | None = None + self.prefill_cumulative_length: int = 0 + self.prefill_keys: torch.Tensor | None = None + self.prefill_values: torch.Tensor | None = None + + # all_keys and all_values are contiguous buffers that hold the full cache to keep dynamo addresses static. + # prefill_keys, prefill_values, keys, and values are views of the buffers: + # all_keys = [prefill_keys, keys] + # all_values = [prefill_values, values] + self.all_keys: torch.Tensor | None = None + self.all_values: torch.Tensor | None = None def update( self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs @@ -1467,106 +1477,98 @@ def update( if not self.is_initialized: self.lazy_initialization(key_states, value_states) - kv_length = key_states.shape[-2] - # Prefill - if self.prefill_length is None or self.cumulative_length_int < self.prefill_length: + if self.prefill_length is None or self.prefill_cumulative_length < self.prefill_length: + kv_length = key_states.shape[-2] + if self.prefill_length is None: - self.prefill_length = self.cumulative_length_int + kv_length + self.prefill_length = min(self.prefill_max_cache_len, self.prefill_cumulative_length + kv_length) # Resize buffer if necessary - required_length = min(self.max_cache_len, self.prefill_length + self.sliding_window) - if self.keys.shape[-2] < required_length: - self._allocate_key_value_buffers(required_length, copy_existing=True) + if self.prefill_keys.shape[-2] != self.prefill_length: + self._allocate_key_value_buffers(self.prefill_length, copy_existing=True) - # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs - # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten - cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length + prefill_kv_length = min(kv_length, self.prefill_length - self.prefill_cumulative_length) + prefill_key_states = key_states[..., :prefill_kv_length, :] + prefill_value_states = value_states[..., :prefill_kv_length, :] + + cache_position = torch.arange(prefill_kv_length, device=self.device) + self.prefill_cumulative_length try: - self.keys.index_copy_(2, cache_position, key_states) - self.values.index_copy_(2, cache_position, value_states) + self.prefill_keys.index_copy_(2, cache_position, prefill_key_states) + self.prefill_values.index_copy_(2, cache_position, prefill_value_states) except NotImplementedError: # Fallback for devices like MPS where index_copy_ might not be supported. - self.keys[:, :, cache_position] = key_states - self.values[:, :, cache_position] = value_states + self.prefill_keys[:, :, cache_position] = prefill_key_states + self.prefill_values[:, :, cache_position] = prefill_value_states - # Keep both the int (control flow) and the tensor (cudagraph-safe indexing) versions in sync. - self.cumulative_length_int += kv_length - self.cumulative_length.add_(kv_length) + self.prefill_cumulative_length += prefill_kv_length - # Very important to return the `self` tensors here, as they have the static dynamo address - return self.keys, self.values + if prefill_kv_length == kv_length: + return self.all_keys, self.all_values + + key_states = key_states[..., prefill_kv_length:, :] + value_states = value_states[..., prefill_kv_length:, :] # Decode - # Call `StaticSlidingWindowLayer.update` to reduce code duplication. This requires setting temporary - # attributes to match expectations of parent class. - keys, values = self.keys, self.values - window = slice(self.prefill_length, self.prefill_length + self.sliding_window) - sliding_keys_buffer, sliding_values_buffer = keys[:, :, window, :], values[:, :, window, :] - self.keys, self.values = sliding_keys_buffer, sliding_values_buffer - self.cumulative_length -= self.prefill_length - self.cumulative_length_int -= self.prefill_length - max_cache_len, self.max_cache_len = self.max_cache_len, self.sliding_window - try: - sliding_keys, sliding_values = super().update(key_states, value_states, *args, **kwargs) - finally: - self.keys, self.values = keys, values - self.cumulative_length += self.prefill_length - self.cumulative_length_int += self.prefill_length - self.max_cache_len = max_cache_len + sliding_keys, sliding_values = super().update(key_states, value_states, *args, **kwargs) # The parent returned the original buffers - if sliding_keys is sliding_keys_buffer: - return self.keys, self.values - # The parent returned concatenated states + if sliding_keys is self.keys: + return self.all_keys, self.all_values + + # The parent returned concatenated full_key_states/full_value_states return ( - torch.cat((self.keys[:, :, : self.prefill_length, :], sliding_keys), dim=-2), - torch.cat((self.values[:, :, : self.prefill_length, :], sliding_values), dim=-2), + torch.cat((self.prefill_keys, sliding_keys), dim=-2), + torch.cat((self.prefill_values, sliding_values), dim=-2), ) def get_mask_sizes(self, query_length: int) -> tuple[int, int]: """Return the length and offset of the cache, used to generate the attention mask""" - is_full = ( - self.prefill_length is not None and self.cumulative_length_int >= self.prefill_length + self.sliding_window - ) - - kv_offset = 0 - # Prefill: the buffer is not necessarily allocated yet, so its size cannot be used - if self.prefill_length is None or self.cumulative_length_int < self.prefill_length: + # Prefill + if self.prefill_length is None or self.prefill_cumulative_length < self.prefill_length: + kv_offset = 0 prefill_length = ( - self.cumulative_length_int + query_length if self.prefill_length is None else self.prefill_length + min(self.prefill_max_cache_len, self.prefill_cumulative_length + query_length) + if self.prefill_length is None + else self.prefill_length + ) + kv_length = max( + prefill_length + self.sliding_window, + # sliding window returns concatenated states + self.prefill_cumulative_length + self.cumulative_length_int + query_length, ) - kv_length = min(self.max_cache_len, prefill_length + self.sliding_window) - # Decode: cache is already full - elif is_full: - kv_offset = max(self.cumulative_length_int - self.prefill_length - self.sliding_window + 1, 0) - kv_length = self.prefill_length + self.sliding_window - 1 + query_length - # Decode: cache not yet full, but becoming full on this update - elif ( - self.prefill_length is not None - and self.cumulative_length_int + query_length > self.prefill_length + self.sliding_window - ): - kv_length = self.cumulative_length_int + query_length - # Decode: cache not yet full but we return the local size as it's static + # Decode else: - kv_length = self.keys.shape[-2] + kv_length, kv_offset = super().get_mask_sizes(query_length) + kv_length += self.prefill_length + # Returned kv_offset is with respect to sliding window keys. + # Remove kv_offset from kv_idx to retrieve the prefill indices. return kv_length, kv_offset def get_max_length(self) -> int: """Return the maximum cache shape of the cache""" + max_cache_len = self.prefill_max_cache_len + self.max_cache_len if self.prefill_length is None: - return self.max_cache_len - return min(self.max_cache_len, self.prefill_length + self.sliding_window) + return max_cache_len + return min(max_cache_len, self.prefill_length + self.sliding_window) + + def get_seq_length(self) -> int: + return self.prefill_cumulative_length + super().get_seq_length() def set_prefill_length(self, prefill_length: int) -> None: """Declare how many leading tokens are prefill states, before they are cached.""" if self.prefill_length is None: - self.prefill_length = prefill_length + self.prefill_length = min(prefill_length, self.prefill_max_cache_len) def reset(self) -> None: super().reset() + if self.is_initialized: + # This also zeros the self.prefill_keys, self.keys, etc. views + self.all_keys.zero_() + self.all_values.zero_() self.prefill_length = None + self.prefill_cumulative_length = 0 def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None: self.dtype, self.device = key_states.dtype, key_states.device @@ -1579,29 +1581,37 @@ def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tens torch._dynamo.mark_static_address(self.cumulative_length) prefill_length = key_states.shape[-2] if self.prefill_length is None else self.prefill_length - self._allocate_key_value_buffers(min(self.max_cache_len, prefill_length + self.sliding_window)) + self._allocate_key_value_buffers(min(self.prefill_max_cache_len, prefill_length)) self.is_initialized = True - def _allocate_key_value_buffers(self, physical_length: int, copy_existing: bool = False) -> None: - """(Re)allocate the static key/value buffers to `physical_length` slots and (re)tag the static address. - - Only ever called from the eager prefill pass (initial allocation or chunked-prefill growth), never from a - compiled decode step, so reallocating here is safe for cudagraphs. - """ - new_keys = torch.zeros( - (self.batch_size, self.num_heads, physical_length, self.k_head_dim), dtype=self.dtype, device=self.device + def _allocate_key_value_buffers(self, prefill_length: int, copy_existing: bool = False) -> None: + """(Re)allocate static key value buffers.""" + total_length = prefill_length + self.sliding_window + new_all_keys = torch.zeros( + (self.batch_size, self.num_heads, total_length, self.k_head_dim), dtype=self.dtype, device=self.device ) - new_values = torch.zeros( - (self.batch_size, self.num_heads, physical_length, self.v_head_dim), dtype=self.dtype, device=self.device + new_all_values = torch.zeros( + (self.batch_size, self.num_heads, total_length, self.v_head_dim), dtype=self.dtype, device=self.device ) if copy_existing: - old_length = self.keys.shape[-2] - new_keys[:, :, :old_length, :].copy_(self.keys) - new_values[:, :, :old_length, :].copy_(self.values) - self.keys = new_keys - self.values = new_values + old_length = self.all_keys.shape[-2] + new_all_keys[:, :, :old_length, :].copy_(self.all_keys) + new_all_values[:, :, :old_length, :].copy_(self.all_values) + self.all_keys = new_all_keys + self.all_values = new_all_values + + # Views into all_keys/all_values + self.prefill_keys = self.all_keys[:, :, :prefill_length, :] + self.prefill_values = self.all_values[:, :, :prefill_length, :] + self.keys = self.all_keys[:, :, prefill_length:, :] + self.values = self.all_values[:, :, prefill_length:, :] + if not is_torchdynamo_compiling(): + torch._dynamo.mark_static_address(self.all_keys) + torch._dynamo.mark_static_address(self.all_values) + torch._dynamo.mark_static_address(self.prefill_keys) + torch._dynamo.mark_static_address(self.prefill_values) torch._dynamo.mark_static_address(self.keys) torch._dynamo.mark_static_address(self.values) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 54c4128d86a7..dcbbee3b3a83 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -733,6 +733,7 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: if self.prefill_length is None or self.prefill_cumulative_length < self.prefill_length: kv_offset = 0 kv_length = self.prefill_cumulative_length + query_length + # Decode else: kv_length, kv_offset = super().get_mask_sizes(query_length) kv_length += self.prefill_length @@ -741,7 +742,7 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: # Remove kv_offset from kv_idx to retrieve the prefill indices. return kv_length, kv_offset - def get_seq_length(self): + def get_seq_length(self) -> int: return self.prefill_cumulative_length + super().get_seq_length() def get_max_length(self) -> int: @@ -803,7 +804,7 @@ def crop(self, max_length: int) -> None: class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): """ A static cache layer that stores the key and value states as static tensors of shape - `[batch_size, num_heads, min(max_cache_len, prefill_size + sliding_window), head_dim]`. + `[batch_size, num_heads, min(max_cache_len, prefill_length + sliding_window), head_dim]`. It lazily allocates its full backing tensors, and then mutates them in-place. Built for `torch.compile` support. @@ -826,11 +827,20 @@ class UnlimitedOcrStaticReferenceSlidingWindowLayer(StaticSlidingWindowLayer): def __init__(self, max_cache_len: int, sliding_window: int, **kwargs): super().__init__(max_cache_len=max_cache_len, sliding_window=sliding_window, **kwargs) - # Keep `max_cache_len` as max value for length bookkeeping. - # The physical buffer doesn't exceed `prefill_length + sliding_window`. - self.max_cache_len = max_cache_len - self.sliding_window = sliding_window + # Effective sliding window length is determined by max_cache_len in parent + self.sliding_window = self.max_cache_len + self.prefill_max_cache_len = max_cache_len - self.sliding_window self.prefill_length: int | None = None + self.prefill_cumulative_length: int = 0 + self.prefill_keys: torch.Tensor | None = None + self.prefill_values: torch.Tensor | None = None + + # all_keys and all_values are contiguous buffers that hold the full cache to keep dynamo addresses static. + # prefill_keys, prefill_values, keys, and values are views of the buffers: + # all_keys = [prefill_keys, keys] + # all_values = [prefill_values, values] + self.all_keys: torch.Tensor | None = None + self.all_values: torch.Tensor | None = None def update( self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs @@ -849,106 +859,98 @@ def update( if not self.is_initialized: self.lazy_initialization(key_states, value_states) - kv_length = key_states.shape[-2] - # Prefill - if self.prefill_length is None or self.cumulative_length_int < self.prefill_length: + if self.prefill_length is None or self.prefill_cumulative_length < self.prefill_length: + kv_length = key_states.shape[-2] + if self.prefill_length is None: - self.prefill_length = self.cumulative_length_int + kv_length + self.prefill_length = min(self.prefill_max_cache_len, self.prefill_cumulative_length + kv_length) # Resize buffer if necessary - required_length = min(self.max_cache_len, self.prefill_length + self.sliding_window) - if self.keys.shape[-2] < required_length: - self._allocate_key_value_buffers(required_length, copy_existing=True) + if self.prefill_keys.shape[-2] != self.prefill_length: + self._allocate_key_value_buffers(self.prefill_length, copy_existing=True) - # Note: very important to use the tensor version of the cumulative length here, as otherwise cudagraphs - # (triggered by mode="reduced_overhead") will lead to random crashes, as the int would be overwritten - cache_position = torch.arange(kv_length, device=self.device) + self.cumulative_length + prefill_kv_length = min(kv_length, self.prefill_length - self.prefill_cumulative_length) + prefill_key_states = key_states[..., :prefill_kv_length, :] + prefill_value_states = value_states[..., :prefill_kv_length, :] + + cache_position = torch.arange(prefill_kv_length, device=self.device) + self.prefill_cumulative_length try: - self.keys.index_copy_(2, cache_position, key_states) - self.values.index_copy_(2, cache_position, value_states) + self.prefill_keys.index_copy_(2, cache_position, prefill_key_states) + self.prefill_values.index_copy_(2, cache_position, prefill_value_states) except NotImplementedError: # Fallback for devices like MPS where index_copy_ might not be supported. - self.keys[:, :, cache_position] = key_states - self.values[:, :, cache_position] = value_states + self.prefill_keys[:, :, cache_position] = prefill_key_states + self.prefill_values[:, :, cache_position] = prefill_value_states - # Keep both the int (control flow) and the tensor (cudagraph-safe indexing) versions in sync. - self.cumulative_length_int += kv_length - self.cumulative_length.add_(kv_length) + self.prefill_cumulative_length += prefill_kv_length - # Very important to return the `self` tensors here, as they have the static dynamo address - return self.keys, self.values + if prefill_kv_length == kv_length: + return self.all_keys, self.all_values + + key_states = key_states[..., prefill_kv_length:, :] + value_states = value_states[..., prefill_kv_length:, :] # Decode - # Call `StaticSlidingWindowLayer.update` to reduce code duplication. This requires setting temporary - # attributes to match expectations of parent class. - keys, values = self.keys, self.values - window = slice(self.prefill_length, self.prefill_length + self.sliding_window) - sliding_keys_buffer, sliding_values_buffer = keys[:, :, window, :], values[:, :, window, :] - self.keys, self.values = sliding_keys_buffer, sliding_values_buffer - self.cumulative_length -= self.prefill_length - self.cumulative_length_int -= self.prefill_length - max_cache_len, self.max_cache_len = self.max_cache_len, self.sliding_window - try: - sliding_keys, sliding_values = super().update(key_states, value_states, *args, **kwargs) - finally: - self.keys, self.values = keys, values - self.cumulative_length += self.prefill_length - self.cumulative_length_int += self.prefill_length - self.max_cache_len = max_cache_len + sliding_keys, sliding_values = super().update(key_states, value_states, *args, **kwargs) # The parent returned the original buffers - if sliding_keys is sliding_keys_buffer: - return self.keys, self.values - # The parent returned concatenated states + if sliding_keys is self.keys: + return self.all_keys, self.all_values + + # The parent returned concatenated full_key_states/full_value_states return ( - torch.cat((self.keys[:, :, : self.prefill_length, :], sliding_keys), dim=-2), - torch.cat((self.values[:, :, : self.prefill_length, :], sliding_values), dim=-2), + torch.cat((self.prefill_keys, sliding_keys), dim=-2), + torch.cat((self.prefill_values, sliding_values), dim=-2), ) def get_mask_sizes(self, query_length: int) -> tuple[int, int]: """Return the length and offset of the cache, used to generate the attention mask""" - is_full = ( - self.prefill_length is not None and self.cumulative_length_int >= self.prefill_length + self.sliding_window - ) - - kv_offset = 0 - # Prefill: the buffer is not necessarily allocated yet, so its size cannot be used - if self.prefill_length is None or self.cumulative_length_int < self.prefill_length: + # Prefill + if self.prefill_length is None or self.prefill_cumulative_length < self.prefill_length: + kv_offset = 0 prefill_length = ( - self.cumulative_length_int + query_length if self.prefill_length is None else self.prefill_length + min(self.prefill_max_cache_len, self.prefill_cumulative_length + query_length) + if self.prefill_length is None + else self.prefill_length + ) + kv_length = max( + prefill_length + self.sliding_window, + # sliding window returns concatenated states + self.prefill_cumulative_length + self.cumulative_length_int + query_length, ) - kv_length = min(self.max_cache_len, prefill_length + self.sliding_window) - # Decode: cache is already full - elif is_full: - kv_offset = max(self.cumulative_length_int - self.prefill_length - self.sliding_window + 1, 0) - kv_length = self.prefill_length + self.sliding_window - 1 + query_length - # Decode: cache not yet full, but becoming full on this update - elif ( - self.prefill_length is not None - and self.cumulative_length_int + query_length > self.prefill_length + self.sliding_window - ): - kv_length = self.cumulative_length_int + query_length - # Decode: cache not yet full but we return the local size as it's static + # Decode else: - kv_length = self.keys.shape[-2] + kv_length, kv_offset = super().get_mask_sizes(query_length) + kv_length += self.prefill_length + # Returned kv_offset is with respect to sliding window keys. + # Remove kv_offset from kv_idx to retrieve the prefill indices. return kv_length, kv_offset def get_max_length(self) -> int: """Return the maximum cache shape of the cache""" + max_cache_len = self.prefill_max_cache_len + self.max_cache_len if self.prefill_length is None: - return self.max_cache_len - return min(self.max_cache_len, self.prefill_length + self.sliding_window) + return max_cache_len + return min(max_cache_len, self.prefill_length + self.sliding_window) + + def get_seq_length(self) -> int: + return self.prefill_cumulative_length + super().get_seq_length() def set_prefill_length(self, prefill_length: int) -> None: """Declare how many leading tokens are prefill states, before they are cached.""" if self.prefill_length is None: - self.prefill_length = prefill_length + self.prefill_length = min(prefill_length, self.prefill_max_cache_len) def reset(self) -> None: super().reset() + if self.is_initialized: + # This also zeros the self.prefill_keys, self.keys, etc. views + self.all_keys.zero_() + self.all_values.zero_() self.prefill_length = None + self.prefill_cumulative_length = 0 def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None: self.dtype, self.device = key_states.dtype, key_states.device @@ -961,29 +963,37 @@ def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tens torch._dynamo.mark_static_address(self.cumulative_length) prefill_length = key_states.shape[-2] if self.prefill_length is None else self.prefill_length - self._allocate_key_value_buffers(min(self.max_cache_len, prefill_length + self.sliding_window)) + self._allocate_key_value_buffers(min(self.prefill_max_cache_len, prefill_length)) self.is_initialized = True - def _allocate_key_value_buffers(self, physical_length: int, copy_existing: bool = False) -> None: - """(Re)allocate the static key/value buffers to `physical_length` slots and (re)tag the static address. - - Only ever called from the eager prefill pass (initial allocation or chunked-prefill growth), never from a - compiled decode step, so reallocating here is safe for cudagraphs. - """ - new_keys = torch.zeros( - (self.batch_size, self.num_heads, physical_length, self.k_head_dim), dtype=self.dtype, device=self.device + def _allocate_key_value_buffers(self, prefill_length: int, copy_existing: bool = False) -> None: + """(Re)allocate static key value buffers.""" + total_length = prefill_length + self.sliding_window + new_all_keys = torch.zeros( + (self.batch_size, self.num_heads, total_length, self.k_head_dim), dtype=self.dtype, device=self.device ) - new_values = torch.zeros( - (self.batch_size, self.num_heads, physical_length, self.v_head_dim), dtype=self.dtype, device=self.device + new_all_values = torch.zeros( + (self.batch_size, self.num_heads, total_length, self.v_head_dim), dtype=self.dtype, device=self.device ) if copy_existing: - old_length = self.keys.shape[-2] - new_keys[:, :, :old_length, :].copy_(self.keys) - new_values[:, :, :old_length, :].copy_(self.values) - self.keys = new_keys - self.values = new_values + old_length = self.all_keys.shape[-2] + new_all_keys[:, :, :old_length, :].copy_(self.all_keys) + new_all_values[:, :, :old_length, :].copy_(self.all_values) + self.all_keys = new_all_keys + self.all_values = new_all_values + + # Views into all_keys/all_values + self.prefill_keys = self.all_keys[:, :, :prefill_length, :] + self.prefill_values = self.all_values[:, :, :prefill_length, :] + self.keys = self.all_keys[:, :, prefill_length:, :] + self.values = self.all_values[:, :, prefill_length:, :] + if not is_torchdynamo_compiling(): + torch._dynamo.mark_static_address(self.all_keys) + torch._dynamo.mark_static_address(self.all_values) + torch._dynamo.mark_static_address(self.prefill_keys) + torch._dynamo.mark_static_address(self.prefill_values) torch._dynamo.mark_static_address(self.keys) torch._dynamo.mark_static_address(self.values) From 20654394976dcf57ddf7e01785696f7d8303a972 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:00:29 +0100 Subject: [PATCH 181/195] add comment --- src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py | 2 ++ src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 9622af93ba74..60210d283418 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1332,6 +1332,7 @@ def update( if prefill_kv_length == kv_length: return self.prefill_keys, self.prefill_values + # Old speculative methods can draft tokens before prefill is complete key_states = key_states[..., prefill_kv_length:, :] value_states = value_states[..., prefill_kv_length:, :] @@ -1506,6 +1507,7 @@ def update( if prefill_kv_length == kv_length: return self.all_keys, self.all_values + # Old speculative methods can draft tokens before prefill is complete key_states = key_states[..., prefill_kv_length:, :] value_states = value_states[..., prefill_kv_length:, :] diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index dcbbee3b3a83..5abb0ea19578 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -714,6 +714,7 @@ def update( if prefill_kv_length == kv_length: return self.prefill_keys, self.prefill_values + # Old speculative methods can draft tokens before prefill is complete key_states = key_states[..., prefill_kv_length:, :] value_states = value_states[..., prefill_kv_length:, :] @@ -888,6 +889,7 @@ def update( if prefill_kv_length == kv_length: return self.all_keys, self.all_values + # Old speculative methods can draft tokens before prefill is complete key_states = key_states[..., prefill_kv_length:, :] value_states = value_states[..., prefill_kv_length:, :] From 80d08e67613dcf17946ba765b5cf4161c5b3793f Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:44:40 +0100 Subject: [PATCH 182/195] update test utils --- tests/generation/test_utils.py | 18 ++++++------ .../test_modeling_unlimited_ocr.py | 28 +++++++------------ 2 files changed, 19 insertions(+), 27 deletions(-) diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py index ddbd6863541c..45e63127d910 100644 --- a/tests/generation/test_utils.py +++ b/tests/generation/test_utils.py @@ -2736,13 +2736,6 @@ def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_l num_hidden_layers -= config.num_kv_shared_layers self.assertEqual(num_hidden_layers, len(past_key_values)) - def check_attention_shapes(layer, attention_shape): - # Remove the seq_length dim for cross-attention cache (it changes based on the model) - keys = layer.keys if seq_length is not None else layer.keys[:, :, 0, :] - values = layer.values if seq_length is not None else layer.values[:, :, 0, :] - self.assertEqual(keys.shape, attention_shape) - self.assertEqual(values.shape, attention_shape) - def check_linear_attention_shapes(layer, num_conv_states, conv_shape, recurrent_shape): # assert we have as many conv states as necessary self.assertEqual(num_conv_states, layer.number_of_states) @@ -2768,14 +2761,21 @@ def check_linear_attention_shapes(layer, num_conv_states, conv_shape, recurrent_ attention_shape = self._get_attention_shape(batch_size, seq_length, layer_config) # Mamba + Attention layer cache if type(layer) in (LinearAttentionAndFullAttentionLayer, LinearAttentionAndSlidingWindowAttentionLayer): - check_attention_shapes(layer, attention_shape) + self._check_attention_shapes(layer, attention_shape, seq_length) check_linear_attention_shapes(layer, num_conv_states, conv_shape, recurrent_shape) # Mamba only layer cache elif type(layer) is LinearAttentionLayer: check_linear_attention_shapes(layer, num_conv_states, conv_shape, recurrent_shape) # Attention only layer type else: - check_attention_shapes(layer, attention_shape) + self._check_attention_shapes(layer, attention_shape, seq_length) + + def _check_attention_shapes(self, layer, attention_shape, seq_length): + # Remove the seq_length dim for cross-attention cache (it changes based on the model) + keys = layer.keys if seq_length is not None else layer.keys[:, :, 0, :] + values = layer.values if seq_length is not None else layer.values[:, :, 0, :] + self.assertEqual(keys.shape, attention_shape) + self.assertEqual(values.shape, attention_shape) def _get_attention_shape(self, batch_size: int, seq_length: int | None, config): # Only pure mamba models do not have num_attention_heads defined in config, so it can never be 1 in practice for attention models diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index ef15eed7c75d..8b200fb1e217 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -145,31 +145,23 @@ def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_l # max cache length can be smaller than sequence length max_length = past_key_values.get_max_length() seq_length = min(seq_length, max_length) if max_length >= 0 else seq_length + super()._check_past_key_values_for_generate(batch_size, past_key_values, seq_length, config) - if not any(hasattr(layer, "prefill_keys") for layer in past_key_values.layers): - super()._check_past_key_values_for_generate(batch_size, past_key_values, seq_length, config) - return - + def _check_attention_shapes(self, layer, attention_shape, seq_length): # Super method assumes that there is only layer.keys/values but reference sliding window layers # can have layer.prefill_keys/prefill_values as well which results in wrong assertions in the # super method. def seq_len(tensor): return tensor.shape[-2] if tensor.dim() > 1 else 0 - config = config.get_text_config(decoder=True) - self.assertEqual(config.num_hidden_layers, len(past_key_values)) - for layer_idx, layer in enumerate(past_key_values.layers): - layer_config = config.per_layer_config[layer_idx] - attention_shape = self._get_attention_shape(batch_size, seq_length, layer_config) - combined_keys_length = seq_len(layer.prefill_keys) + seq_len(layer.keys) - combined_values_length = seq_len(layer.prefill_values) + seq_len(layer.values) - self.assertEqual( - (*layer.prefill_keys.shape[:-2], combined_keys_length, layer.prefill_keys.shape[-1]), attention_shape - ) - self.assertEqual( - (*layer.prefill_values.shape[:-2], combined_values_length, layer.prefill_values.shape[-1]), - attention_shape, - ) + combined_keys_length = seq_len(layer.prefill_keys) + seq_len(layer.keys) + combined_values_length = seq_len(layer.prefill_values) + seq_len(layer.values) + self.assertEqual( + (*layer.prefill_keys.shape[:-2], combined_keys_length, layer.prefill_keys.shape[-1]), attention_shape + ) + self.assertEqual( + (*layer.prefill_values.shape[:-2], combined_values_length, layer.prefill_values.shape[-1]), attention_shape + ) def _check_generate_cache_sliding_window_too_small(self, cache_implementation: str, prefill_max_new_tokens: int): """Test that reference sliding window cache works correctly when decoding more than sliding_window tokens at once.""" From 1a4109aead54d1ab40654433326d40d890019084 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:55:37 +0100 Subject: [PATCH 183/195] clarify comment --- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 4 ++-- .../models/unlimited_ocr/modular_unlimited_ocr.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 60210d283418..5a060fc3161c 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1536,8 +1536,8 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: ) kv_length = max( prefill_length + self.sliding_window, - # sliding window returns concatenated states - self.prefill_cumulative_length + self.cumulative_length_int + query_length, + # Query exceeds prefill and maybe also sliding window + self.prefill_cumulative_length + super().get_seq_length() + query_length, ) # Decode else: diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 5abb0ea19578..a5bf3180e050 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -918,8 +918,8 @@ def get_mask_sizes(self, query_length: int) -> tuple[int, int]: ) kv_length = max( prefill_length + self.sliding_window, - # sliding window returns concatenated states - self.prefill_cumulative_length + self.cumulative_length_int + query_length, + # Query exceeds prefill and maybe also sliding window + self.prefill_cumulative_length + super().get_seq_length() + query_length, ) # Decode else: From 5102060ab8d28f49b8f8696cba74fabcf8369dab Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 31 Jul 2026 17:00:40 +0100 Subject: [PATCH 184/195] fix-repo --- .../unlimited_ocr/modeling_unlimited_ocr.py | 41 ++++++++----------- 1 file changed, 16 insertions(+), 25 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 5a060fc3161c..b27ec13764c3 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -21,7 +21,6 @@ import math from collections.abc import Callable from dataclasses import dataclass -from typing import Optional import torch import torch.nn.functional as F @@ -31,12 +30,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache, DynamicSlidingWindowLayer, StaticSlidingWindowLayer from ...configuration_utils import PreTrainedConfig -from ...integrations import ( - use_experts_implementation, - use_kernel_forward_from_hub, - use_kernel_func_from_hub, - use_kernelized_func, -) +from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func from ...masking_utils import ( BlockMask, and_masks, @@ -60,6 +54,7 @@ torch_compilable_check, torch_int, ) +from ...utils.deprecation import deprecate_kwarg from ...utils.generic import maybe_autocast, merge_with_config_defaults from ...utils.output_capturing import capture_outputs from .configuration_unlimited_ocr import ( @@ -977,7 +972,7 @@ def rotate_half(x): return torch.cat((-x2, x1), dim=-1) -@use_kernel_func_from_hub("rotary_pos_emb") +@use_kernel_forward_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. @@ -1621,6 +1616,7 @@ def _allocate_key_value_buffers(self, prefill_length: int, copy_existing: bool = class UnlimitedOcrTextRotaryEmbedding(nn.Module): inv_freq: torch.Tensor # fix linting for `register_buffer` + @deprecate_kwarg("device", version="5.18") def __init__(self, config: UnlimitedOcrTextConfig, device=None): super().__init__() self.max_seq_len_cached = config.max_position_embeddings @@ -1638,20 +1634,15 @@ def __init__(self, config: UnlimitedOcrTextConfig, device=None): self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) @staticmethod + @deprecate_kwarg("device", version="5.18") def compute_default_rope_parameters( - config: UnlimitedOcrTextConfig | None = None, - device: Optional["torch.device"] = None, - seq_len: int | None = None, - ) -> tuple["torch.Tensor", float]: + config: UnlimitedOcrTextConfig, device=None, **kwargs + ) -> 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). @@ -1660,22 +1651,22 @@ def compute_default_rope_parameters( 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 + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) + return inv_freq.to(device), 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) + inv_freq_expanded = ( + self.inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1).to(dtype=torch.float, device=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) + # Disable any outside autocast context if any, to really force fp32 + with maybe_autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() * self.attention_scaling sin = emb.sin() * self.attention_scaling @@ -1926,7 +1917,7 @@ def get_placeholder_mask( """ 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) + torch.full((), self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) ) special_image_mask = special_image_mask.all(-1) else: From 54b13d550c908eb6d50910f25822981904e2f0ee Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 21 Aug 2026 10:40:31 +0100 Subject: [PATCH 185/195] update --- docs/source/en/model_doc/unlimited_ocr.md | 2 +- .../image_processing_unlimited_ocr.py | 23 +++++++++---------- .../unlimited_ocr/modeling_unlimited_ocr.py | 18 +++++++-------- .../unlimited_ocr/modular_unlimited_ocr.py | 18 +++++++-------- tests/generation/test_utils.py | 6 ++--- .../test_modeling_unlimited_ocr.py | 8 +++---- 6 files changed, 35 insertions(+), 40 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 290aa0cc91a7..2d0b9ca06f3c 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -16,7 +16,7 @@ limitations under the License. ⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer. --> -*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-07-31.* +*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-08-21.* # UnlimitedOcr diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index ea66ee4df057..6fb35b9897c2 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -238,8 +238,8 @@ def _preprocess( num_local_patches_grouped = {} grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) - for shape, stacked_images in grouped_images.items(): - height, width = shape[-2:] + for key, stacked_images in grouped_images.items(): + height, width = stacked_images.shape[-2:] num_images = stacked_images.shape[0] if crop_to_patches and max(height, width) > tile_size: num_columns, num_rows = get_optimal_tiled_canvas( @@ -256,12 +256,12 @@ def _preprocess( flat_patches = self.rescale_and_normalize( flat_patches, do_rescale, rescale_factor, do_normalize, image_mean, image_std ) - local_patches_grouped[shape] = flat_patches.reshape(stacked_patches.shape) + local_patches_grouped[key] = flat_patches.reshape(stacked_patches.shape) else: num_columns, num_rows, num_patches = 1, 1, 0 - local_patches_grouped[shape] = [None] * num_images - patches_grid_grouped[shape] = [[num_columns, num_rows]] * num_images - num_local_patches_grouped[shape] = [num_patches] * num_images + local_patches_grouped[key] = [None] * num_images + patches_grid_grouped[key] = [[num_columns, num_rows]] * num_images + num_local_patches_grouped[key] = [num_patches] * num_images ordered_local = reorder_images(local_patches_grouped, grouped_images_index) patches_grid = reorder_images(patches_grid_grouped, grouped_images_index) @@ -274,14 +274,14 @@ def _preprocess( global_target_size = max(size.height, size.width) processed_global_grouped = {} - for shape, stacked in grouped_images.items(): + for key, stacked in grouped_images.items(): # Different from DeepseekOcr2 which crops and pads all images if not crop_to_patches and global_target_size <= maximum_pad_value: stacked = self.resize( stacked, SizeDict(height=global_target_size, width=global_target_size), resample=resample ) else: - height, width = shape[-2:] + height, width = stacked.shape[-2:] scale = global_target_size / max(height, width) new_height = round(height * scale) new_width = round(width * scale) @@ -290,7 +290,7 @@ def _preprocess( stacked = self.rescale_and_normalize( stacked, do_rescale, rescale_factor, do_normalize, image_mean, image_std ) - processed_global_grouped[shape] = stacked + processed_global_grouped[key] = stacked all_pixel_values_global = reorder_images(processed_global_grouped, grouped_images_index) data = {"pixel_values": all_pixel_values_global} @@ -304,12 +304,11 @@ def _preprocess( tensor_type=return_tensors, ) - def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None) -> int: + def get_number_of_image_patches(self, height: int, width: int, images_kwargs: dict | None = None) -> int: """ Returns the number of image patches for a given image size (1 global + local patches). """ - if images_kwargs is None: - images_kwargs = {} + images_kwargs = images_kwargs or {} min_patches = images_kwargs.get("min_patches", self.min_patches) max_patches = images_kwargs.get("max_patches", self.max_patches) tile_size = images_kwargs.get("tile_size", self.tile_size) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index b27ec13764c3..e0e4bc5d261b 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -276,11 +276,6 @@ class UnlimitedOcrPreTrainedModel(PreTrainedModel): base_model_prefix = "model" input_modalities = ("image", "text") supports_gradient_checkpointing = True - _no_split_modules = [ - "UnlimitedOcrEncoderLayer", - "UnlimitedOcrSamVisionLayer", - "UnlimitedOcrTextDecoderLayer", - ] _skip_keys_device_placement = ["past_key_values"] # SAM uses rel-pos bias, incompatible with flash attention. _supports_flash_attn = False @@ -290,6 +285,11 @@ class UnlimitedOcrPreTrainedModel(PreTrainedModel): # SAM doesn't support flex attention _supports_flex_attn = False _supports_attention_backend = True + _no_split_modules = [ + "UnlimitedOcrEncoderLayer", + "UnlimitedOcrSamVisionLayer", + "UnlimitedOcrTextDecoderLayer", + ] @torch.no_grad() def _init_weights(self, module): @@ -782,7 +782,7 @@ def __init__(self, config: UnlimitedOcrVisionConfig): self.num_patches = (self.image_size // self.patch_size) ** 2 self.num_positions = self.num_patches + 1 self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) - self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False) + self.position_ids = nn.Buffer(torch.arange(self.num_positions).expand((1, -1)), persistent=False) def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: """ @@ -1614,8 +1614,6 @@ def _allocate_key_value_buffers(self, prefill_length: int, copy_existing: bool = class UnlimitedOcrTextRotaryEmbedding(nn.Module): - inv_freq: torch.Tensor # fix linting for `register_buffer` - @deprecate_kwarg("device", version="5.18") def __init__(self, config: UnlimitedOcrTextConfig, device=None): super().__init__() @@ -1630,8 +1628,8 @@ def __init__(self, config: UnlimitedOcrTextConfig, device=None): 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) + self.inv_freq = nn.Buffer(inv_freq, persistent=False) + self.original_inv_freq = nn.Buffer(inv_freq.clone(), persistent=False) @staticmethod @deprecate_kwarg("device", version="5.18") diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index a5bf3180e050..1f69fb58e126 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -134,8 +134,8 @@ def _preprocess( num_local_patches_grouped = {} grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) - for shape, stacked_images in grouped_images.items(): - height, width = shape[-2:] + for key, stacked_images in grouped_images.items(): + height, width = stacked_images.shape[-2:] num_images = stacked_images.shape[0] if crop_to_patches and max(height, width) > tile_size: num_columns, num_rows = get_optimal_tiled_canvas( @@ -152,12 +152,12 @@ def _preprocess( flat_patches = self.rescale_and_normalize( flat_patches, do_rescale, rescale_factor, do_normalize, image_mean, image_std ) - local_patches_grouped[shape] = flat_patches.reshape(stacked_patches.shape) + local_patches_grouped[key] = flat_patches.reshape(stacked_patches.shape) else: num_columns, num_rows, num_patches = 1, 1, 0 - local_patches_grouped[shape] = [None] * num_images - patches_grid_grouped[shape] = [[num_columns, num_rows]] * num_images - num_local_patches_grouped[shape] = [num_patches] * num_images + local_patches_grouped[key] = [None] * num_images + patches_grid_grouped[key] = [[num_columns, num_rows]] * num_images + num_local_patches_grouped[key] = [num_patches] * num_images ordered_local = reorder_images(local_patches_grouped, grouped_images_index) patches_grid = reorder_images(patches_grid_grouped, grouped_images_index) @@ -170,14 +170,14 @@ def _preprocess( global_target_size = max(size.height, size.width) processed_global_grouped = {} - for shape, stacked in grouped_images.items(): + for key, stacked in grouped_images.items(): # Different from DeepseekOcr2 which crops and pads all images if not crop_to_patches and global_target_size <= maximum_pad_value: stacked = self.resize( stacked, SizeDict(height=global_target_size, width=global_target_size), resample=resample ) else: - height, width = shape[-2:] + height, width = stacked.shape[-2:] scale = global_target_size / max(height, width) new_height = round(height * scale) new_width = round(width * scale) @@ -186,7 +186,7 @@ def _preprocess( stacked = self.rescale_and_normalize( stacked, do_rescale, rescale_factor, do_normalize, image_mean, image_std ) - processed_global_grouped[shape] = stacked + processed_global_grouped[key] = stacked all_pixel_values_global = reorder_images(processed_global_grouped, grouped_images_index) data = {"pixel_values": all_pixel_values_global} diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py index fedae855c753..c574c8e3ce40 100644 --- a/tests/generation/test_utils.py +++ b/tests/generation/test_utils.py @@ -2827,16 +2827,16 @@ def check_linear_attention_shapes(layer, num_conv_states, conv_shape, recurrent_ k_shape, v_shape = self._get_attention_shape(batch_size, seq_length, layer_config) # Mamba + Attention layer cache if type(layer) in (LinearAttentionAndFullAttentionLayer, LinearAttentionAndSlidingWindowAttentionLayer): - self._check_attention_shapes(layer, k_shape, v_shape) + self._check_attention_shapes(layer, seq_length, k_shape, v_shape) check_linear_attention_shapes(layer, num_conv_states, conv_shape, recurrent_shape) # Mamba only layer cache elif type(layer) is LinearAttentionLayer: check_linear_attention_shapes(layer, num_conv_states, conv_shape, recurrent_shape) # Attention only layer type else: - self._check_attention_shapes(layer, k_shape, v_shape) + self._check_attention_shapes(layer, seq_length, k_shape, v_shape) - def _check_attention_shapes(self, layer, k_shape, v_shape): + def _check_attention_shapes(self, layer, seq_length, k_shape, v_shape): # Remove the seq_length dim for cross-attention cache (it changes based on the model) keys = layer.keys if seq_length is not None else layer.keys[:, :, 0, :] values = layer.values if seq_length is not None else layer.values[:, :, 0, :] diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 8b200fb1e217..a2a7f034c6ef 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -147,7 +147,7 @@ def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_l seq_length = min(seq_length, max_length) if max_length >= 0 else seq_length super()._check_past_key_values_for_generate(batch_size, past_key_values, seq_length, config) - def _check_attention_shapes(self, layer, attention_shape, seq_length): + def _check_attention_shapes(self, layer, seq_length, k_shape, v_shape): # Super method assumes that there is only layer.keys/values but reference sliding window layers # can have layer.prefill_keys/prefill_values as well which results in wrong assertions in the # super method. @@ -156,11 +156,9 @@ def seq_len(tensor): combined_keys_length = seq_len(layer.prefill_keys) + seq_len(layer.keys) combined_values_length = seq_len(layer.prefill_values) + seq_len(layer.values) + self.assertEqual((*layer.prefill_keys.shape[:-2], combined_keys_length, layer.prefill_keys.shape[-1]), k_shape) self.assertEqual( - (*layer.prefill_keys.shape[:-2], combined_keys_length, layer.prefill_keys.shape[-1]), attention_shape - ) - self.assertEqual( - (*layer.prefill_values.shape[:-2], combined_values_length, layer.prefill_values.shape[-1]), attention_shape + (*layer.prefill_values.shape[:-2], combined_values_length, layer.prefill_values.shape[-1]), v_shape ) def _check_generate_cache_sliding_window_too_small(self, cache_implementation: str, prefill_max_new_tokens: int): From adb0e0e11a67d8aeeb00fa1393ce04f8755d607f Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:11:56 +0100 Subject: [PATCH 186/195] fixes --- .../deepseek_ocr2/modeling_deepseek_ocr2.py | 3 ++- .../deepseek_ocr2/modular_deepseek_ocr2.py | 3 ++- .../unlimited_ocr/modeling_unlimited_ocr.py | 17 ++++++++--------- .../unlimited_ocr/modular_unlimited_ocr.py | 14 +++++--------- 4 files changed, 17 insertions(+), 20 deletions(-) diff --git a/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py index c20ec1a90b31..05ff156d049e 100644 --- a/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py @@ -54,9 +54,10 @@ logger = logging.get_logger(__name__) +@auto_docstring @dataclass class DeepseekOcr2ModelOutputWithPooling(BaseModelOutputWithPooling): - """ + r""" 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*): diff --git a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py index a853ab2f5fee..6085654976d6 100644 --- a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py @@ -649,9 +649,10 @@ def __post_init__(self, **kwargs): super().__post_init__(**kwargs) +@auto_docstring @dataclass class DeepseekOcr2ModelOutputWithPooling(BaseModelOutputWithPooling): - """ + r""" 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*): diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index e0e4bc5d261b..1962a047b14f 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -70,9 +70,10 @@ logger = logging.get_logger(__name__) +@auto_docstring @dataclass class UnlimitedOcrModelOutputWithPooling(BaseModelOutputWithPooling): - """ + r""" 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*): @@ -914,8 +915,9 @@ def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwarg return BaseModelOutput(last_hidden_state=encoder_outputs.last_hidden_state) +@auto_docstring(custom_intro="Vision model encoding images first with SAM followed by an additional (CLIP) model.") class UnlimitedOcrVisionModel(UnlimitedOcrPreTrainedModel): - """Vision model encoding images first with SAM followed by an additional (CLIP) model.""" + """Vision pipeline: SAM ViT-B (with neck)""" def __init__(self, config: UnlimitedOcrVisionConfig): super().__init__(config) @@ -1365,7 +1367,7 @@ def get_max_length(self) -> int: return -1 return self.prefill_length + self.sliding_window - def set_prefill_length(self, prefill_length: int) -> None: + def set_prefill_length(self, prefill_length: int) -> None: # trf-ignore: TRF033 """Declare how many leading tokens are prefill states, before they are cached.""" if self.prefill_length is None: self.prefill_length = prefill_length @@ -1553,7 +1555,7 @@ def get_max_length(self) -> int: def get_seq_length(self) -> int: return self.prefill_cumulative_length + super().get_seq_length() - def set_prefill_length(self, prefill_length: int) -> None: + def set_prefill_length(self, prefill_length: int) -> None: # trf-ignore: TRF033 """Declare how many leading tokens are prefill states, before they are cached.""" if self.prefill_length is None: self.prefill_length = min(prefill_length, self.prefill_max_cache_len) @@ -1776,12 +1778,8 @@ def forward( causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), + "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), } - # The reference sliding window layers are not always activated depending on the config - if "reference_sliding_attention" in self.config.layer_types: - causal_mask_mapping["reference_sliding_attention"] = create_reference_sliding_window_causal_mask( - **mask_kwargs - ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) @@ -1818,6 +1816,7 @@ def __init__(self, config: UnlimitedOcrConfig): self.vision_tower = UnlimitedOcrVisionModel(config.vision_config) self.multi_modal_projector = nn.Linear( + # trf-ignore: TRF030 config.vision_config.sam_config.downsample_channels[-1] + config.vision_config.encoder_config.hidden_size, config.text_config.hidden_size, ) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 1f69fb58e126..72ab64714b60 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -618,9 +618,8 @@ def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwarg return BaseModelOutput(last_hidden_state=encoder_outputs.last_hidden_state) +@auto_docstring(custom_intro="Vision model encoding images first with SAM followed by an additional (CLIP) model.") class UnlimitedOcrVisionModel(DeepseekOcr2VisionModel): - """Vision model encoding images first with SAM followed by an additional (CLIP) model.""" - def __init__(self, config: UnlimitedOcrVisionConfig): super().__init__(config) del self.query_768_resolution @@ -752,7 +751,7 @@ def get_max_length(self) -> int: return -1 return self.prefill_length + self.sliding_window - def set_prefill_length(self, prefill_length: int) -> None: + def set_prefill_length(self, prefill_length: int) -> None: # trf-ignore: TRF033 """Declare how many leading tokens are prefill states, before they are cached.""" if self.prefill_length is None: self.prefill_length = prefill_length @@ -940,7 +939,7 @@ def get_max_length(self) -> int: def get_seq_length(self) -> int: return self.prefill_cumulative_length + super().get_seq_length() - def set_prefill_length(self, prefill_length: int) -> None: + def set_prefill_length(self, prefill_length: int) -> None: # trf-ignore: TRF033 """Declare how many leading tokens are prefill states, before they are cached.""" if self.prefill_length is None: self.prefill_length = min(prefill_length, self.prefill_max_cache_len) @@ -1083,12 +1082,8 @@ def forward( causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), + "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), } - # The reference sliding window layers are not always activated depending on the config - if "reference_sliding_attention" in self.config.layer_types: - causal_mask_mapping["reference_sliding_attention"] = create_reference_sliding_window_causal_mask( - **mask_kwargs - ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) @@ -1117,6 +1112,7 @@ class UnlimitedOcrModel(DeepseekOcr2Model): def __init__(self, config: UnlimitedOcrConfig): super().__init__(config) self.multi_modal_projector = nn.Linear( + # trf-ignore: TRF030 config.vision_config.sam_config.downsample_channels[-1] + config.vision_config.encoder_config.hidden_size, config.text_config.hidden_size, ) From 6c0b38cb01c003d5f8de27cbbab6cf2b9393b27e Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:32:49 +0100 Subject: [PATCH 187/195] revert --- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 7 ++++++- .../models/unlimited_ocr/modular_unlimited_ocr.py | 7 ++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 1962a047b14f..090acdf42ebd 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1778,8 +1778,13 @@ def forward( causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), - "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), } + # CODEPATH: create_sliding_window_causal_mask errors if config.sliding_window is None. + # Without the check here a model with full attention crashes. + if self.config.sliding_window is not None: + causal_mask_mapping["reference_sliding_attention"] = create_reference_sliding_window_causal_mask( + **mask_kwargs + ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 72ab64714b60..e7f2263c7efc 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -1082,8 +1082,13 @@ def forward( causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), - "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), } + # CODEPATH: create_sliding_window_causal_mask errors if config.sliding_window is None. + # Without the check here a model with full attention crashes. + if self.config.sliding_window is not None: + causal_mask_mapping["reference_sliding_attention"] = create_reference_sliding_window_causal_mask( + **mask_kwargs + ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) From 6a087ca9858b7581f5b03cebcd7c60ab0c44b522 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Mon, 24 Aug 2026 09:29:23 +0200 Subject: [PATCH 188/195] cleanup docstrings --- docs/source/en/model_doc/unlimited_ocr.md | 2 +- .../models/deepseek_ocr2/modeling_deepseek_ocr2.py | 3 +-- .../models/deepseek_ocr2/modular_deepseek_ocr2.py | 3 +-- .../models/unlimited_ocr/modeling_unlimited_ocr.py | 6 ++---- .../models/unlimited_ocr/modular_unlimited_ocr.py | 12 ++++++++++++ 5 files changed, 17 insertions(+), 9 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 2d0b9ca06f3c..ca4c66266041 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -16,7 +16,7 @@ limitations under the License. ⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer. --> -*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-08-21.* +*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-08-24.* # UnlimitedOcr diff --git a/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py index 05ff156d049e..4377246984ae 100644 --- a/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/modeling_deepseek_ocr2.py @@ -952,9 +952,8 @@ def forward( return BaseModelOutputWithPast(last_hidden_state=hidden_states) +@auto_docstring(custom_intro="Vision pipeline: SAM ViT-B (with neck)") 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) diff --git a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py index 6085654976d6..8e21ba4d63cb 100644 --- a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py @@ -870,9 +870,8 @@ def forward( return BaseModelOutputWithPast(last_hidden_state=hidden_states) +@auto_docstring(custom_intro="Vision pipeline: SAM ViT-B (with neck)") 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) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 090acdf42ebd..0e54dae10aa7 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -89,7 +89,7 @@ class UnlimitedOcrModelOutputWithPooling(BaseModelOutputWithPooling): @auto_docstring( custom_intro=""" - Base class for Llava outputs, with hidden states and attentions. + Base class for UnlimitedOcr outputs, with hidden states and attentions. """ ) @dataclass @@ -917,8 +917,6 @@ def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwarg @auto_docstring(custom_intro="Vision model encoding images first with SAM followed by an additional (CLIP) model.") class UnlimitedOcrVisionModel(UnlimitedOcrPreTrainedModel): - """Vision pipeline: SAM ViT-B (with neck)""" - def __init__(self, config: UnlimitedOcrVisionConfig): super().__init__(config) self.sam_encoder = UnlimitedOcrSamVisionEncoder(config.sam_config) @@ -1809,7 +1807,7 @@ def forward( @auto_docstring( custom_intro=""" - The Llava-Next model which consists of a vision backbone and a language model without language modeling head. + The UnlimitedOcr model which consists of a vision backbone and a language model. """ ) class UnlimitedOcrModel(UnlimitedOcrPreTrainedModel): diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index e7f2263c7efc..0e0f1bc2bfe4 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -14,6 +14,7 @@ import math import re from collections.abc import Callable +from dataclasses import dataclass import torch from huggingface_hub.dataclasses import strict @@ -481,6 +482,12 @@ class UnlimitedOcrModelOutputWithPooling(DeepseekOcr2ModelOutputWithPooling): pass +@auto_docstring( + custom_intro=""" + Base class for UnlimitedOcr outputs, with hidden states and attentions. + """ +) +@dataclass class UnlimitedOcrModelOutputWithPast(DeepseekOcr2ModelOutputWithPast): pass @@ -1111,6 +1118,11 @@ def forward( ) +@auto_docstring( + custom_intro=""" + The UnlimitedOcr model which consists of a vision backbone and a language model. + """ +) class UnlimitedOcrModel(DeepseekOcr2Model): _keys_to_ignore_on_load_unexpected = {"lm_head"} # unused and untied weight in original checkpoint From fb34ec1ab2aca7e7e2befd902163b03256226a2e Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 25 Aug 2026 16:52:23 +0200 Subject: [PATCH 189/195] fixes --- docs/source/en/model_doc/unlimited_ocr.md | 2 +- src/transformers/generation/utils.py | 3 +- .../configuration_deepseek_ocr2.py | 8 +- .../deepseek_ocr2/modular_deepseek_ocr2.py | 8 +- .../configuration_unlimited_ocr.py | 3 +- .../unlimited_ocr/generation_unlimited_ocr.py | 2 + .../image_processing_unlimited_ocr.py | 6 +- .../unlimited_ocr/modeling_unlimited_ocr.py | 56 ++------- .../unlimited_ocr/modular_unlimited_ocr.py | 108 +++++++----------- .../unlimited_ocr/processing_unlimited_ocr.py | 11 +- tests/generation/test_logits_process.py | 25 ++++ .../test_modeling_unlimited_ocr.py | 2 +- tests/test_processing_common.py | 10 +- 13 files changed, 112 insertions(+), 132 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index ca4c66266041..8cd1d734bd2b 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -65,7 +65,7 @@ inputs = processor.apply_chat_template( output = model.generate( **inputs, - max_new_tokens=32768, + max_new_tokens=512, no_repeat_ngram_size=35, no_repeat_ngram_window_size=128, ) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index edfe829b82f3..5d099cf581f9 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -161,7 +161,8 @@ "image_sizes_videos", "pixel_attention_mask", "pixel_values_images", - "num_local_patches", + "pixel_values_local", + "patches_grid", ) diff --git a/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py index d12eb50849db..78393a62f375 100644 --- a/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/configuration_deepseek_ocr2.py @@ -167,14 +167,14 @@ class DeepseekOcr2VisionConfig(PreTrainedConfig): def __post_init__(self, **kwargs): if self.sam_config is None: - self.sam_config = DeepseekOcr2SamVisionConfig() + self.sam_config = self.sub_configs["sam_config"]() elif isinstance(self.sam_config, dict): - self.sam_config = DeepseekOcr2SamVisionConfig(**self.sam_config) + self.sam_config = self.sub_configs["sam_config"](**self.sam_config) if self.encoder_config is None: - self.encoder_config = DeepseekOcr2VisionEncoderConfig() + self.encoder_config = self.sub_configs["encoder_config"]() elif isinstance(self.encoder_config, dict): - self.encoder_config = DeepseekOcr2VisionEncoderConfig(**self.encoder_config) + self.encoder_config = self.sub_configs["encoder_config"](**self.encoder_config) super().__post_init__(**kwargs) diff --git a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py index 8e21ba4d63cb..564d71bb4cd0 100644 --- a/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py +++ b/src/transformers/models/deepseek_ocr2/modular_deepseek_ocr2.py @@ -556,14 +556,14 @@ class DeepseekOcr2VisionConfig(PreTrainedConfig): def __post_init__(self, **kwargs): if self.sam_config is None: - self.sam_config = DeepseekOcr2SamVisionConfig() + self.sam_config = self.sub_configs["sam_config"]() elif isinstance(self.sam_config, dict): - self.sam_config = DeepseekOcr2SamVisionConfig(**self.sam_config) + self.sam_config = self.sub_configs["sam_config"](**self.sam_config) if self.encoder_config is None: - self.encoder_config = DeepseekOcr2VisionEncoderConfig() + self.encoder_config = self.sub_configs["encoder_config"]() elif isinstance(self.encoder_config, dict): - self.encoder_config = DeepseekOcr2VisionEncoderConfig(**self.encoder_config) + self.encoder_config = self.sub_configs["encoder_config"](**self.encoder_config) super().__post_init__(**kwargs) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index b266591fc02d..b2cfe615153a 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -308,7 +308,8 @@ def __post_init__(self, **kwargs): self.vision_config = self.sub_configs["vision_config"](**self.vision_config) if self.text_config is None: - # The reference implementation defines the text config values on the main config + # For backward compatibility. The reference implementation defines the text config values + # on the main config self.text_config = self.sub_configs["text_config"](**kwargs) elif isinstance(self.text_config, dict): self.text_config = self.sub_configs["text_config"](**self.text_config) diff --git a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py index 67f095fc2e7d..1ac4a797c45f 100644 --- a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py @@ -98,6 +98,7 @@ def _get_logits_processor( negative_prompt_ids: torch.Tensor | None = None, negative_prompt_attention_mask: torch.Tensor | None = None, ) -> LogitsProcessorList: + # Save original config value as we'll have to override it temporarily in this method. no_repeat_ngram_size = generation_config.no_repeat_ngram_size no_repeat_ngram_window_size = getattr(generation_config, "no_repeat_ngram_window_size", None) use_sliding_window_processor = False @@ -115,6 +116,7 @@ def _get_logits_processor( generation_config.no_repeat_ngram_size = None try: + # try/finally to make sure we revert temporary config modification processors = super()._get_logits_processor( generation_config=generation_config, input_ids_seq_length=input_ids_seq_length, diff --git a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py index 6fb35b9897c2..8efacd406750 100644 --- a/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/image_processing_unlimited_ocr.py @@ -44,13 +44,13 @@ class UnlimitedOcrImageProcessorKwargs(ImagesKwargs, total=False): set to `True`. Can be overridden by the `max_patches` parameter in the `preprocess` method. tile_size (`int`, *optional*, defaults to `640`): The size of each local tile. Must match the model's query embedding size. + background_color (`list[int]`, *optional*, defaults to `[127, 127, 127]`): + The background color for padding. maximum_pad_value (`int`, *optional*, defaults to `640`): If `crop_to_patches` is `False` and `max(size.height, size.width)` is smaller than or equal to this value, the image is resized directly to a square of `max(size.height, size.width)` without preserving the aspect ratio. Otherwise, the image is resized while preserving the aspect ratio and then padded to a square with `background_color`. - background_color (`list[int]`, *optional*, defaults to `[127, 127, 127]`): - The background color for padding. """ crop_to_patches: bool @@ -155,6 +155,7 @@ class UnlimitedOcrImageProcessor(TorchvisionBackend): background_color = [127, 127, 127] model_input_names = ["pixel_values", "num_local_patches", "patches_grid"] maximum_pad_value = 640 + skip_tensor_conversion = ["num_local_patches", "patches_grid"] def __init__(self, **kwargs: Unpack[UnlimitedOcrImageProcessorKwargs]): super().__init__(**kwargs) @@ -302,6 +303,7 @@ def _preprocess( return BatchFeature( data=data, tensor_type=return_tensors, + skip_tensor_conversion=self.skip_tensor_conversion, ) def get_number_of_image_patches(self, height: int, width: int, images_kwargs: dict | None = None) -> int: diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 0e54dae10aa7..092edbcfaa81 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -679,7 +679,7 @@ def __init__(self, config: UnlimitedOcrVisionEncoderConfig): self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) - # Required for repeat_kv(..., num_key_value_groups) + # Required for GQA compatibility self.num_key_value_groups = 1 def forward( @@ -837,7 +837,7 @@ def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: class_embeds = self.class_embedding.expand(batch_size, 1, -1) embeddings = torch.cat([class_embeds, patch_embeds], dim=1) - # always interpolate + # always interpolate compared to clip where this can be optional embeddings = embeddings + self.interpolate_pos_encoding( embeddings, grid_height * self.patch_size, grid_width * self.patch_size ) @@ -902,6 +902,7 @@ def __init__(self, config: UnlimitedOcrVisionEncoderConfig): self.encoder = UnlimitedOcrEncoder(config) self.post_init() + @merge_with_config_defaults @capture_outputs @auto_docstring def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: @@ -1840,7 +1841,7 @@ def get_image_features( pixel_values: torch.FloatTensor, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - patches_grid: torch.Tensor | None = None, + patches_grid: list[list[int]] | torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> "UnlimitedOcrModelOutputWithPooling": r""" @@ -1848,7 +1849,7 @@ def get_image_features( 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]`. - patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): The patches grid `(num_columns, num_rows)` per image. Required if `pixel_values_local` is passed. """ if isinstance(num_local_patches, torch.Tensor): @@ -1939,7 +1940,7 @@ def forward( pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - patches_grid: torch.Tensor | None = None, + patches_grid: list[list[int]] | torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -1952,7 +1953,7 @@ def forward( 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. - patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): The patches grid `(num_columns, num_rows)` per image. """ if inputs_embeds is None: @@ -1997,9 +1998,6 @@ def __init__(self, config: UnlimitedOcrConfig): 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 - @can_return_tuple @auto_docstring def get_image_features( @@ -2007,7 +2005,7 @@ def get_image_features( pixel_values: torch.FloatTensor, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - patches_grid: torch.Tensor | None = None, + patches_grid: list[list[int]] | torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: r""" @@ -2017,7 +2015,7 @@ def get_image_features( 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]`. - patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): The patches grid `(num_columns, num_rows)` per image. Required if `pixel_values_local` is passed. """ return self.model.get_image_features( @@ -2036,7 +2034,7 @@ def forward( pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - patches_grid: torch.Tensor | None = None, + patches_grid: list[list[int]] | torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -2051,7 +2049,7 @@ def forward( 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. - patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): The patches grid `(num_columns, num_rows)` per image. Example single-page OCR: @@ -2156,38 +2154,6 @@ def forward( 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, - patches_grid=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 - model_inputs["patches_grid"] = patches_grid - - return model_inputs - __all__ = [ "UnlimitedOcrForConditionalGeneration", diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 0e0f1bc2bfe4..21936cf7e484 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -50,6 +50,7 @@ is_torchdynamo_compiling, torch_int, ) +from ...utils.generic import merge_with_config_defaults from ...utils.output_capturing import capture_outputs from ..clip.configuration_clip import CLIPVisionConfig from ..clip.modeling_clip import CLIPAttention, CLIPEncoderLayer, CLIPVisionEmbeddings, CLIPVisionModel @@ -92,13 +93,13 @@ class UnlimitedOcrImageProcessorKwargs(DeepseekOcr2ImageProcessorKwargs): set to `True`. Can be overridden by the `max_patches` parameter in the `preprocess` method. tile_size (`int`, *optional*, defaults to `640`): The size of each local tile. Must match the model's query embedding size. + background_color (`list[int]`, *optional*, defaults to `[127, 127, 127]`): + The background color for padding. maximum_pad_value (`int`, *optional*, defaults to `640`): If `crop_to_patches` is `False` and `max(size.height, size.width)` is smaller than or equal to this value, the image is resized directly to a square of `max(size.height, size.width)` without preserving the aspect ratio. Otherwise, the image is resized while preserving the aspect ratio and then padded to a square with `background_color`. - background_color (`list[int]`, *optional*, defaults to `[127, 127, 127]`): - The background color for padding. """ maximum_pad_value: int @@ -109,6 +110,7 @@ class UnlimitedOcrImageProcessor(DeepseekOcr2ImageProcessor): maximum_pad_value = 640 max_patches = 32 model_input_names = ["pixel_values", "num_local_patches", "patches_grid"] + skip_tensor_conversion = ["num_local_patches", "patches_grid"] def _preprocess( self, @@ -199,6 +201,7 @@ def _preprocess( return BatchFeature( data=data, tensor_type=return_tensors, + skip_tensor_conversion=self.skip_tensor_conversion, ) @@ -212,6 +215,29 @@ class UnlimitedOcrProcessorKwargs(DeepseekOcr2ProcessorKwargs): class UnlimitedOcrProcessor(DeepseekOcr2Processor): valid_processor_kwargs = UnlimitedOcrProcessorKwargs + skip_tensor_conversion = [*ProcessorMixin.skip_tensor_conversion, "num_local_patches", "patches_grid"] + + def __init__( + self, + image_processor=None, + tokenizer=None, + chat_template=None, + patch_size=16, + downsample_ratio=4, + **kwargs, + ): + super().__init__( + image_processor, + tokenizer, + chat_template=chat_template, + patch_size=patch_size, + downsample_ratio=downsample_ratio, + **kwargs, + ) + self.detections_pattern = re.compile( + r"<\|det\|>(\S+) \[(\d+), (\d+), (\d+), (\d+)\]<\|/det\|>(.*?)(?=<\|det\|>||\Z)", + flags=re.DOTALL, + ) def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> TextInput: image_size = kwargs.get("size") or self.image_processor.size @@ -270,11 +296,7 @@ def decode(self, *args, return_detections: bool = False, **kwargs): return decoded def _parse_detections(self, decoded: str) -> list[dict]: - matches = re.findall( - r"<\|det\|>(\S+) \[(\d+), (\d+), (\d+), (\d+)\]<\|/det\|>(.*?)(?=<\|det\|>||\Z)", - decoded, - flags=re.DOTALL, - ) + matches = self.detections_pattern.findall(decoded) detections = [] for region_type, x1, y1, x2, y2, text in matches: detections.append( @@ -368,19 +390,6 @@ class UnlimitedOcrVisionConfig(DeepseekOcr2VisionConfig): "encoder_config": UnlimitedOcrVisionEncoderConfig, } - def __post_init__(self, **kwargs): - if self.sam_config is None: - self.sam_config = self.sub_configs["sam_config"]() - elif isinstance(self.sam_config, dict): - self.sam_config = self.sub_configs["sam_config"](**self.sam_config) - - if self.encoder_config is None: - self.encoder_config = self.sub_configs["encoder_config"]() - elif isinstance(self.encoder_config, dict): - self.encoder_config = self.sub_configs["encoder_config"](**self.encoder_config) - - PreTrainedConfig.__post_init__(self, **kwargs) - @auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict @@ -471,7 +480,8 @@ def __post_init__(self, **kwargs): self.vision_config = self.sub_configs["vision_config"](**self.vision_config) if self.text_config is None: - # The reference implementation defines the text config values on the main config + # For backward compatibility. The reference implementation defines the text config values + # on the main config self.text_config = self.sub_configs["text_config"](**kwargs) elif isinstance(self.text_config, dict): self.text_config = self.sub_configs["text_config"](**self.text_config) @@ -522,7 +532,7 @@ class UnlimitedOcrSamVisionEncoder(DeepseekOcr2SamVisionEncoder): class UnlimitedOcrAttention(CLIPAttention): def __init__(self, config: UnlimitedOcrVisionEncoderConfig): super().__init__(config) - # Required for repeat_kv(..., num_key_value_groups) + # Required for GQA compatibility self.num_key_value_groups = 1 @@ -587,7 +597,7 @@ def forward(self, patch_embeds: torch.Tensor) -> torch.Tensor: class_embeds = self.class_embedding.expand(batch_size, 1, -1) embeddings = torch.cat([class_embeds, patch_embeds], dim=1) - # always interpolate + # always interpolate compared to clip where this can be optional embeddings = embeddings + self.interpolate_pos_encoding( embeddings, grid_height * self.patch_size, grid_width * self.patch_size ) @@ -612,6 +622,7 @@ def __init__(self, config: UnlimitedOcrVisionEncoderConfig): super().__init__(config) del self.post_layernorm + @merge_with_config_defaults @capture_outputs @auto_docstring def forward(self, patch_embeds: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: @@ -1140,7 +1151,7 @@ def get_image_features( pixel_values: torch.FloatTensor, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - patches_grid: torch.Tensor | None = None, + patches_grid: list[list[int]] | torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> "UnlimitedOcrModelOutputWithPooling": r""" @@ -1148,7 +1159,7 @@ def get_image_features( 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]`. - patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): The patches grid `(num_columns, num_rows)` per image. Required if `pixel_values_local` is passed. """ if isinstance(num_local_patches, torch.Tensor): @@ -1214,7 +1225,7 @@ def forward( pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - patches_grid: torch.Tensor | None = None, + patches_grid: list[list[int]] | torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -1227,7 +1238,7 @@ def forward( 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. - patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): The patches grid `(num_columns, num_rows)` per image. """ if inputs_embeds is None: @@ -1272,9 +1283,6 @@ def __init__(self, config: UnlimitedOcrConfig): 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 - @can_return_tuple @auto_docstring def get_image_features( @@ -1282,7 +1290,7 @@ def get_image_features( pixel_values: torch.FloatTensor, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - patches_grid: torch.Tensor | None = None, + patches_grid: list[list[int]] | torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: r""" @@ -1292,7 +1300,7 @@ def get_image_features( 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]`. - patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): The patches grid `(num_columns, num_rows)` per image. Required if `pixel_values_local` is passed. """ return self.model.get_image_features( @@ -1311,7 +1319,7 @@ def forward( pixel_values: torch.FloatTensor | None = None, pixel_values_local: torch.FloatTensor | None = None, num_local_patches: list[int] | torch.Tensor | None = None, - patches_grid: torch.Tensor | None = None, + patches_grid: list[list[int]] | torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -1326,7 +1334,7 @@ def forward( 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. - patches_grid (`torch.Tensor` of shape `(num_images, 2)`, *optional*): + patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): The patches grid `(num_columns, num_rows)` per image. Example single-page OCR: @@ -1431,38 +1439,6 @@ def forward( 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, - patches_grid=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 - model_inputs["patches_grid"] = patches_grid - - return model_inputs - __all__ = [ "UnlimitedOcrConfig", diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index 8852c265f457..d826b3b002b3 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -37,6 +37,7 @@ class UnlimitedOcrProcessorKwargs(ProcessingKwargs, total=False): @auto_docstring class UnlimitedOcrProcessor(ProcessorMixin): valid_processor_kwargs = UnlimitedOcrProcessorKwargs + skip_tensor_conversion = [*ProcessorMixin.skip_tensor_conversion, "num_local_patches", "patches_grid"] def __init__( self, @@ -58,6 +59,10 @@ def __init__( self.downsample_ratio = downsample_ratio self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token) super().__init__(image_processor, tokenizer, chat_template=chat_template, **kwargs) + self.detections_pattern = re.compile( + r"<\|det\|>(\S+) \[(\d+), (\d+), (\d+), (\d+)\]<\|/det\|>(.*?)(?=<\|det\|>||\Z)", + flags=re.DOTALL, + ) def prepare_inputs_layout(self, images=None, text=None, videos=None, audio=None, **kwargs): images, text, videos, audio = super().prepare_inputs_layout( @@ -136,11 +141,7 @@ def decode(self, *args, return_detections: bool = False, **kwargs): return decoded def _parse_detections(self, decoded: str) -> list[dict]: - matches = re.findall( - r"<\|det\|>(\S+) \[(\d+), (\d+), (\d+), (\d+)\]<\|/det\|>(.*?)(?=<\|det\|>||\Z)", - decoded, - flags=re.DOTALL, - ) + matches = self.detections_pattern.findall(decoded) detections = [] for region_type, x1, y1, x2, y2, text in matches: detections.append( diff --git a/tests/generation/test_logits_process.py b/tests/generation/test_logits_process.py index 12ea0f8b880f..19336ff19d24 100644 --- a/tests/generation/test_logits_process.py +++ b/tests/generation/test_logits_process.py @@ -652,6 +652,31 @@ def test_no_repeat_ngram_dist_processor(self): self.assertFalse(torch.all(scores == filtered_scores_2_gram)) self.assertFalse(torch.all(scores == filtered_scores_3_gram)) + def test_no_repeat_ngram_dist_processor_banned_token_reappears(self): + vocab_size = 10 + batch_size = 1 + + # The 3-gram (1, 2, 5) repeats the current (1, 2) suffix and bans 5. The later 3-gram (7, 8, 5) ends in the + # same token without repeating the suffix, and must not undo that ban. + input_ids = torch.tensor([[1, 2, 5, 7, 8, 5, 1, 2]], device=torch_device, dtype=torch.long) + scores = self._get_uniform_logits(batch_size, vocab_size) + + filtered_scores = NoRepeatNGramLogitsProcessor(3)(input_ids, scores) + + self.assertListEqual(torch.isinf(filtered_scores).nonzero()[:, 1].tolist(), [5]) + + def test_no_repeat_ngram_dist_processor_sequence_shorter_than_ngram(self): + vocab_size = 3 + batch_size = 2 + + # The sequences are one token short of holding a 3-gram, so nothing can be banned yet + input_ids = torch.tensor([[1, 2], [0, 1]], device=torch_device, dtype=torch.long) + scores = self._get_uniform_logits(batch_size, vocab_size) + + filtered_scores = NoRepeatNGramLogitsProcessor(3)(input_ids, scores) + + self.assertFalse(torch.isinf(filtered_scores).any()) + def test_sliding_window_no_repeat_ngram_dist_processor(self): vocab_size = 3 # The (0, 1) bigram appears at the start, so a full-sequence processor would forbid token 1 diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index a2a7f034c6ef..2f1734281cc0 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -254,7 +254,7 @@ def _check_manual_forward_cache(self, cache_implementation: str): def test_manual_forward_dynamic_cache(self): self._check_manual_forward_cache(cache_implementation="dynamic") - def test_manual_forward_loop_static_cache(self): + def test_manual_forward_static_cache(self): self._check_manual_forward_cache(cache_implementation="static") def test_generate_cache_chunked_prefill(self): diff --git a/tests/test_processing_common.py b/tests/test_processing_common.py index 8e756923c5e4..9615abb80c4c 100644 --- a/tests/test_processing_common.py +++ b/tests/test_processing_common.py @@ -1120,7 +1120,10 @@ def test_flat_kwarg_applied_when_modality_dict_lacks_it(self): inputs = processor(**inputs_dict, **extra_kwargs, text_kwargs={}, return_tensors="pt") for k, v in inputs.items(): - self.assertIsInstance(v, torch.Tensor, msg=f"{k} should be a torch.Tensor") + if k in processor.skip_tensor_conversion: + self.assertNotIsInstance(v, torch.Tensor, msg=f"{k} should not be a torch.Tensor") + else: + self.assertIsInstance(v, torch.Tensor, msg=f"{k} should be a torch.Tensor") def test_args_overlap_kwargs(self): if "image_processor" not in self.processor_class.get_attributes(): @@ -1724,7 +1727,10 @@ def _test_apply_chat_template( return_tensor_to_type = {"pt": torch.Tensor, "np": np.ndarray, None: list} for k in out_dict: - self.assertIsInstance(out_dict[k], return_tensor_to_type[return_tensors]) + if k in processor.skip_tensor_conversion: + self.assertNotIsInstance(out_dict[k], (torch.Tensor, np.ndarray)) + else: + self.assertIsInstance(out_dict[k], return_tensor_to_type[return_tensors]) # Test continue from final message assistant_message = { From 1ab76a93d1ad3a2389d556e075f467ebddc64f3f Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:11:05 +0200 Subject: [PATCH 190/195] remove tensor option --- .../unlimited_ocr/modeling_unlimited_ocr.py | 37 ++++++++-------- .../unlimited_ocr/modular_unlimited_ocr.py | 42 +++++++++---------- .../unlimited_ocr/processing_unlimited_ocr.py | 5 +-- 3 files changed, 38 insertions(+), 46 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 092edbcfaa81..f72aa08a0ae6 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1840,21 +1840,18 @@ 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, - patches_grid: list[list[int]] | torch.Tensor | None = None, + num_local_patches: list[int] | None = None, + patches_grid: list[list[int]] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> "UnlimitedOcrModelOutputWithPooling": 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*): + num_local_patches (`list[int]`, *optional*): Number of local patches per image, e.g. `[6, 0, 4]`. - patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): + patches_grid (`list[list[int]]`, *optional*): The patches grid `(num_columns, num_rows)` per image. Required if `pixel_values_local` is passed. """ - if isinstance(num_local_patches, torch.Tensor): - num_local_patches = num_local_patches.tolist() - batch_size = pixel_values.shape[0] pixel_values = pixel_values.to(self.vision_tower.dtype) @@ -1888,7 +1885,7 @@ def get_image_features( local_features = per_image_local[idx] if local_features is not None and local_features.shape[0] > 0: - num_columns, num_rows = int(patches_grid[idx][0]), int(patches_grid[idx][1]) + num_columns, num_rows = patches_grid[idx] num_queries_local = int(local_features.shape[1] ** 0.5) local_grid_shape = (num_rows * num_queries_local, -1, hidden_size) local_grid = local_features.reshape( @@ -1939,8 +1936,8 @@ def forward( 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, - patches_grid: list[list[int]] | torch.Tensor | None = None, + num_local_patches: list[int] | None = None, + patches_grid: list[list[int]] | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -1951,9 +1948,9 @@ def forward( 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*): + num_local_patches (`list[int]`, *optional*): Number of local patches per image in the batch. - patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): + patches_grid (`list[list[int]]`, *optional*): The patches grid `(num_columns, num_rows)` per image. """ if inputs_embeds is None: @@ -2004,8 +2001,8 @@ 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, - patches_grid: list[list[int]] | torch.Tensor | None = None, + num_local_patches: list[int] | None = None, + patches_grid: list[list[int]] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: r""" @@ -2013,9 +2010,9 @@ def get_image_features( 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*): + num_local_patches (`list[int]`, *optional*): Number of local patches per image, e.g. `[6, 0, 4]`. - patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): + patches_grid (`list[list[int]]`, *optional*): The patches grid `(num_columns, num_rows)` per image. Required if `pixel_values_local` is passed. """ return self.model.get_image_features( @@ -2033,8 +2030,8 @@ def forward( 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, - patches_grid: list[list[int]] | torch.Tensor | None = None, + num_local_patches: list[int] | None = None, + patches_grid: list[list[int]] | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -2047,9 +2044,9 @@ def forward( 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*): + num_local_patches (`list[int]`, *optional*): Number of local patches per image in the batch. - patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): + patches_grid (`list[list[int]]`, *optional*): The patches grid `(num_columns, num_rows)` per image. Example single-page OCR: diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 21936cf7e484..e96d02124eea 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -247,10 +247,9 @@ def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> T num_queries_global = math.ceil(size // self.patch_size / self.downsample_ratio) num_queries_local = math.ceil(tile_size // self.patch_size / self.downsample_ratio) - num_columns = int(image_inputs["patches_grid"][image_idx][0]) - num_rows = int(image_inputs["patches_grid"][image_idx][1]) + num_columns, num_rows = image_inputs["patches_grid"][image_idx] num_tokens = num_queries_global * (num_queries_global + 1) + 1 - if int(image_inputs["num_local_patches"][image_idx]) > 0: + if image_inputs["num_local_patches"][image_idx] > 0: num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) return self.image_token * num_tokens @@ -1150,21 +1149,18 @@ 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, - patches_grid: list[list[int]] | torch.Tensor | None = None, + num_local_patches: list[int] | None = None, + patches_grid: list[list[int]] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> "UnlimitedOcrModelOutputWithPooling": 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*): + num_local_patches (`list[int]`, *optional*): Number of local patches per image, e.g. `[6, 0, 4]`. - patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): + patches_grid (`list[list[int]]`, *optional*): The patches grid `(num_columns, num_rows)` per image. Required if `pixel_values_local` is passed. """ - if isinstance(num_local_patches, torch.Tensor): - num_local_patches = num_local_patches.tolist() - batch_size = pixel_values.shape[0] pixel_values = pixel_values.to(self.vision_tower.dtype) @@ -1198,7 +1194,7 @@ def get_image_features( local_features = per_image_local[idx] if local_features is not None and local_features.shape[0] > 0: - num_columns, num_rows = int(patches_grid[idx][0]), int(patches_grid[idx][1]) + num_columns, num_rows = patches_grid[idx] num_queries_local = int(local_features.shape[1] ** 0.5) local_grid_shape = (num_rows * num_queries_local, -1, hidden_size) local_grid = local_features.reshape( @@ -1224,8 +1220,8 @@ def forward( 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, - patches_grid: list[list[int]] | torch.Tensor | None = None, + num_local_patches: list[int] | None = None, + patches_grid: list[list[int]] | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -1236,9 +1232,9 @@ def forward( 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*): + num_local_patches (`list[int]`, *optional*): Number of local patches per image in the batch. - patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): + patches_grid (`list[list[int]]`, *optional*): The patches grid `(num_columns, num_rows)` per image. """ if inputs_embeds is None: @@ -1289,8 +1285,8 @@ 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, - patches_grid: list[list[int]] | torch.Tensor | None = None, + num_local_patches: list[int] | None = None, + patches_grid: list[list[int]] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: r""" @@ -1298,9 +1294,9 @@ def get_image_features( 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*): + num_local_patches (`list[int]`, *optional*): Number of local patches per image, e.g. `[6, 0, 4]`. - patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): + patches_grid (`list[list[int]]`, *optional*): The patches grid `(num_columns, num_rows)` per image. Required if `pixel_values_local` is passed. """ return self.model.get_image_features( @@ -1318,8 +1314,8 @@ def forward( 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, - patches_grid: list[list[int]] | torch.Tensor | None = None, + num_local_patches: list[int] | None = None, + patches_grid: list[list[int]] | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, @@ -1332,9 +1328,9 @@ def forward( 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*): + num_local_patches (`list[int]`, *optional*): Number of local patches per image in the batch. - patches_grid (`list[list[int]]` or `torch.Tensor`, *optional*): + patches_grid (`list[list[int]]`, *optional*): The patches grid `(num_columns, num_rows)` per image. Example single-page OCR: diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index d826b3b002b3..34d98983e1b6 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -92,10 +92,9 @@ def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> T num_queries_global = math.ceil(size // self.patch_size / self.downsample_ratio) num_queries_local = math.ceil(tile_size // self.patch_size / self.downsample_ratio) - num_columns = int(image_inputs["patches_grid"][image_idx][0]) - num_rows = int(image_inputs["patches_grid"][image_idx][1]) + num_columns, num_rows = image_inputs["patches_grid"][image_idx] num_tokens = num_queries_global * (num_queries_global + 1) + 1 - if int(image_inputs["num_local_patches"][image_idx]) > 0: + if image_inputs["num_local_patches"][image_idx] > 0: num_tokens += (num_rows * num_queries_local) * (num_columns * num_queries_local + 1) return self.image_token * num_tokens From f81a69d53e4b55921158d357874fc47ba2b1dd6b Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 25 Aug 2026 18:24:04 +0200 Subject: [PATCH 191/195] skip export tests for now --- tests/exporters/test_export.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/exporters/test_export.py b/tests/exporters/test_export.py index 5d074da906f7..1726be0fb3a0 100644 --- a/tests/exporters/test_export.py +++ b/tests/exporters/test_export.py @@ -119,6 +119,14 @@ "exported prefill returns only `logits` while eager surfaces the populated KV cache. " "Same shape as Voxtral. TODO: align the generate-decomposition path." ), + "UnlimitedOcrForConditionalGeneration": ( + "The reference sliding window cache writes into `prefill_keys`/`keys`, which are views of a " + "single contiguous buffer. The test exports models with grad enabled, so the first in-place write " + "of grad-carrying states turns that buffer into a leaf requiring grad and the next write " + "raises `a leaf Variable that requires grad is being used in an in-place operation`. Exporting in " + "a `torch.no_grad` context works correctly. Alternatively, kv values have to be detached before " + "updating the cache." + ), }, # Every backend, dynamic-shape only. "dynamic": { @@ -163,6 +171,11 @@ ), "PixioModel": ("Lowering exceeds the 10-minute test timeout."), "PixioBackbone": "Same `timeout` failure as `PixioModel`.", + "UnlimitedOcrModel": ( + "The vision component hits `DispatchError: No ONNX function found for aten.mul.Scalar` " + "(no decompositions registered for the real-valued input). Torch export works." + ), + "UnlimitedOcrForConditionalGeneration": "Same `aten.mul.Scalar` dispatch failure as `UnlimitedOcrModel`.", }, # ONNX, generate path only. "onnx.generate": { From 21f6dbf099f71424af2e6807fc88ae20c39523af Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Mon, 14 Sep 2026 11:41:58 +0100 Subject: [PATCH 192/195] use config default --- docs/source/en/model_doc/unlimited_ocr.md | 27 ++++--------------- .../unlimited_ocr/modeling_unlimited_ocr.py | 2 +- 2 files changed, 6 insertions(+), 23 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index 8cd1d734bd2b..eca7821d9b6a 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -16,7 +16,7 @@ limitations under the License. ⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer. --> -*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-08-24.* +*This model was published in HF papers on 2026-06-23 and contributed to Hugging Face Transformers on 2026-09-14.* # UnlimitedOcr @@ -63,12 +63,7 @@ inputs = processor.apply_chat_template( messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" ).to(model.device) -output = model.generate( - **inputs, - max_new_tokens=512, - no_repeat_ngram_size=35, - no_repeat_ngram_window_size=128, -) +output = model.generate(**inputs, max_new_tokens=512) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) # image [383, 87, 497, 171]\ntext [333, 201, 558, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... @@ -105,12 +100,7 @@ inputs = processor.apply_chat_template( processor_kwargs={"padding": True}, ).to(model.device) -output = model.generate( - **inputs, - max_new_tokens=32768, - no_repeat_ngram_size=35, - no_repeat_ngram_window_size=128, -) +output = model.generate(**inputs) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) # image [383, 87, 497, 171]\ntext [333, 201, 558, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... @@ -138,12 +128,7 @@ inputs = processor.apply_chat_template( messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" ).to(model.device) -output = model.generate( - **inputs, - max_new_tokens=32768, - no_repeat_ngram_size=35, - no_repeat_ngram_window_size=128, -) +output = model.generate(**inputs) decoded, detections = processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=False, return_detections=True) # <|det|>image [383, 87, 497, 171]<|/det|>\n<|det|>text [333, 201, 558, 230]<|/det|>R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... @@ -204,9 +189,7 @@ inputs = processor.apply_chat_template( output = model.generate( **inputs, - max_new_tokens=32768, - no_repeat_ngram_size=35, - no_repeat_ngram_window_size=1024, + no_repeat_ngram_window_size=1024, # default: 128, larger window size recommended for long documents ) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) # image [382, 87, 489, 174]\ntitle [333, 201, 556, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SCLUTION FORM... diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index f72aa08a0ae6..90bb93bd8a56 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -1116,7 +1116,7 @@ def forward( ) -> 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 = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts + 1) expert_mask = expert_mask.permute(2, 1, 0) expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() From 9a0e9a7c952297ab6a81ab6b13a1d74f6e3ef262 Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Mon, 14 Sep 2026 15:50:41 +0100 Subject: [PATCH 193/195] re anton --- docs/source/en/model_doc/unlimited_ocr.md | 15 +++++-- .../configuration_unlimited_ocr.py | 2 + .../unlimited_ocr/generation_unlimited_ocr.py | 9 +++- .../unlimited_ocr/modeling_unlimited_ocr.py | 36 +++++++--------- .../unlimited_ocr/modular_unlimited_ocr.py | 43 +++++++++---------- .../unlimited_ocr/processing_unlimited_ocr.py | 2 +- .../test_image_processing_unlimited_ocr.py | 24 ++--------- .../test_modeling_unlimited_ocr.py | 20 ++++----- .../test_processing_unlimited_ocr.py | 6 +-- 9 files changed, 72 insertions(+), 85 deletions(-) diff --git a/docs/source/en/model_doc/unlimited_ocr.md b/docs/source/en/model_doc/unlimited_ocr.md index eca7821d9b6a..a514be2e6ce7 100644 --- a/docs/source/en/model_doc/unlimited_ocr.md +++ b/docs/source/en/model_doc/unlimited_ocr.md @@ -42,6 +42,13 @@ The text model is identical to DeepSeek-OCR-2 with the additional Reference Slid This model was contributed by [guarin](https://huggingface.co/guarin). The original code can be found [here](https://github.com/baidu/Unlimited-OCR). +> [!TIP] +> The original implementation runs the model with [torch.autocast](https://pytorch.org/docs/stable/amp.html#torch.autocast) in bfloat16. Wrap the forward and generate calls in an autocast context to reproduce its outputs. +> +> ```python +> with torch.autocast(device_type=model.device.type, dtype=torch.bfloat16): +> output = model.generate(**inputs) +> ``` @@ -61,7 +68,7 @@ messages = [ ] inputs = processor.apply_chat_template( messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" -).to(model.device) +).to(model.device, model.dtype) output = model.generate(**inputs, max_new_tokens=512) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) @@ -98,7 +105,7 @@ inputs = processor.apply_chat_template( return_dict=True, return_tensors="pt", processor_kwargs={"padding": True}, -).to(model.device) +).to(model.device, model.dtype) output = model.generate(**inputs) processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) @@ -126,7 +133,7 @@ messages = [ ] inputs = processor.apply_chat_template( messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" -).to(model.device) +).to(model.device, model.dtype) output = model.generate(**inputs) decoded, detections = processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=False, return_detections=True) @@ -185,7 +192,7 @@ inputs = processor.apply_chat_template( return_dict=True, return_tensors="pt", processor_kwargs={"crop_to_patches": False}, -).to(model.device) +).to(model.device, model.dtype) output = model.generate( **inputs, diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index b2cfe615153a..990863ccce52 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -271,6 +271,8 @@ def validate_architecture(self): f"The hidden size ({self.hidden_size}) is not a multiple of the number of attention " f"heads ({self.num_attention_heads})." ) + if self.sliding_window is None: + raise ValueError("`sliding_window` cannot be `None` for this model.") @auto_docstring(checkpoint="baidu/Unlimited-OCR") diff --git a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py index 1ac4a797c45f..5251995a89b8 100644 --- a/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/generation_unlimited_ocr.py @@ -50,8 +50,13 @@ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> to class UnlimitedOcrGenerationMixin(GenerationMixin): - r""" - Adds support for the `no_repeat_ngram_window_size` generation option. If set together with `no_repeat_ngram_size`, + r"""Adds support for reference sliding window attention (R-SWA) cache prefill handling and the corresponding + sliding window logits processor. + + The cache combines a fixed prefill with a sliding window cache. The mixin tracks when prefill is completed + and marks that region of the cache as fixed. + + Adds the `no_repeat_ngram_window_size` generation option. If set together with `no_repeat_ngram_size`, n-gram repetitions are blocked only within this many trailing tokens instead of over the whole sequence. """ diff --git a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py index 90bb93bd8a56..92719eb60f28 100644 --- a/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modeling_unlimited_ocr.py @@ -941,7 +941,7 @@ def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwarg >>> image_processor = AutoImageProcessor.from_pretrained(model_id) >>> image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg") - >>> inputs = image_processor(images=image, return_tensors="pt").to(model.device) + >>> inputs = image_processor(images=image, return_tensors="pt").to(model.device, model.dtype) >>> with torch.no_grad(): ... outputs = model(**inputs) @@ -1683,6 +1683,7 @@ def create_reference_sliding_window_causal_mask( and_mask_function: Callable | None = None, block_sequence_ids: torch.Tensor | None = None, layer_idx: int | None = None, + allow_is_causal_skip: bool = False, ) -> torch.Tensor | BlockMask | None: layer = None if past_key_values is not None: @@ -1703,6 +1704,9 @@ def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): return kv_idx - kv_offset < prefill_length prefill_mask_function = and_masks(causal_mask_function, prefill_overlay) + + # Workaround in case an or_mask_function is passed as prefill_mask_function + # is also an or mask. if or_mask_function is not None: prefill_mask_function = or_masks(prefill_mask_function, or_mask_function) @@ -1716,6 +1720,7 @@ def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): and_mask_function=and_mask_function, block_sequence_ids=block_sequence_ids, layer_idx=layer_idx, + allow_is_causal_skip=allow_is_causal_skip, ) @@ -1777,13 +1782,8 @@ def forward( causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), + "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), } - # CODEPATH: create_sliding_window_causal_mask errors if config.sliding_window is None. - # Without the check here a model with full attention crashes. - if self.config.sliding_window is not None: - causal_mask_mapping["reference_sliding_attention"] = create_reference_sliding_window_causal_mask( - **mask_kwargs - ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) @@ -1854,13 +1854,11 @@ def get_image_features( """ batch_size = pixel_values.shape[0] - pixel_values = pixel_values.to(self.vision_tower.dtype) 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: - pixel_values_local = pixel_values_local.to(self.vision_tower.dtype) 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) @@ -1874,10 +1872,11 @@ def get_image_features( hidden_size = global_features.shape[-1] newline = self.image_newline[None, None, :] - view_separator = self.view_separator[None, :] + num_queries_global = int(global_features.shape[1] ** 0.5) + # Different from DeepseekOcr2: each row of image patches gets a newline token appended. all_features = [] - num_queries_global = int(global_features.shape[1] ** 0.5) + view_separator = self.view_separator.to(global_features.device).unsqueeze(0) for idx in range(batch_size): global_grid = global_features[idx].reshape(num_queries_global, num_queries_global, hidden_size) global_grid = torch.cat([global_grid, newline.expand(num_queries_global, 1, hidden_size)], dim=1) @@ -1985,6 +1984,8 @@ def forward( ) +# Cannot inherit from XYZForConditionalGeneration as this results in MRO clashes between GenerationMixin +# and UnlimitedOcrGenerationMixin. Cannot change the inheritance order as this breaks modular generation. @auto_docstring class UnlimitedOcrForConditionalGeneration(UnlimitedOcrPreTrainedModel, UnlimitedOcrGenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} @@ -2066,14 +2067,9 @@ def forward( ... ] >>> inputs = processor.apply_chat_template( ... messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" - ... ).to(model.device) + ... ).to(model.device, model.dtype) - >>> output = model.generate( - ... **inputs, - ... max_new_tokens=32768, - ... no_repeat_ngram_size=35, - ... no_repeat_ngram_window_size=128, - ... ) + >>> output = model.generate(**inputs) >>> processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) >>> # image [383, 87, 497, 171]\ntext [333, 201, 558, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... ``` @@ -2103,12 +2099,10 @@ def forward( ... return_dict=True, ... return_tensors="pt", ... processor_kwargs={"crop_to_patches": False}, - ... ).to(model.device) + ... ).to(model.device, model.dtype) >>> output = model.generate( ... **inputs, - ... max_new_tokens=32768, - ... no_repeat_ngram_size=35, ... no_repeat_ngram_window_size=1024, ... ) >>> processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index e96d02124eea..003e3b1eb1e1 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -215,7 +215,7 @@ class UnlimitedOcrProcessorKwargs(DeepseekOcr2ProcessorKwargs): class UnlimitedOcrProcessor(DeepseekOcr2Processor): valid_processor_kwargs = UnlimitedOcrProcessorKwargs - skip_tensor_conversion = [*ProcessorMixin.skip_tensor_conversion, "num_local_patches", "patches_grid"] + skip_tensor_conversion = ["video_metadata", "text_replacement_offsets", "num_local_patches", "patches_grid"] def __init__( self, @@ -448,6 +448,11 @@ def __post_init__(self, **kwargs): ] super().__post_init__(**kwargs) + def validate_architecture(self): + super().validate_architecture() + if self.sliding_window is None: + raise ValueError("`sliding_window` cannot be `None` for this model.") + @auto_docstring(checkpoint="baidu/Unlimited-OCR") @strict @@ -657,7 +662,7 @@ def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwarg >>> image_processor = AutoImageProcessor.from_pretrained(model_id) >>> image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/ocr_suggestion_form.jpg") - >>> inputs = image_processor(images=image, return_tensors="pt").to(model.device) + >>> inputs = image_processor(images=image, return_tensors="pt").to(model.device, model.dtype) >>> with torch.no_grad(): ... outputs = model(**inputs) @@ -1026,6 +1031,7 @@ def create_reference_sliding_window_causal_mask( and_mask_function: Callable | None = None, block_sequence_ids: torch.Tensor | None = None, layer_idx: int | None = None, + allow_is_causal_skip: bool = False, ) -> torch.Tensor | BlockMask | None: layer = None if past_key_values is not None: @@ -1046,6 +1052,9 @@ def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): return kv_idx - kv_offset < prefill_length prefill_mask_function = and_masks(causal_mask_function, prefill_overlay) + + # Workaround in case an or_mask_function is passed as prefill_mask_function + # is also an or mask. if or_mask_function is not None: prefill_mask_function = or_masks(prefill_mask_function, or_mask_function) @@ -1059,6 +1068,7 @@ def prefill_overlay(batch_idx, head_idx, q_idx, kv_idx): and_mask_function=and_mask_function, block_sequence_ids=block_sequence_ids, layer_idx=layer_idx, + allow_is_causal_skip=allow_is_causal_skip, ) @@ -1099,13 +1109,8 @@ def forward( causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), + "reference_sliding_attention": create_reference_sliding_window_causal_mask(**mask_kwargs), } - # CODEPATH: create_sliding_window_causal_mask errors if config.sliding_window is None. - # Without the check here a model with full attention crashes. - if self.config.sliding_window is not None: - causal_mask_mapping["reference_sliding_attention"] = create_reference_sliding_window_causal_mask( - **mask_kwargs - ) hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) @@ -1163,13 +1168,11 @@ def get_image_features( """ batch_size = pixel_values.shape[0] - pixel_values = pixel_values.to(self.vision_tower.dtype) 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: - pixel_values_local = pixel_values_local.to(self.vision_tower.dtype) 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) @@ -1183,10 +1186,11 @@ def get_image_features( hidden_size = global_features.shape[-1] newline = self.image_newline[None, None, :] - view_separator = self.view_separator[None, :] + num_queries_global = int(global_features.shape[1] ** 0.5) + # Different from DeepseekOcr2: each row of image patches gets a newline token appended. all_features = [] - num_queries_global = int(global_features.shape[1] ** 0.5) + view_separator = self.view_separator.to(global_features.device).unsqueeze(0) for idx in range(batch_size): global_grid = global_features[idx].reshape(num_queries_global, num_queries_global, hidden_size) global_grid = torch.cat([global_grid, newline.expand(num_queries_global, 1, hidden_size)], dim=1) @@ -1269,6 +1273,8 @@ def forward( ) +# Cannot inherit from XYZForConditionalGeneration as this results in MRO clashes between GenerationMixin +# and UnlimitedOcrGenerationMixin. Cannot change the inheritance order as this breaks modular generation. @auto_docstring class UnlimitedOcrForConditionalGeneration(UnlimitedOcrPreTrainedModel, UnlimitedOcrGenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} @@ -1350,14 +1356,9 @@ def forward( ... ] >>> inputs = processor.apply_chat_template( ... messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" - ... ).to(model.device) + ... ).to(model.device, model.dtype) - >>> output = model.generate( - ... **inputs, - ... max_new_tokens=32768, - ... no_repeat_ngram_size=35, - ... no_repeat_ngram_window_size=128, - ... ) + >>> output = model.generate(**inputs) >>> processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) >>> # image [383, 87, 497, 171]\ntext [333, 201, 558, 230]R&D QUALITY IMPROVEMENT\nSUGGESTION/SOLUTION FORM... ``` @@ -1387,12 +1388,10 @@ def forward( ... return_dict=True, ... return_tensors="pt", ... processor_kwargs={"crop_to_patches": False}, - ... ).to(model.device) + ... ).to(model.device, model.dtype) >>> output = model.generate( ... **inputs, - ... max_new_tokens=32768, - ... no_repeat_ngram_size=35, ... no_repeat_ngram_window_size=1024, ... ) >>> processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) diff --git a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py index 34d98983e1b6..51b980eedf94 100644 --- a/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/processing_unlimited_ocr.py @@ -37,7 +37,7 @@ class UnlimitedOcrProcessorKwargs(ProcessingKwargs, total=False): @auto_docstring class UnlimitedOcrProcessor(ProcessorMixin): valid_processor_kwargs = UnlimitedOcrProcessorKwargs - skip_tensor_conversion = [*ProcessorMixin.skip_tensor_conversion, "num_local_patches", "patches_grid"] + skip_tensor_conversion = ["video_metadata", "text_replacement_offsets", "num_local_patches", "patches_grid"] def __init__( self, diff --git a/tests/models/unlimited_ocr/test_image_processing_unlimited_ocr.py b/tests/models/unlimited_ocr/test_image_processing_unlimited_ocr.py index 0a906c1bb3d1..15a368b4bf2a 100644 --- a/tests/models/unlimited_ocr/test_image_processing_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_image_processing_unlimited_ocr.py @@ -14,17 +14,12 @@ import unittest -from transformers import is_vision_available from transformers.testing_utils import require_torch, require_vision -from ...test_image_processing_common import ImageProcessingTestMixin, prepare_image_inputs +from ...test_image_processing_common import ImageProcessingTester, ImageProcessingTestMixin -if is_vision_available(): - pass - - -class UnlimitedOcrImageProcessingTester: +class UnlimitedOcrImageProcessingTester(ImageProcessingTester): def __init__( self, parent, @@ -41,6 +36,7 @@ def __init__( image_std=[0.5, 0.5, 0.5], do_convert_rgb=True, ): + super().__init__() size = size if size is not None else {"height": 512, "width": 512} self.parent = parent self.batch_size = batch_size @@ -67,20 +63,6 @@ def prepare_image_processor_dict(self): "do_convert_rgb": self.do_convert_rgb, } - def expected_output_image_shape(self, images): - return self.num_channels, self.size["height"], self.size["width"] - - def prepare_image_inputs(self, equal_resolution=False, numpify=False, torchify=False): - return prepare_image_inputs( - batch_size=self.batch_size, - num_channels=self.num_channels, - min_resolution=self.min_resolution, - max_resolution=self.max_resolution, - equal_resolution=equal_resolution, - numpify=numpify, - torchify=torchify, - ) - @require_torch @require_vision diff --git a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py index 2f1734281cc0..ac619ca3e2ed 100644 --- a/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_modeling_unlimited_ocr.py @@ -282,17 +282,6 @@ def test_generate_cache_chunked_prefill(self): for layer in reference_layers: self.assertEqual(layer.prefill_length, prompt_length) - def test_generate_without_sliding_window(self): - """With `sliding_window=None` every layer is a full attention layer.""" - model_tester = self.model_tester_class(self, sliding_window=None) - config, inputs_dict = model_tester.prepare_config_and_inputs_for_common() - self.assertEqual(config.text_config.layer_types, ["full_attention"] * config.text_config.num_hidden_layers) - - for model_class in self.all_generative_model_classes: - model = model_class(config).to(torch_device).eval() - out = model.generate(**inputs_dict, max_new_tokens=3, do_sample=False, return_dict_in_generate=True) - self.assertTrue(all(not layer.is_sliding for layer in out.past_key_values.layers)) - @require_torch class UnlimitedOcrIntegrationTest(unittest.TestCase): @@ -331,6 +320,7 @@ def test_small_model_integration_test_document_parsing(self): inputs = self.processor.apply_chat_template( messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt" ).to(model.device) + # The reference implementation uses autocast: https://huggingface.co/baidu/Unlimited-OCR/blob/main/modeling_unlimitedocr.py#L1042 with torch.autocast(device_type=torch_device, dtype=torch.bfloat16): generate_ids = model.generate( **inputs, @@ -341,6 +331,7 @@ def test_small_model_integration_test_document_parsing(self): EXPECTED_DECODED_TEXT = Expectations( { ("cuda", 9): "image [383, 88, 497, 175]\ntitle [333", + ("cuda", 8): "image [383, 88, 497, 171]\ntitle [333", ("cpu", None): "image [383, 87, 497, 171]\ntitle [333", } ).get_expectation() # fmt: skip @@ -373,6 +364,7 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se EXPECTED_DECODED_TEXT = Expectations( { ("cuda", 9): "<|det|>image [383, 88, 497, 175]<|/det|>\n<|det|>title [333", + ("cuda", 8): "<|det|>image [383, 88, 497, 171]<|/det|>\n<|det|>title [333", ("cpu", None): "<|det|>image [383, 87, 497, 171]<|/det|>\n<|det|>title [333", } ).get_expectation() # fmt: skip @@ -381,6 +373,7 @@ def test_small_model_integration_test_document_parsing_no_skip_special_tokens(se EXPECTED_DETECTIONS = Expectations( { ("cuda", 9): [{"region_type": "image", "box": [383, 88, 497, 175], "text": "\n"}], + ("cuda", 8): [{"region_type": "image", "box": [383, 88, 497, 171], "text": "\n"}], ("cpu", None): [{"region_type": "image", "box": [383, 87, 497, 171], "text": "\n"}], } ).get_expectation() # fmt: skip @@ -427,6 +420,10 @@ def test_small_model_integration_test_document_parsing_batched(self): "image [383, 88, 497, 174]\ntitle [333", "header [53, 23, 365, 41]Advanced Template and Styl", ], + ("cuda", 8): [ + "image [383, 87, 497, 171]\ntitle [333", + "header [53, 23, 365, 41]Advanced Template and Styl", + ], ("cpu", None): [ "image [383, 88, 497, 171]\ntitle [333", "header [53, 23, 365, 41]Advanced Template and Styl", @@ -472,6 +469,7 @@ def test_small_model_integration_test_multi_page_document_parsing(self): EXPECTED_DECODED_TEXT = Expectations( { ("cuda", 9): "image [382, 87, 489, 180]\n", + ("cuda", 8): "image [382, 87, 489, 174]\n", ("cpu", None): "image [382, 87, 489, 174]\n", } ).get_expectation() diff --git a/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py b/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py index d95d43b0fcaa..190dadcb03ef 100644 --- a/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py +++ b/tests/models/unlimited_ocr/test_processing_unlimited_ocr.py @@ -33,9 +33,9 @@ class UnlimitedOcrProcessorTest(ProcessorTesterMixin, unittest.TestCase): # Defaults from mixin are too small as a single image expands to 273 image tokens # for this checkpoint (size=1024) - image_text_kwargs_max_length = 320 - image_text_kwargs_override_max_length = 310 - image_unstructured_max_length = 300 + images_text_kwargs_max_length = 320 + images_text_kwargs_override_max_length = 310 + images_unstructured_max_length = 300 def get_messages(self, image): return [ From c1cb494062b5a4e72a1c5dd8dd89f7600707c7fb Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 15 Sep 2026 11:15:20 +0100 Subject: [PATCH 194/195] change order --- tests/generation/test_utils.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py index 03fd1713ce83..91138ffb802f 100644 --- a/tests/generation/test_utils.py +++ b/tests/generation/test_utils.py @@ -2813,6 +2813,13 @@ def _get_recurrent_state_shape(self, batch_size: int, config): state_size = getattr(config, "state_size", None) return (batch_size, intermediate_size, state_size) + def _check_attention_shapes(self, layer, seq_length, k_shape, v_shape): + # Remove the seq_length dim for cross-attention cache (it changes based on the model) + keys = layer.keys if seq_length is not None else layer.keys[:, :, 0, :] + values = layer.values if seq_length is not None else layer.values[:, :, 0, :] + self.assertEqual(keys.shape, k_shape) + self.assertEqual(values.shape, v_shape) + def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_length, config): # Raise a useful error, asking to explicitly override the method if not isinstance(past_key_values, Cache): @@ -2875,13 +2882,6 @@ def check_linear_attention_shapes(layer, num_conv_states, conv_shape, recurrent_ else: self._check_attention_shapes(layer, seq_length, k_shape, v_shape) - def _check_attention_shapes(self, layer, seq_length, k_shape, v_shape): - # Remove the seq_length dim for cross-attention cache (it changes based on the model) - keys = layer.keys if seq_length is not None else layer.keys[:, :, 0, :] - values = layer.values if seq_length is not None else layer.values[:, :, 0, :] - self.assertEqual(keys.shape, k_shape) - self.assertEqual(values.shape, v_shape) - def _get_attention_shape(self, batch_size: int, seq_length: int | None, config) -> tuple[tuple[int, ...], ...]: """Returns the expected shape of the keys and values tensors. They can differs for some models like DeepSeekV2, which uses MLA.""" From 67c57c61e98d3d5d7aa43c8127c4fcf51ab630dd Mon Sep 17 00:00:00 2001 From: guarin <43336610+guarin@users.noreply.github.com> Date: Tue, 15 Sep 2026 13:52:37 +0100 Subject: [PATCH 195/195] Remove None option --- .../models/unlimited_ocr/configuration_unlimited_ocr.py | 2 +- src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py index 990863ccce52..eb3bb9a7a0c3 100644 --- a/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/configuration_unlimited_ocr.py @@ -168,7 +168,7 @@ class UnlimitedOcrTextConfig(PreTrainedConfig): Attention type for each decoder layer. Defaults to `"reference_sliding_attention"` on every layer. sliding_window (`int`, *optional*, defaults to `128`): Sliding window size for reference sliding window attention. If set, every token attends to the last - `sliding_window` and all image and prompt tokens. Set to `None` to use full attention on every layer. + `sliding_window` and all image and prompt tokens. Example: diff --git a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py index 003e3b1eb1e1..c1d6a51f1898 100644 --- a/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py +++ b/src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py @@ -404,7 +404,7 @@ class UnlimitedOcrTextConfig(DeepseekOcr2TextConfig): Attention type for each decoder layer. Defaults to `"reference_sliding_attention"` on every layer. sliding_window (`int`, *optional*, defaults to `128`): Sliding window size for reference sliding window attention. If set, every token attends to the last - `sliding_window` and all image and prompt tokens. Set to `None` to use full attention on every layer. + `sliding_window` and all image and prompt tokens. Example: