trim unused model param tensors

.eval_results/olmocrbench.yaml

@@ -0,0 +1,72 @@
+- dataset:
+    id: allenai/olmOCR-bench
+    task_id: overall
+  value: 83.4
+  source:
+    url: https://huggingface.co/papers/2609.03181
+    name: Jina-OCR-v1 technical report
+    user: jinaai
+- dataset:
+    id: allenai/olmOCR-bench
+    task_id: arxiv_math
+  value: 86.1
+  source:
+    url: https://huggingface.co/papers/2609.03181
+    name: Jina-OCR-v1 technical report
+    user: jinaai
+- dataset:
+    id: allenai/olmOCR-bench
+    task_id: old_scans_math
+  value: 82.3
+  source:
+    url: https://huggingface.co/papers/2609.03181
+    name: Jina-OCR-v1 technical report
+    user: jinaai
+- dataset:
+    id: allenai/olmOCR-bench
+    task_id: table_tests
+  value: 88.8
+  source:
+    url: https://huggingface.co/papers/2609.03181
+    name: Jina-OCR-v1 technical report
+    user: jinaai
+- dataset:
+    id: allenai/olmOCR-bench
+    task_id: old_scans
+  value: 42.6
+  source:
+    url: https://huggingface.co/papers/2609.03181
+    name: Jina-OCR-v1 technical report
+    user: jinaai
+- dataset:
+    id: allenai/olmOCR-bench
+    task_id: multi_column
+  value: 85.5
+  source:
+    url: https://huggingface.co/papers/2609.03181
+    name: Jina-OCR-v1 technical report
+    user: jinaai
+- dataset:
+    id: allenai/olmOCR-bench
+    task_id: long_tiny_text
+  value: 93.2
+  source:
+    url: https://huggingface.co/papers/2609.03181
+    name: Jina-OCR-v1 technical report
+    user: jinaai
+- dataset:
+    id: allenai/olmOCR-bench
+    task_id: headers_footers
+  value: 88.7
+  source:
+    url: https://huggingface.co/papers/2609.03181
+    name: Jina-OCR-v1 technical report
+    user: jinaai
+- dataset:
+    id: allenai/olmOCR-bench
+    task_id: baseline
+  value: 99.9
+  source:
+    url: https://huggingface.co/papers/2609.03181
+    name: Jina-OCR-v1 technical report
+    user: jinaai

README.md

@@ -3,15 +3,39 @@ library_name: transformers
 pipeline_tag: image-text-to-text
 license: cc-by-nc-4.0
 language:
-- multilingual
+  - multilingual
 tags:
-- multimodal
-- multilingual
-- ocr
-- vision-language
+  - multimodal
+  - multilingual
+  - ocr
+  - vision-language
+  - document-intelligence
 base_model:
-- deepseek-ai/DeepSeek-OCR
+  - deepseek-ai/DeepSeek-OCR
 inference: false
+model-index:
+  - name: Jina-OCR-v1
+    results:
+      - task:
+          type: document-parsing
+          name: Document Parsing
+        dataset:
+          name: OmniDocBench v1.6
+          type: opendatalab/OmniDocBench
+        metrics:
+          - type: overall
+            value: 91.14
+            name: Overall Score
+      - task:
+          type: ocr
+          name: OCR
+        dataset:
+          name: OlmOCR Bench
+          type: allenai/olmOCR-bench
+        metrics:
+          - type: accuracy
+            value: 83.4
+            name: Overall Score
 ---
 
 <p align="center">
@@ -20,38 +44,32 @@ inference: false
 
 # Jina-OCR-v1: Efficient Document Parsing with Speculative Decoding and Reward-Dense Post-Training
 
-[Blog](#) | [API](https://jina.ai) | AWS | Azure | GCP | [Arxiv](#)
+[Blog](#) | [API](https://jina.ai) | AWS | Azure | GCP | [Arxiv](https://arxiv.org/abs/2609.03181)
 
-`jina-ocr-v1` is an end-to-end document parsing model designed for high-quality OCR at an *efficient serving point*. The model builds on `DeepSeek-OCR` and inherits two of its efficiency-oriented components: a **DeepEncoder** vision tower that represents a **1024×1024** global view with **256 visual tokens** and augments it with dynamic local tiles, and a **3B-parameter mixture-of-experts** decoder with approximately **570M active parameters** per token. On top of this backbone, we add a **FastMTP** speculative decoding head that recursively shares one dense draft block across **K=3** prediction steps. *Greedy verification* accepts the longest token-equality prefix, preserving the verifier’s greedy output sequence *by construction*.
+`jina-ocr-v1` is an end-to-end document parsing model designed for high-quality OCR at an _efficient serving point_. The model builds on `DeepSeek-OCR` and inherits two of its efficiency-oriented components: a **DeepEncoder** vision tower that represents a **1024×1024** global view with **256 visual tokens** and augments it with dynamic local tiles, and a **3B-parameter mixture-of-experts** decoder with approximately **570M active parameters** per token. On top of this backbone, we add a **FastMTP** speculative decoding head that recursively shares one dense draft block across **K=3** prediction steps. _Greedy verification_ accepts the longest token-equality prefix, preserving the verifier’s greedy output sequence _by construction_.
 
 This repository ships the model weights **and** the custom code needed to load them (`trust_remote_code=True`). There is no separate training or serving package to install. Transformers uses the modeling files in this snapshot; vLLM needs a one-time architecture registration from `deepseek_ocr_mtp.py` (see [Using vLLM](#vllm-fastmtp)). A single script covers both backends: [`example.py`](example.py).
 
-| | |
-|---|---|
-| Backbone | DeepSeek-OCR (DeepEncoder + 3B MoE, ~570M active) |
-| Acceleration | FastMTP, one dense draft head, K=3 (vLLM only) |
-| Local run | [`example.py`](example.py) — Transformers or vLLM |
-| Hosted | OpenAI-compatible `https://api.jina.ai/v1/chat/completions` |
+|              |                                                             |
+| ------------ | ----------------------------------------------------------- |
+| Backbone     | DeepSeek-OCR (DeepEncoder + 3B MoE, ~570M active)           |
+| Acceleration | FastMTP, one dense draft head, K=3 (vLLM only)              |
+| Local run    | [`example.py`](example.py) — Transformers or vLLM           |
+| Hosted       | OpenAI-compatible `https://api.jina.ai/v1/chat/completions` |
 
 ```bash
 python example.py --backend transformers --image document.png
 python example.py --backend vllm --image document.png
 ```
 
-
 ## Prompt
 
 ```text
 Transcribe the provided document image into a clean Markdown format, preserving the natural reading order.
-Convert all formulas into LaTeX format. Inline formulas should be enclosed in `$ $`. Display (block) formulas should be enclosed in `$$ $$`.
-Convert tables into HTML format (using <table border="1"><tr><td> tags).
-Ignore all graphical content in the image document. Do not attempt to describe or convert the images.
-Remove the headers and footers, but keep references and footnotes.
 ```
 
 Imported as `DEFAULT_OCR_PROMPT` from `deepseek_ocr_mtp.py` / the processor.
 
-
 ## Hosted API
 
 Key from [jina.ai](https://jina.ai). Image URL or `data:image/...;base64,...`. Add `"stream": true` to stream.
@@ -80,7 +98,6 @@ Local file:
 
 A cold start returns HTTP 503 — retry after 30–60s.
 
-
 ## Transformers
 
 Install `transformers`, `torch`, `torchvision`, `Pillow`. One PIL image per call. `eos` / `pad` come from `generation_config.json`; do not pass a fresh `GenerationConfig`.
@@ -106,8 +123,9 @@ print(processor.decode_ocr(output, inputs['input_ids']))
 
 `prepare_ocr_inputs(image, prompt=..., device=...)` overrides the default prompt.
 
-**No FastMTP here.** `generate()` is the 3B MoE decoder only. MTP keys (`mtp_module.*`, `mtp_embed_tokens.*`) are ignored on load.
+`generate()` attaches `SlidingWindowNoRepeatNgramProcessor` with defaults `no_repeat_ngram_size=35`, `ngram_window=1024`, and whitelist `{128821, 128822}` (`<td>`, `</td>`). Override with those keyword arguments, or disable with `no_repeat_ngram_size=0`.
 
+**No FastMTP here.** `generate()` is the 3B MoE decoder only. MTP keys (`mtp_module.*`, `mtp_embed_tokens.*`) are ignored on load.
 
 ## vLLM (FastMTP)
 
@@ -117,12 +135,12 @@ Requires [vLLM](https://github.com/vllm-project/vllm) ≥ 0.21. Call `register()
 import sys
 from huggingface_hub import snapshot_download
 from PIL import Image
-from vllm import LLM, SamplingParams
+from vllm import LLM
 
 sys.path.insert(0, snapshot_download('jinaai/jina-ocr-v1'))
 # local checkout: sys.path.insert(0, '/path/to/jina-ocr-v1')
 
-from deepseek_ocr_mtp import DEFAULT_OCR_PROMPT, log_spec_stats, register, vllm_llm_kwargs
+from deepseek_ocr_mtp import DEFAULT_OCR_PROMPT, log_spec_stats, register, vllm_llm_kwargs, vllm_sampling_params
 
 register()
 
@@ -144,7 +162,7 @@ outputs = llm.chat(
             {'type': 'text', 'text': DEFAULT_OCR_PROMPT},
         ],
     }],
-    sampling_params=SamplingParams(temperature=0.0, max_tokens=4096, repetition_penalty=1.05),
+    sampling_params=vllm_sampling_params(max_tokens=4096),
 )
 log_spec_stats(llm)
 print(outputs[0].outputs[0].text)
@@ -152,22 +170,23 @@ print(outputs[0].outputs[0].text)
 
 `vllm_llm_kwargs()` sets `trust_remote_code`, the OCR architecture override, and `speculative_config.method="eagle"`. FastMTP was trained with recursive hidden-state feedback; vLLM's default `method="mtp"` re-grounds every draft step on the target and breaks that. `register()` also maps `EagleDeepSeekMTPModel` (vLLM prepends `Eagle` when `method="eagle"`).
 
+`vllm_sampling_params()` sets `temperature=0.0`, `repetition_penalty=1.05`, and `repetition_detection` (`max_pattern_size=35`, `min_pattern_size=35`, `min_count=10`). That uses vLLM's built-in n-gram stop instead of a custom logits processor, so it works with FastMTP. Pass `repetition_detection=None` to disable.
+
 `log_spec_stats(llm)` flushes `SpecDecoding metrics` after a short run — vLLM otherwise prints them on a ~10s interval.
 
 Chat content types `image`, `image_pil`, and `image_url` all insert `<image>`. Do not pass `use_fast=False`; this checkpoint is `LlamaTokenizerFast`.
 
-
 ## Citation
 
 ```bibtex
-@misc{wang2026jinaocr,
-    title={Jina-OCR-v1: Efficient Document Parsing with Speculative Decoding and Reward-Dense Post-Training},
-    author={Feng Wang and Alejandro Barón García and Emilia Garcia Casademont and Han Xiao},
+@misc{garcía2026jinaocr,
+    title={Jina-OCR-v1: Efficient Document Parsing with Speculative Decoding and Dense Verifiable Rewards}, 
+    author={Alejandro Barón García and Feng Wang and Emilia Garcia Casademont and Han Xiao},
     year={2026},
-    eprint={xxx},
+    eprint={2609.03181},
     archivePrefix={arXiv},
     primaryClass={cs.CL},
-    url={#},
+    url={https://arxiv.org/abs/2609.03181}, 
 }
 ```
 

deepseek_ocr_mtp.py

@@ -42,7 +42,7 @@ _STACKED_PARAMS = [
     ("gate_up_proj", "up_proj", 1),
     ("qkv_proj", "q_proj", "q"),
     ("qkv_proj", "k_proj", "k"),
-    ("qkv_proj", "v_proj", "v")
+    ("qkv_proj", "v_proj", "v"),
 ]
 
 
@@ -148,7 +148,6 @@ class DeepSeekOCRMultiTokenPredictor(DeepSeekMultiTokenPredictor):
         return logits
 
 
-
 # This file postpones annotations, so vLLM cannot infer dims from live
 # types. Pass them explicitly or class import raises.
 @support_torch_compile(
@@ -321,9 +320,7 @@ _OURS = {
 }
 
 _NESTED_TEXT_CONFIG_KEYS = ("text_config", "language_config")
-_NESTED_TEXT_SKIP = frozenset(
-    (*_NESTED_TEXT_CONFIG_KEYS, "vision_config", "projector_config", "auto_map")
-)
+_NESTED_TEXT_SKIP = frozenset((*_NESTED_TEXT_CONFIG_KEYS, "vision_config", "projector_config", "auto_map"))
 
 
 def _flatten_nested_text_config(hf_config) -> None:
@@ -403,6 +400,38 @@ def register() -> None:
     _registered = True
 
 
+def _default_repetition_detection():
+    """Built-in n-gram loop stopper; works with FastMTP (custom logits processors do not)."""
+    try:
+        from vllm.sampling_params import RepetitionDetectionParams
+    except ImportError:
+        return None
+    return RepetitionDetectionParams(
+        max_pattern_size=35,
+        min_pattern_size=35,
+        min_count=10,
+    )
+
+
+def vllm_sampling_params(**kwargs):
+    """``SamplingParams`` with OCR defaults. Pass ``repetition_detection=None`` to disable."""
+    from vllm import SamplingParams
+
+    kwargs.setdefault("temperature", 0.0)
+    kwargs.setdefault("repetition_penalty", 1.05)
+    if "repetition_detection" not in kwargs:
+        detection = _default_repetition_detection()
+        if detection is not None:
+            kwargs["repetition_detection"] = detection
+    try:
+        return SamplingParams(**kwargs)
+    except TypeError as exc:
+        if "repetition_detection" in kwargs and "repetition_detection" in str(exc):
+            kwargs.pop("repetition_detection")
+            return SamplingParams(**kwargs)
+        raise
+
+
 def vllm_llm_kwargs(
     model: str,
     *,
@@ -417,6 +446,8 @@ def vllm_llm_kwargs(
 
     FastMTP (``mtp_recursive=True``) sets ``speculative_config.method="eagle"``
     so each draft step consumes the previous step's hidden state.
+    Use ``vllm_sampling_params()`` for n-gram repetition detection; vLLM rejects
+    custom logits processors while speculative decoding is enabled.
     """
     kwargs = {
         "model": model,
@@ -461,12 +492,7 @@ def log_spec_stats(llm) -> None:
 
 DEFAULT_OCR_PROMPT = (
     "Transcribe the provided document image into a clean Markdown format, "
-    "preserving the natural reading order.\n"
-    "Convert all formulas into LaTeX format. Inline formulas should be enclosed in `$ $`. "
-    "Display (block) formulas should be enclosed in `$$ $$`.\n"
-    'Convert tables into HTML format (using <table border="1"><tr><td> tags).\n'
-    "Ignore all graphical content in the image document. Do not attempt to describe or convert the images.\n"
-    "Remove the headers and footers, but keep references and footnotes."
+    "preserving the natural reading order."
 )
 
 

example.py

@@ -40,9 +40,14 @@ def run_vllm(
     max_new_tokens: int,
     num_speculative_tokens: int,
 ) -> str:
-    from vllm import LLM, SamplingParams
-
-    from deepseek_ocr_mtp import DEFAULT_OCR_PROMPT, log_spec_stats, register, vllm_llm_kwargs
+    from deepseek_ocr_mtp import (
+        DEFAULT_OCR_PROMPT,
+        log_spec_stats,
+        register,
+        vllm_llm_kwargs,
+        vllm_sampling_params,
+    )
+    from vllm import LLM
 
     register()
     llm = LLM(
@@ -64,7 +69,7 @@ def run_vllm(
     ]
     outputs = llm.chat(
         messages,
-        sampling_params=SamplingParams(temperature=0.0, max_tokens=max_new_tokens, repetition_penalty=1.05),
+        sampling_params=vllm_sampling_params(max_tokens=max_new_tokens),
     )
     if num_speculative_tokens > 0:
         log_spec_stats(llm)

model-00002-of-00002.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:71a318aefacd340c3fdbc478b222988ea21ea8eec7488eb4ca9dff9c6867394a
-size 3113227120
+oid sha256:809f8de9023bb5217d0aa396deb9e093f89862e85d66a28720bad1a6627f86e5
+size 2451310608

model.safetensors.index.json

@@ -1,6 +1,6 @@
 {
   "metadata": {
-    "total_size": 7406392320
+    "total_size": 6744476160
   },
   "weight_map": {
     "lm_head.weight": "model-00001-of-00002.safetensors",
@@ -2713,7 +2713,6 @@
     "model.vision_model.transformer.layers.9.self_attn.out_proj.weight": "model-00002-of-00002.safetensors",
     "model.vision_model.transformer.layers.9.self_attn.qkv_proj.bias": "model-00002-of-00002.safetensors",
     "model.vision_model.transformer.layers.9.self_attn.qkv_proj.weight": "model-00002-of-00002.safetensors",
-    "mtp_embed_tokens.weight": "model-00002-of-00002.safetensors",
     "mtp_module.heads.0.eh_proj.weight": "model-00002-of-00002.safetensors",
     "mtp_module.heads.0.enorm.weight": "model-00002-of-00002.safetensors",
     "mtp_module.heads.0.hnorm.weight": "model-00002-of-00002.safetensors",
@@ -2725,8 +2724,6 @@
     "mtp_module.heads.0.mtp_block.self_attn.k_proj.weight": "model-00002-of-00002.safetensors",
     "mtp_module.heads.0.mtp_block.self_attn.o_proj.weight": "model-00002-of-00002.safetensors",
     "mtp_module.heads.0.mtp_block.self_attn.q_proj.weight": "model-00002-of-00002.safetensors",
-    "mtp_module.heads.0.mtp_block.self_attn.v_proj.weight": "model-00002-of-00002.safetensors",
-    "mtp_module.heads.0.shared_head.local_head.weight": "model-00002-of-00002.safetensors",
-    "mtp_module.heads.0.shared_head.norm.weight": "model-00002-of-00002.safetensors"
+    "mtp_module.heads.0.mtp_block.self_attn.v_proj.weight": "model-00002-of-00002.safetensors"
   }
-}
+}

modeling_deepseekocr.py

@@ -7,7 +7,7 @@ import torch.nn as nn
 from PIL import Image, ImageDraw, ImageFont, ImageOps
 from torch.nn import CrossEntropyLoss
 from tqdm import tqdm
-from transformers import TextStreamer
+from transformers import LogitsProcessorList, TextStreamer
 from transformers.modeling_outputs import BaseModelOutputWithPast, MoeCausalLMOutputWithPast
 from transformers.utils import logging
 
@@ -33,6 +33,7 @@ class CausalLMOutputWithPastForMTP(MoeCausalLMOutputWithPast):
     position_ids: Optional[torch.Tensor] = None
     mtp_hidden_states: Optional[list[torch.Tensor]] = None
 
+
 class NoEOSTextStreamer(TextStreamer):
     def on_finalized_text(self, text: str, stream_end: bool = False):
 
@@ -41,6 +42,44 @@ class NoEOSTextStreamer(TextStreamer):
         print(text, flush=True, end="")
 
 
+# Match Unlimited-OCR generate() defaults.
+_DEFAULT_NGRAM_SIZE = 35
+_DEFAULT_NGRAM_WINDOW = 1024
+_DEFAULT_NGRAM_WHITELIST = frozenset({128821, 128822})  # <td>, </td>
+
+
+class SlidingWindowNoRepeatNgramProcessor:
+    """Ban tokens that would complete an n-gram already seen in the last ``window`` tokens."""
+
+    def __init__(self, ngram_size, window, whitelist_token_ids=None):
+        self.ngram_size = ngram_size
+        self.window = window
+        self.whitelist = set(whitelist_token_ids) if whitelist_token_ids else set()
+
+    def __call__(self, input_ids, scores):
+        for batch_idx in range(input_ids.shape[0]):
+            sequence = input_ids[batch_idx].tolist()
+            if len(sequence) < self.ngram_size:
+                continue
+            search_start = max(0, len(sequence) - self.window)
+            search_end = len(sequence) - self.ngram_size + 1
+            if search_end <= search_start:
+                continue
+            if self.ngram_size > 1:
+                current_prefix = tuple(sequence[-(self.ngram_size - 1) :])
+            else:
+                current_prefix = tuple()
+            banned = set()
+            for idx in range(search_start, search_end):
+                ngram = sequence[idx : idx + self.ngram_size]
+                if self.ngram_size == 1 or tuple(ngram[:-1]) == current_prefix:
+                    banned.add(ngram[-1])
+            banned.difference_update(self.whitelist)
+            for token_id in banned:
+                scores[batch_idx, token_id] = float("-inf")
+        return scores
+
+
 class DeepseekOCRConfig(DeepseekV2Config):
     model_type = "deepseek_vl_v2"
 
@@ -392,6 +431,26 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
     def prepare_inputs_for_generation(self, input_ids, **kwargs):
         return _prepare_inputs_for_generation(self, input_ids, **kwargs)
 
+    def generate(self, *args, **kwargs):
+        ngram_size = kwargs.pop("no_repeat_ngram_size", _DEFAULT_NGRAM_SIZE)
+        window_size = kwargs.pop("ngram_window", _DEFAULT_NGRAM_WINDOW)
+        whitelist_token_ids = kwargs.pop("whitelist_token_ids", _DEFAULT_NGRAM_WHITELIST)
+        processors = kwargs.get("logits_processor")
+        if processors is None:
+            processors = LogitsProcessorList()
+        elif not isinstance(processors, (list, tuple)):
+            processors = LogitsProcessorList([processors])
+        else:
+            processors = LogitsProcessorList(processors)
+        already = any(isinstance(p, SlidingWindowNoRepeatNgramProcessor) for p in processors)
+        if ngram_size and window_size and not already:
+            processors.append(
+                SlidingWindowNoRepeatNgramProcessor(ngram_size, window_size, whitelist_token_ids)
+            )
+        if processors:
+            kwargs["logits_processor"] = processors
+        return super().generate(*args, **kwargs)
+
     def disable_torch_init(self):
         """
         Disable the redundant torch default initialization to accelerate model creation.

modeling_deepseekv2.py

@@ -1214,9 +1214,7 @@ class _LlamaRotaryEmbedding(nn.Module):
         seq_len = x.shape[-2]
         if position_ids is None or position_ids.dim() != 2:
             position_ids = torch.arange(seq_len, device=x.device).unsqueeze(0)
-        inv_freq = 1.0 / (
-            self.base ** (torch.arange(0, self.dim, 2, device=x.device, dtype=torch.float32) / self.dim)
-        )
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=x.device, dtype=torch.float32) / self.dim))
         inv_freq_expanded = inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1)
         position_ids_expanded = position_ids[:, None, :].to(dtype=torch.float32)
         freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
@@ -1905,7 +1903,7 @@ class DeepseekV2Model(DeepseekV2PreTrainedModel):
 
 
 class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel, GenerationMixin):
-    _tied_weights_keys = ["lm_head.weight"]
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
 
     def __init__(self, config):
         super().__init__(config)

processing_deepseek_ocr.py

@@ -28,12 +28,7 @@ DS_OCR_BASE_SIZE = 1024
 
 DEFAULT_OCR_PROMPT = (
     "Transcribe the provided document image into a clean Markdown format, "
-    "preserving the natural reading order.\n"
-    "Convert all formulas into LaTeX format. Inline formulas should be enclosed in `$ $`. "
-    "Display (block) formulas should be enclosed in `$$ $$`.\n"
-    'Convert tables into HTML format (using <table border="1"><tr><td> tags).\n'
-    "Ignore all graphical content in the image document. Do not attempt to describe or convert the images.\n"
-    "Remove the headers and footers, but keep references and footnotes."
+    "preserving the natural reading order."
 )
 
 DS_OCR_DEFAULT_CHAT_TEMPLATE = dedent(