Add Unlimited OCR by guarin · Pull Request #46836 · huggingface/transformers

guarin · 2026-06-23T07:56:45Z

What does this PR do?

WIP branch to add Unlimited OCR: https://huggingface.co/baidu/Unlimited-OCR

Hub repo PR to update files: https://huggingface.co/baidu/Unlimited-OCR/discussions/13/files

Code Agent Policy

The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. We are currently bottlenecked by our ability to review and respond to them. As a result,
we ask that new users do not submit pure code agent PRs at this time.
You may use code agents in drafting or to help you diagnose issues. We'd also ask autonomous "OpenClaw"-like agents
not to open any PRs or issues for the moment.

PRs that appear to be fully agent-written will probably be closed without review, and we may block users who do this
repeatedly or maliciously.

This is a rapidly-evolving situation that's causing significant shockwaves in the open-source community. As a result,
this policy is likely to be updated regularly in the near future. For more information, please read CONTRIBUTING.md.

I confirm that this is not a pure code agent PR.

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline and the
Pull Request checks?
Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case.
Did you make sure to update the documentation with your changes according to the guidelines?
Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

🤖 mlinter review state

hash: 04267b78
commit: 54b13d55
review: Add Unlimited OCR #46836 (review)

HuggingFaceDocBuilderDev · 2026-06-23T08:06:42Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

…fsx/guarin_flueck/transformers into add-unlimited-ocr

Cyrilvallez

Nice, added a few remarks! One of the main question I have since the beginning, is how does this cache work in multi-turns? I.e. when restarting from a non-empty cache, how is the new prefill handled? Should it be appendend the the existing "prefill-part" of the cache, discarded as if it's simple decode step,...?

guarin · 2026-08-01T15:07:31Z

Nice, added a few remarks! One of the main question I have since the beginning, is how does this cache work in multi-turns? I.e. when restarting from a non-empty cache, how is the new prefill handled? Should it be appendend the the existing "prefill-part" of the cache, discarded as if it's simple decode step,...?

Assuming each turn is a request to run OCR on a new document the cache should be wiped including prefill. Then prefill filled with the image tokens and prompt from the new document. But I guess this is not a common scenario for this model. I would assume all uses are single-turn.

github-actions

Model linter — first pass

transformers-mlinter found 6 item(s) in the model files this PR touches. These are structural conventions a maintainer would otherwise flag by hand.

This is automated and advisory — it does not block merging.

rule	count	what it checks
`TRF057`	2	Public model, config, output and processor classes and their public methods must be decorated with @auto_docstring.
`TRF033`	2	Hyperparameters must be set on the config, not mutated through a set_* method.
`TRF041`	1	A config-gated branch must carry a `# CODEPATH:` note saying which checkpoints diverge.
`TRF030`	1	Reaching more than two levels into the config means the module was handed the wrong config.

github-actions · 2026-08-21T09:46:26Z

+            return -1
+        return self.prefill_length + self.sliding_window
+
+    def set_prefill_length(self, prefill_length: int) -> None:


TRF033 — TRF033: UnlimitedOcrDynamicReferenceSlidingWindowLayer.set_prefill_length mutates a hyperparameter after construction. Put the value on the config and read it where it is used.

Why this matters

A setter makes the model's behaviour depend on call order: the value is not in the config, so it is not saved, not restored by from_pretrained, and not visible to anything planning device maps or parallelism. Users then have to know to call it, and forgetting is silent.

class AcmeTriangleAttention(nn.Module): - def set_chunk_size(self, chunk_size): - self.chunk_size = chunk_size + def __init__(self, config): + super().__init__() + self.chunk_size = config.chunk_size

Suppress with # trf-ignore: TRF033 if intentional.

had to ignore, doesn't make sense here

github-actions · 2026-08-21T09:46:26Z

+    def get_seq_length(self) -> int:
+        return self.prefill_cumulative_length + super().get_seq_length()
+
+    def set_prefill_length(self, prefill_length: int) -> None:


TRF033 — TRF033: UnlimitedOcrStaticReferenceSlidingWindowLayer.set_prefill_length mutates a hyperparameter after construction. Put the value on the config and read it where it is used.

Why this matters

A setter makes the model's behaviour depend on call order: the value is not in the config, so it is not saved, not restored by from_pretrained, and not visible to anything planning device maps or parallelism. Users then have to know to call it, and forgetting is silent.

class AcmeTriangleAttention(nn.Module): - def set_chunk_size(self, chunk_size): - self.chunk_size = chunk_size + def __init__(self, config): + super().__init__() + self.chunk_size = config.chunk_size

Suppress with # trf-ignore: TRF033 if intentional.

had to ignore, doesn't make sense here

github-actions · 2026-08-21T09:46:26Z

+                "full_attention": create_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:


TRF041 — TRF041: branch on self.config.layer_types has no # CODEPATH: note. Add one naming the checkpoints that take each path, or delete the branch.

Why this matters

Every config-gated branch is a second architecture living in the same file, and the reader cannot tell from the code whether both halves are reachable. That is why reviewers ask "is this ever used?", "are they all needed?" and "why are there so many cases?" on almost every new model, and why dead experimental branches survive for releases. This rule does not forbid the branch; it borrows Rust's // SAFETY: discipline and makes the author write down the justification next to it. A branch nobody can name a checkpoint for is a branch to delete, and the note makes that obvious at review time instead of three rounds later.

+ # CODEPATH: ESMC-6B ships pre-normalised embeddings, the 300M/600M checkpoints do not. if config.use_embedding_norm: hidden_states = self.embedding_norm(hidden_states) - if config.msa_encoder_enabled: - hidden_states = self.msa_encoder(hidden_states) + # no released checkpoint sets msa_encoder_enabled -> branch removed

Suppress with # trf-ignore: TRF041 if intentional.

Changed code but I vaguely remember some issue when always creating the masking function. Rule seems very strict here.

Yeah had to revert as create_sliding_window_causal_mask raises on any config that has sliding_window=None

github-actions · 2026-08-21T09:46:26Z

+    def __init__(self, config: UnlimitedOcrConfig):
+        super().__init__(config)
+        self.multi_modal_projector = nn.Linear(
+            config.vision_config.sam_config.downsample_channels[-1] + config.vision_config.encoder_config.hidden_size,


TRF030 — TRF030: config.vision_config.encoder_config.hidden_size reaches 3 levels into the config. Pass the relevant sub-config to the module instead of walking the hierarchy.

Why this matters

A module that walks config.diffusion_config.atom_encoder_config.hidden_size is coupled to the whole config hierarchy rather than to its own slice of it, so it cannot be reused, tested or given a different sub-config. Pass the relevant sub-config down and the chain collapses to one hop.

class AcmeAtomEncoder(nn.Module): def __init__(self, config): super().__init__() - self.norm = AcmeLayerNorm(config.diffusion_config.atom_encoder_config.hidden_size) + self.norm = AcmeLayerNorm(config.hidden_size)

Suppress with # trf-ignore: TRF030 if intentional.

hmm doesn't really make sense in this case as this is the top-level model.

Yea to ignore

guarin · 2026-08-24T07:33:21Z

+    # TODO: remove revision before merge
+    revision = "refs/pr/13"


Will update this and merge the hub pr once this PR is overall approved. Same for processing todos.

vasqu

Looks pretty solid! Have a few comments but I think it's nothing crazy

vasqu · 2026-08-25T09:09:46Z

+
+output = model.generate(
+    **inputs,
+    max_new_tokens=32768,


Let's shorten it, it's only an example and it should run "quick" ish

vasqu · 2026-08-25T09:10:29Z

+    no_repeat_ngram_size=35,
+    no_repeat_ngram_window_size=128,


Potentially a generation config instead that has it? Looks like a recommended setting so could be a default imo as per config

I tried this and if I recall correctly I ran into issues with adding the non-standard no_repeat_ngram_window_size to the generation config. The value also changes depending on setting so maybe best to set it explicitly?

If it changes across settings yea but a default wouldnt hurt no? Yea you have to manually add them onto the json, there is no good api atm 😢

vasqu · 2026-08-25T09:13:24Z

        ("unispeech-sat", "UniSpeechSatConfig"),
        ("univnet", "UnivNetConfig"),
+        ("unlimited_ocr", "UnlimitedOcrConfig"),
+        ("unlimited_ocr_sam_vision_model", "UnlimitedOcrSamVisionConfig"),


Ah so there is no auto model for it, that explains things. Might be worth to add but wouldn't push for it in this PR for now

vasqu · 2026-08-25T09:16:29Z

+            # Set to None to avoid also adding the default `NoRepeatNGramLogitsProcessor`
+            generation_config.no_repeat_ngram_size = None
+
+        try:


Why do we need the try finally?

To make sure we revert the config modification which leaks outside of this method. Added a comment.

vasqu · 2026-08-25T09:17:01Z

+        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


small comment to show why this is needed

vasqu · 2026-08-25T10:28:00Z

The test has evolved quite a bit over the time can you check if this is still needed

Sadly still needed because of custom prefill+sliding window cache shape

vasqu · 2026-08-25T10:29:21Z

+    def test_manual_forward_dynamic_cache(self):
+        self._check_manual_forward_cache(cache_implementation="dynamic")
+
+    def test_manual_forward_loop_static_cache(self):


Suggested change

def test_manual_forward_loop_static_cache(self):

def test_manual_forward_static_cache(self):

vasqu · 2026-08-25T10:30:16Z

+
+    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)


Might be a bit over the top -> integration tests maybe that make sure it works properly maybe?

Would love to keep them as is. They really help quickly spotting errors when modifying the custom cache.

vasqu · 2026-08-25T10:30:55Z

+        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):


this is weird, why do we autocast?

Because the original code autocasts and otherwise the expected values do not match the original implementation: https://huggingface.co/baidu/Unlimited-OCR/blob/main/modeling_unlimitedocr.py#L1042

Lets add a small comment then please, I doubt users would do this on their own (e.g. in the docs we should note / tip about this)

Probably best to update the example scripts to use bf16 and autocast then?

Imo, we can add a note / warning that the original model prefers autocast but honestly the examples should be simple just using normal bf16 / auto dtype

vasqu · 2026-08-25T10:31:20Z

+
+        EXPECTED_DECODED_TEXT = Expectations(
+            {
+                ("cuda", 9): "<PAGE>image [382, 87, 489, 180]\n",


does it also work for a10 (run-slow)?

Will add values for a10 and check run-slow before merging. I anyway have to adjust expectation values depending on the state of #47773 when I merge this PR

guarin · 2026-08-25T15:00:42Z

    "image_sizes_videos",
    "pixel_attention_mask",
    "pixel_values_images",
-    "num_local_patches",


Duplicate (it is already further above)

guarin · 2026-08-25T15:02:25Z

    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"]


Image processors do not have skip_tensor_conversion. Only ProcessorMixin has it. Would probably be a good addition in a follow-up pr.

github-actions · 2026-08-25T16:25:20Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: auto, deepseek_ocr2, unlimited_ocr

github-actions · 2026-08-25T16:37:38Z

CI recap

Dashboard: View test results in Grafana
Latest run: 32864226814:1
Result: failure | Jobs: 14 | Tests: 182,902 | Failures: 60 | Duration: 4m 54s

vasqu

Some small comments but overall very ready to get merged soon 🤗

vasqu · 2026-08-31T17:30:04Z

+    no_repeat_ngram_size=35,
+    no_repeat_ngram_window_size=128,


vasqu · 2026-08-31T17:30:51Z

    "pixel_values_images",
-    "num_local_patches",
+    "pixel_values_local",
+    "patches_grid",


potentially different PR (to first cover the deepseek model) or is it only now with unlimited ocr?

patches_grid is unlimited ocr only, pixel_values_local is also in deepseek and step3p7 but in both forwards it is only used within if pixel_values is not None: and pixel_values is dropped so it was effectively ignored.

vasqu · 2026-08-31T17:33:51Z

+
+class UnlimitedOcrGenerationMixin(GenerationMixin):
+    r"""
+    Adds support for the `no_repeat_ngram_window_size` generation option. If set together with `no_repeat_ngram_size`,


The differences are a bit more, e.g. the prefill handling :D

vasqu · 2026-08-31T17:36:05Z

+
+class UnlimitedOcrProcessor(DeepseekOcr2Processor):
+    valid_processor_kwargs = UnlimitedOcrProcessorKwargs
+    skip_tensor_conversion = [*ProcessorMixin.skip_tensor_conversion, "num_local_patches", "patches_grid"]


Let's avoid the *ProcessorMixin.skip_tensor_conversion, it's not resolved in modular and it looks really weird in processing then

vasqu · 2026-08-31T17:39:53Z

+        del self.post_layernorm
+
+    @can_return_tuple
+    @capture_outputs


Yea you're not the first one. The decorators are currently quite hard have to agree, it's an old tale with Raushan, Pablo, and Ilyas atp

Yea, better docs could help but I feel like it's still messy even then. Anyways should be a different PR and a whole design question overall (could be an issue if you want to keep track)

vasqu · 2026-08-31T17:47:22Z

+
+
+@auto_docstring
+class UnlimitedOcrForConditionalGeneration(UnlimitedOcrPreTrainedModel, UnlimitedOcrGenerationMixin):


Let's add a comment for this tho so it's known for outsiders later on

vasqu · 2026-08-31T17:48:12Z

+
+
+@auto_docstring
+class UnlimitedOcrForConditionalGeneration(UnlimitedOcrPreTrainedModel, UnlimitedOcrGenerationMixin):


Does UnlimitedOcrForConditionalGeneration(XXXForConditionalGeneration, UnlimitedOcrGenerationMixin) maybe work? Just a last chance/try

vasqu · 2026-08-31T17:49:04Z

-                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, seq_length, k_shape, v_shape):


Let's keep it as in the same order no? I.e. where the original method was

vasqu · 2026-08-31T17:49:33Z

+    pass
+
+
+class UnlimitedOcrImageProcessingTester:


Sync with main for your own PR :p

vasqu · 2026-08-31T17:50:27Z

+        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):


Add Unlimited OCR

13b0585

guarin added 3 commits June 23, 2026 12:15

update

3d0a006

add CLIP

16a7260

add r-swa

ebf9b3b

molbap added the New model label Jun 23, 2026

guarin added 24 commits June 23, 2026 15:22

fix tests

a531dd8

update expected values

8cd0c32

Merge branch 'add-unlimited-ocr' of hpc-cluster-hopper-login-node-1:/…

83ab5d8

…fsx/guarin_flueck/transformers into add-unlimited-ocr

enable _can_compile_fullgraph

b6a226e

fix downsample outputs

e2d8d50

use temp repo for tests

e5f2c46

fix mlp_layer_types None

0e47343

Merge branch 'add-unlimited-ocr' of hpc-cluster-hopper-login-node-1:/…

a26c74a

…fsx/guarin_flueck/transformers into add-unlimited-ocr

add partial conversion mapping

38f3199

handle sparse layers

6c42faa

update conversion mapping

1263664

update num_hidden_layers

b7c0276

ignore lm_head

34abc9f

update conversion mapping

6cf169b

Merge branch 'add-unlimited-ocr' of hpc-cluster-hopper-login-node-1:/…

0366fc3

…fsx/guarin_flueck/transformers into add-unlimited-ocr

draft docs

73281a7

Merge branch 'main' into add-unlimited-ocr

ec0c401

add default config values

8ecbf2c

update text config handling

96a794a

update post init

d4aa878

update

bffafd0

fix to dtype

16dcdbb

update docs

3bba66b

Merge branch 'add-unlimited-ocr' of hpc-cluster-hopper-login-node-1:/…

c81df8b

…fsx/guarin_flueck/transformers into add-unlimited-ocr

guarin added 2 commits July 30, 2026 10:56

update tests

0ad2158

fix-repo

1cf342a

Cyrilvallez reviewed Jul 31, 2026

View reviewed changes

Comment thread src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py

Comment thread src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py Outdated

Comment thread src/transformers/models/unlimited_ocr/modular_unlimited_ocr.py

guarin added 7 commits July 31, 2026 08:32

cleanup

76076d7

match static cache impl with dynamic one

5e21344

add comment

2065439

update test utils

80d08e6

clarify comment

1a4109a

Merge branch 'main' into add-unlimited-ocr

124ddef

fix-repo

5102060

guarin added 3 commits August 5, 2026 15:08

Merge branch 'main' into add-unlimited-ocr

f127652

Merge branch 'main' into add-unlimited-ocr

1172fac

update

54b13d5

github-actions Bot reviewed Aug 21, 2026

View reviewed changes

guarin added 3 commits August 21, 2026 11:11

fixes

adb0e0e

revert

6c0b38c

cleanup docstrings

6a087ca

guarin commented Aug 24, 2026

View reviewed changes

vasqu reviewed Aug 25, 2026

View reviewed changes

guarin added 2 commits August 25, 2026 14:08

Merge branch 'main' into add-unlimited-ocr

2788fbe

fixes

fb34ec1

guarin commented Aug 25, 2026

View reviewed changes

remove tensor option

1ab76a9

zucchini-nlp mentioned this pull request Aug 25, 2026

[WIP] unlimited-ocr #48255

Closed

skip export tests for now

f81a69d

vasqu approved these changes Aug 31, 2026

View reviewed changes

	def test_manual_forward_loop_static_cache(self):
	def test_manual_forward_static_cache(self):



		@auto_docstring
		class UnlimitedOcrForConditionalGeneration(UnlimitedOcrPreTrainedModel, UnlimitedOcrGenerationMixin):

Uh oh!

Conversation

guarin commented Jun 23, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Code Agent Policy

Before submitting

Who can review?

Uh oh!

HuggingFaceDocBuilderDev commented Jun 23, 2026

Uh oh!

Cyrilvallez left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

guarin commented Aug 1, 2026

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

Model linter — first pass

Uh oh!

Uh oh!

Uh oh!

github-actions Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

guarin Aug 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

guarin commented Jun 23, 2026 •

edited by github-actions Bot

Loading

guarin Aug 21, 2026 •

edited

Loading