Baberu OCR public release (Apache-2.0): 115M multilingual manga OCR + ONNX

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/hero-en.png filter=lfs diff=lfs merge=lfs -text
+assets/hero-ja.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or Derivative
+          Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2026 Baberu OCR authors
+
+   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.

README.md

@@ -0,0 +1,176 @@
+---
+license: apache-2.0
+language:
+- ja
+- en
+- zh
+tags:
+- ocr
+- manga
+- japanese
+- image-to-text
+- vision-language-model
+- dinov2
+- onnx
+pipeline_tag: image-to-text
+---
+
+# Baberu OCR
+
+🇬🇧 English · 🇯🇵 [日本語](./README_ja.md)
+
+**A 115M-parameter manga speech-bubble OCR that beats its own 0.9B teachers on independent ground truth — in Japanese, English, and Chinese — and ships as small as 121 MB (ONNX).**
+
+![Baberu OCR — 115M beats its teachers across ja/en/zh](./assets/hero-en.png)
+
+---
+
+## Highlights
+
+- 🏆 **Beats its own teachers** on independent (non-circular) ground truth: `manga-ocr` on Japanese, and the **0.9B `PaddleOCR-VL`** on English & Chinese.
+- 🌏 **One model, three languages** (ja / en / zh). Lowest lenient character-error-rate (lCER) of every model compared, in **all three**.
+- 🪶 **115M parameters — ~1/8 the size** of the 0.9B VLMs it beats. Char-level (14,630 vocab), robust to SFX, symbols and mixed full/half-width.
+- ⚡ **Runs fully local, no cloud.** Ships as ONNX from **121 MB** (int4+int8) or **242 MB** lossless — plus a self-contained runner that needs only `onnxruntime`, `numpy`, `pillow`.
+- ✅ **Apache-2.0.**
+
+> **Honest scope:** this is SOTA for **modern manga/doujinshi speech bubbles at ≤115M**. On far-out-of-domain print (1940s Western comics) a general 0.9B VLM still wins — see [Generalization](#generalization-honest-limit).
+
+---
+
+## The model
+
+```
+bubble crop
+  → DINOv2-ViT-B/14  (86M)            # vision encoder
+  → 256 vision tokens → MLP projector (768→512)
+  → [ vision | BOS … ] → 6-layer GQA decoder (~26M)   # custom char-level LM
+  → character logits (vocab 14,630)
+```
+
+- **Vision = DINOv2** (Meta, Apache-2.0). **Language side = a small custom char-level decoder** (6 layers, GQA 8h/2kv, SwiGLU, RMSNorm-sandwich, 1D RoPE, tied embeddings) trained from scratch.
+- **Training recipe that made it strong:** first train with the **vision encoder frozen**, then **unfreeze the vision encoder and continue** — this second stage gives a large accuracy jump, and saturates after only ~30k steps (≈10% of an epoch). char-level means no tokenizer OOV problems, which helps SFX and mixed scripts.
+- **115M total.** PyTorch (CUDA / Apple-MPS / CPU) and ONNX (CPU / CUDA / DirectML / web).
+
+---
+
+## Benchmarks (independent ground truth)
+
+Teacher-generated labels would be circular, so we score against **independent** GT: Japanese = the human-labeled **Manga109-v2026**; English/Chinese = **re-annotated by hand** (teacher labels not used). Primary metric = **lCER** (lenient char-error-rate: NFKC + whitespace-strip + drop decorative symbols + unify dot-runs), lower is better. Baberu was **not** trained on any of these sets.
+
+**Japanese — Manga109-v2026 (n=2000)**
+
+| model | size | lCER ↓ | nCER ↓ |
+|---|---|---|---|
+| **Baberu OCR** | **115M** | **0.0345** | 0.0871 |
+| PaddleOCR-VL | 0.9B | 0.0368 | **0.0808** |
+| manga-ocr *(its teacher)* | 110M | 0.0422 | 0.0954 |
+| PP-OCRv5 | ~10M | 0.6317 | 0.6491 |
+
+Baberu has the **best lCER**, and **beats its teacher manga-ocr on every metric**. *(Honest note: the 0.9B PaddleOCR-VL edges Baberu on strict nCER 0.0808 vs 0.0871 and on exact-match. `PaddleOCR-VL-For-Manga` posts a lower number here but is **excluded — it was trained on Manga109, so that score is a data leak**.)*
+
+**English — in-domain doujin, hand-annotated (n=80)**
+
+| model | size | lCER ↓ | exact ↑ |
+|---|---|---|---|
+| **Baberu OCR** | **115M** | **0.0177** | **0.89** |
+| PaddleOCR-VL *(its teacher)* | 0.9B | 0.0268 | 0.88 |
+| PP-OCRv5 | ~10M | 0.0975 | 0.63 |
+| PaddleOCR-VL-For-Manga | 0.9B | 0.1140 | 0.41 |
+| manga-ocr | 110M | 0.2987 | 0.24 |
+
+**Chinese — in-domain doujin, hand-annotated (n=77)**
+
+| model | size | lCER ↓ | nCER ↓ |
+|---|---|---|---|
+| **Baberu OCR** | **115M** | **0.0868** | **0.1244** |
+| PaddleOCR-VL *(its teacher)* | 0.9B | 0.1013 | 0.1808 |
+| PaddleOCR-VL-For-Manga | 0.9B | 0.1168 | 0.1605 |
+| manga-ocr | 110M | 0.2506 | 0.2518 |
+| PP-OCRv5 | ~10M | 0.5639 | 0.5778 |
+
+→ **Baberu beats the 0.9B teacher in both English and Chinese**, at 1/8 the size.
+
+### Generalization (honest limit)
+
+Public, never-seen gold GT — **COMICS Text+** (1940s Western comics, n=480):
+
+| model | size | lCER ↓ |
+|---|---|---|
+| PaddleOCR-VL | 0.9B | **0.0256** |
+| PP-OCRv5 | ~10M | 0.1332 |
+| **Baberu OCR** | **115M** | 0.1488 |
+
+Baberu still reads this far-out-of-domain set zero-shot, but on clean vintage print the general 0.9B model wins. Baberu's strength is **modern manga/doujinshi bubbles**, by design.
+
+---
+
+## Quick start
+
+Input is **one speech-bubble crop** (use a text detector upstream), not a whole page.
+
+### ONNX — easiest & smallest (no PyTorch)
+
+```bash
+pip install onnxruntime numpy pillow
+huggingface-cli download genshiai-daichi/baberu-ocr --local-dir baberu-ocr
+cd baberu-ocr
+
+python onnx_infer.py --image bubble.jpg                          # 242 MB, lossless
+python onnx_infer.py --vision vision_int4.onnx --image bubble.jpg # 121 MB, smallest
+```
+
+The runner is fully self-contained (`onnx_infer.py`): preprocessing, KV-cache decode loop, and the published decode settings (greedy + `repetition_penalty=1.2` + symbol-aware content-run cap) are all included, so you get the same text as the PyTorch model.
+
+### PyTorch — for fine-tuning / `transformers`
+
+```bash
+pip install torch transformers safetensors pillow
+huggingface-cli download genshiai-daichi/baberu-ocr --local-dir baberu-ocr
+cd baberu-ocr
+python inference.py --model . --image bubble.jpg
+```
+
+```python
+from PIL import Image
+from inference import BaberuOCR
+ocr = BaberuOCR(".")                 # local download dir
+print(ocr(Image.open("bubble.jpg")))
+```
+
+---
+
+## Sizes & quantization (ONNX)
+
+| component | fp16 | int8 | int4 |
+|---|---|---|---|
+| vision (86M) | 173 MB | 92 MB | **52 MB** |
+| decoder (prefill+step) | — | **69 MB** | (89 MB) |
+
+Two recommended tiers (both keep the decoder at **int8 — which is accuracy-free here**):
+
+| tier | files | size | nCER (ja) |
+|---|---|---|---|
+| **lossless** | `vision_fp16.onnx` + `decoder_*_int8.onnx` | **242 MB** | 0.0867 |
+| **smallest** | `vision_int4.onnx` + `decoder_*_int8.onnx` | **121 MB** | 0.0893 |
+
+Notes from the measured sweep (see `quantization-results.md`): **decoder int8 is free** (nCER 0.0867 ≈ fp32 0.0871 — greedy decoding is robust to weight-quantization noise); the only real cost is vision quantization, and even there **int4 weight-only beats int8-dynamic** (smaller *and* more accurate). Both tiers still beat the teacher manga-ocr (0.0954). The strictly-dominated variants (vision int8-dynamic, decoder int4) are intentionally not shipped.
+
+---
+
+## Limitations (honest)
+
+- **Single bubble crop in, text out** — you need a text-detection stage upstream. Not a page-level OCR.
+- Weak on **dakuten/handakuten** confusions (ポ↔ボ), some **rare kanji**, and dropped small kana.
+- **Far-out-of-domain** (clean vintage/Western print) → a general large VLM is better.
+- Adult content appears in the training distribution; the model itself only transcribes text.
+
+## Training data & license
+
+- Trained on Japanese **manga / doujinshi** speech-bubble crops. Text labels were distilled from `manga-ocr` (ja) and `PaddleOCR-VL` (zh/en). **Training images are not redistributed** — only the model weights and code are released.
+- **License: Apache-2.0** (weights + code). Compatible with the upstream components used (DINOv2, manga-ocr, PaddleOCR all Apache-2.0).
+
+## Acknowledgements
+
+- [DINOv2](https://github.com/facebookresearch/dinov2) (Meta) — vision encoder
+- [manga-ocr](https://github.com/kha-white/manga-ocr) — the de-facto Japanese manga OCR baseline / ja teacher
+- [PaddleOCR-VL](https://github.com/PaddlePaddle/PaddleOCR) — zh/en teacher

README_ja.md

@@ -0,0 +1,176 @@
+---
+license: apache-2.0
+language:
+- ja
+- en
+- zh
+tags:
+- ocr
+- manga
+- japanese
+- image-to-text
+- vision-language-model
+- dinov2
+- onnx
+pipeline_tag: image-to-text
+---
+
+# Baberu OCR
+
+🇯🇵 日本語 · 🇬🇧 [English](./README.md)
+
+**115M パラメータのマンガ吹き出し OCR。自分の教師（0.9B級）を独立GTで上回る — 日本語・英語・中国語すべてで。ONNX なら最小 121MB。**
+
+![Baberu OCR — 115Mで教師超え、ja/en/zh 全制覇](./assets/hero-ja.png)
+
+---
+
+## ハイライト
+
+- 🏆 **自分の教師に勝つ**（身内採点でない独立GTで）: 日本語は `manga-ocr`、英語・中国語は **0.9B の `PaddleOCR-VL`**。
+- 🌏 **1モデルで3言語**（ja / en / zh）。比較した全モデル中、**3言語すべてで最小の lCER**（緩和文字誤り率）。
+- 🪶 **115M ＝ 勝った相手(0.9B)の約1/8**。char-level（語彙14,630）で擬音・記号・全半角混在に強い。
+- ⚡ **完全ローカル動作**。ONNX は **121MB**（int4+int8）/ **242MB**（無劣化）。`onnxruntime`・`numpy`・`pillow` だけで動く自己完結ランナー同梱。
+- ✅ **Apache-2.0**。
+
+> **正直なスコープ:** これは「**モダンなマンガ/同人の吹き出し × ≤115M**」での SOTA。遠いドメイン（1940年代の欧米コミック等）の clean な印字では汎用 0.9B が依然上 → [汎化の上限](#汎化の上限正直な開示)。
+
+---
+
+## モデル概要
+
+```
+吹き出し crop
+  → DINOv2-ViT-B/14  (86M)            # 視覚エンコーダ
+  → 256 vision tokens → MLP projector (768→512)
+  → [ vision | BOS … ] → 6層 GQA decoder (~26M)   # 独自の char-level LM
+  → 文字ロジット（語彙 14,630）
+```
+
+- **視覚 = DINOv2**（Meta, Apache-2.0）。**言語側 = 独自の小型 char-level decoder**（6層, GQA 8h/2kv, SwiGLU, RMSNorm sandwich, 1D RoPE, tied embedding）をフルスクラッチ学習。
+- **強さの肝＝学習レシピ:** まず**視覚エンコーダを凍結**して学習 → その後**視覚を解凍して継続学習**すると、この2段目で精度が大きく跳ね上がる（しかも解凍後 ~30k step ＝ 1エポックの約10% で頭打ち）。char-level なので tokenizer 由来の未知語が無く、擬音・多言語混在に強い。
+- **合計 115M**。PyTorch（CUDA / Apple-MPS / CPU）と ONNX（CPU / CUDA / DirectML / web）対応。
+
+---
+
+## ベンチマーク（独立GT）
+
+教師が作ったラベルで測ると「身内採点」になるため、**独立GT**で採点: 日本語は人手ラベルの **Manga109-v2026**、英語/中国語は**人手で再注釈**（教師ラベル不使用）。主指標 = **lCER**（緩和文字誤り率: NFKC＋空白除去＋装飾記号除去＋`…/・・・`統一、低いほど良い）。Baberu はいずれも**未学習**。
+
+**日本語 — Manga109-v2026 (n=2000)**
+
+| model | size | lCER ↓ | nCER ↓ |
+|---|---|---|---|
+| **Baberu OCR** | **115M** | **0.0345** | 0.0871 |
+| PaddleOCR-VL | 0.9B | 0.0368 | **0.0808** |
+| manga-ocr *(教師)* | 110M | 0.0422 | 0.0954 |
+| PP-OCRv5 | ~10M | 0.6317 | 0.6491 |
+
+**lCER 最良**、そして**教師 manga-ocr には全指標で勝利**。*（正直な注記: strict nCER と完全一致では 0.9B の PaddleOCR-VL が僅かに上 0.0808 vs 0.0871。`PaddleOCR-VL-For-Manga` はさらに低い値を出すが **Manga109 で学習済＝データリークのため除外**。）*
+
+**英語 — in-domain doujin・人手GT (n=80)**
+
+| model | size | lCER ↓ | 完全一致 ↑ |
+|---|---|---|---|
+| **Baberu OCR** | **115M** | **0.0177** | **0.89** |
+| PaddleOCR-VL *(教師)* | 0.9B | 0.0268 | 0.88 |
+| PP-OCRv5 | ~10M | 0.0975 | 0.63 |
+| PaddleOCR-VL-For-Manga | 0.9B | 0.1140 | 0.41 |
+| manga-ocr | 110M | 0.2987 | 0.24 |
+
+**中国語 — in-domain doujin・人手GT (n=77)**
+
+| model | size | lCER ↓ | nCER ↓ |
+|---|---|---|---|
+| **Baberu OCR** | **115M** | **0.0868** | **0.1244** |
+| PaddleOCR-VL *(教師)* | 0.9B | 0.1013 | 0.1808 |
+| PaddleOCR-VL-For-Manga | 0.9B | 0.1168 | 0.1605 |
+| manga-ocr | 110M | 0.2506 | 0.2518 |
+| PP-OCRv5 | ~10M | 0.5639 | 0.5778 |
+
+→ **英語・中国語とも、教師 0.9B を 1/8 サイズで上回る**。
+
+### 汎化の上限（正直な開示）
+
+公開・未学習の gold GT **COMICS Text+**（1940年代の欧米コミック, n=480）:
+
+| model | size | lCER ↓ |
+|---|---|---|
+| PaddleOCR-VL | 0.9B | **0.0256** |
+| PP-OCRv5 | ~10M | 0.1332 |
+| **Baberu OCR** | **115M** | 0.1488 |
+
+遠いドメインでも zero-shot で読めるが、clean な古印字では汎用 0.9B が勝つ。Baberu の本領は**モダンなマンガ/同人吹き出し**（設計どおり）。
+
+---
+
+## クイックスタート
+
+入力は**吹き出し1個の crop**（上流にテキスト検出器が必要）。ページ全体ではない。
+
+### ONNX — 最も簡単＆最小（PyTorch不要）
+
+```bash
+pip install onnxruntime numpy pillow
+huggingface-cli download genshiai-daichi/baberu-ocr --local-dir baberu-ocr
+cd baberu-ocr
+
+python onnx_infer.py --image bubble.jpg                          # 242MB, 無劣化
+python onnx_infer.py --vision vision_int4.onnx --image bubble.jpg # 121MB, 最小
+```
+
+`onnx_infer.py` は自己完結（前処理・KVキャッシュのデコードループ・公開済みデコード設定 = greedy ＋ `repetition_penalty=1.2` ＋ 記号自由のcontent-run cap を全部内蔵）なので、PyTorch版と同じテキストが得られる。
+
+### PyTorch — fine-tuning / `transformers` 向け
+
+```bash
+pip install torch transformers safetensors pillow
+huggingface-cli download genshiai-daichi/baberu-ocr --local-dir baberu-ocr
+cd baberu-ocr
+python inference.py --model . --image bubble.jpg
+```
+
+```python
+from PIL import Image
+from inference import BaberuOCR
+ocr = BaberuOCR(".")                 # ローカルのDLディレクトリ
+print(ocr(Image.open("bubble.jpg")))
+```
+
+---
+
+## サイズと量子化（ONNX）
+
+| 部品 | fp16 | int8 | int4 |
+|---|---|---|---|
+| vision (86M) | 173MB | 92MB | **52MB** |
+| decoder (prefill+step) | — | **69MB** | (89MB) |
+
+推奨2ティア（どちらも **decoder は int8 ＝ ここでは精度タダ** に固定）:
+
+| ティア | ファイル | サイズ | nCER (ja) |
+|---|---|---|---|
+| **無劣化** | `vision_fp16.onnx` + `decoder_*_int8.onnx` | **242MB** | 0.0867 |
+| **最小** | `vision_int4.onnx` + `decoder_*_int8.onnx` | **121MB** | 0.0893 |
+
+実測スイープの要点（`quantization-results.md`）: **decoder int8 は劣化ゼロ**（nCER 0.0867 ≈ fp32 0.0871 ＝ greedy デコードは重み量子化ノイズに頑健）。コストが出るのは vision 量子化だけで、そこでも **int4(weight-only) が int8(dynamic) より小さく かつ 正確**。どちらのティアも教師 manga-ocr (0.0954) に勝つ。劣勢な構成（vision int8-dynamic / decoder int4）は同梱しない。
+
+---
+
+## 制限（正直な開示）
+
+- **入力は吹き出し1 crop**（上流にテキスト検出が必要）。ページOCRではない。
+- **濁点/半濁点**の取り違え（ポ↔ボ）、一部の**稀少漢字**、小書きカナ脱落に弱い。
+- **遠いドメイン**（clean な欧米/古印字）では汎用大型VLMが上。
+- 学習分布に成人向けコンテンツを含むが、モデルはテキストの転写のみを行う。
+
+## 学習データ・ライセンス
+
+- 日本語の**マンガ/同人**吹き出し crop で学習。テキストラベルは `manga-ocr`（ja）と `PaddleOCR-VL`（zh/en）から蒸留。**学習画像は再配布しない** — 公開するのはモデル重みとコードのみ。
+- **ライセンス: Apache-2.0**（重み＋コード）。上流（DINOv2 / manga-ocr / PaddleOCR、いずれも Apache-2.0）と整合。
+
+## 謝辞
+
+- [DINOv2](https://github.com/facebookresearch/dinov2)（Meta）— 視覚エンコーダ
+- [manga-ocr](https://github.com/kha-white/manga-ocr) — 日本語マンガOCRの事実上の標準 / ja教師
+- [PaddleOCR-VL](https://github.com/PaddlePaddle/PaddleOCR) — zh/en 教師

config.json

@@ -0,0 +1,39 @@
+{
+  "architectures": [
+    "BaberuOCRModel"
+  ],
+  "attention_dropout": 0.0,
+  "attn_logit_softcap": null,
+  "bos_token_id": 1,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "final_logit_softcap": 30.0,
+  "freeze_vision_encoder": false,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 512,
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "max_position_embeddings": 2048,
+  "model_type": "baberu_ocr",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 6,
+  "num_key_value_heads": 2,
+  "pad_token_id": 0,
+  "projector_act": "gelu",
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 10000.0,
+  "sandwich_norm": true,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.5.3",
+  "unk_token_id": 3,
+  "use_cache": true,
+  "vision_hidden_size": 768,
+  "vision_image_size": 224,
+  "vision_model_name": "facebook/dinov2-base",
+  "vision_num_tokens": 256,
+  "vision_patch_size": 14,
+  "vision_rope_2d": false,
+  "vocab_size": 14630,
+  "z_loss_weight": 0.0001
+}

configuration_baberu.py

@@ -0,0 +1,129 @@
+"""Configuration for baberu-ocr v2.
+
+Single config that drives both the text-only decoder (BaberuCausalLM, used
+for Step 1 pretraining) and the full vision-grounded OCR model
+(BaberuOCRModel = DINOv2 + MLP projector + decoder).
+
+See docs/design/ocr-v2-design.md for the full architecture rationale.
+"""
+
+from __future__ import annotations
+
+from transformers import PretrainedConfig
+
+
+class BaberuOCRConfig(PretrainedConfig):
+    """Config for baberu-ocr v2 (~115M params, multilingual JA+ZH+EN).
+
+    Defaults match the May 2026 SOTA design:
+      - Vision: DINOv2-ViT-B/14 (86M, Apache 2.0)
+      - Projector: MLP 768 -> hidden_size
+      - Decoder: 6-layer GQA causal LM, hidden=512, FFN=1536
+      - RoPE 1D (text), 2D-RoPE for vision optional
+      - RMSNorm sandwich-norm, SwiGLU, logit soft-cap, z-loss, tied embeddings
+
+    Use ``vision_model_name=None`` to disable the vision branch
+    (text-only mode for Step 1 pretraining).
+    """
+
+    model_type = "baberu_ocr"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        # ---- Vision encoder ----
+        vision_model_name: str | None = "facebook/dinov2-base",
+        vision_hidden_size: int = 768,
+        vision_image_size: int = 224,
+        vision_patch_size: int = 14,
+        vision_num_tokens: int = 256,  # (224/14)^2
+        freeze_vision_encoder: bool = True,
+        # ---- Projector ----
+        projector_act: str = "gelu",
+        # ---- Decoder ----
+        vocab_size: int = 14630,
+        hidden_size: int = 512,
+        intermediate_size: int = 1536,
+        num_hidden_layers: int = 6,
+        num_attention_heads: int = 8,
+        num_key_value_heads: int = 2,
+        head_dim: int | None = None,  # default = hidden_size // num_attention_heads
+        max_position_embeddings: int = 2048,
+        hidden_act: str = "silu",
+        # ---- Norm ----
+        rms_norm_eps: float = 1e-6,
+        sandwich_norm: bool = True,
+        # ---- RoPE ----
+        rope_theta: float = 10000.0,
+        vision_rope_2d: bool = False,  # 2D RoPE for vision prefix (planned)
+        # ---- Stability ----
+        final_logit_softcap: float = 30.0,
+        attn_logit_softcap: float | None = None,
+        z_loss_weight: float = 1e-4,
+        # ---- Embeddings ----
+        tie_word_embeddings: bool = True,
+        # ---- Init ----
+        initializer_range: float = 0.02,
+        # ---- Special tokens ----
+        pad_token_id: int = 0,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        unk_token_id: int = 3,
+        # ---- Cache / dtype ----
+        use_cache: bool = True,
+        attention_dropout: float = 0.0,
+        **kwargs,
+    ):
+        self.vision_model_name = vision_model_name
+        self.vision_hidden_size = vision_hidden_size
+        self.vision_image_size = vision_image_size
+        self.vision_patch_size = vision_patch_size
+        self.vision_num_tokens = vision_num_tokens
+        self.freeze_vision_encoder = freeze_vision_encoder
+
+        self.projector_act = projector_act
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim or (hidden_size // num_attention_heads)
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_act = hidden_act
+
+        self.rms_norm_eps = rms_norm_eps
+        self.sandwich_norm = sandwich_norm
+
+        self.rope_theta = rope_theta
+        self.vision_rope_2d = vision_rope_2d
+
+        self.final_logit_softcap = final_logit_softcap
+        self.attn_logit_softcap = attn_logit_softcap
+        self.z_loss_weight = z_loss_weight
+
+        self.initializer_range = initializer_range
+        self.use_cache = use_cache
+        self.attention_dropout = attention_dropout
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        # Stored separately because PretrainedConfig doesn't claim unk.
+        self.unk_token_id = unk_token_id
+
+    # Convenience for sanity checks
+    def validate(self) -> None:
+        assert self.num_attention_heads % self.num_key_value_heads == 0, (
+            f"num_attention_heads ({self.num_attention_heads}) must be divisible "
+            f"by num_key_value_heads ({self.num_key_value_heads}) for GQA"
+        )
+        assert self.head_dim * self.num_attention_heads == self.hidden_size, (
+            f"head_dim ({self.head_dim}) * num_attention_heads "
+            f"({self.num_attention_heads}) must equal hidden_size ({self.hidden_size})"
+        )

generation_config.json

@@ -0,0 +1,12 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "max_new_tokens": 128,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 0,
+  "repetition_penalty": 1.2,
+  "transformers_version": "5.5.3",
+  "use_cache": true
+}

inference.py

@@ -0,0 +1,161 @@
+"""Minimal, canonical inference for the custom Baberu OCR VLM (DINOv2 + projector
++ 6-layer GQA char decoder, ~115M). This is the source-of-truth example shipped
+with the model card.
+
+DECODING — IMPORTANT (see docs/design/ocr-indomain-adjudication-2026-06-19.md):
+  Two complementary anti-loop guards, both ON by default — they reproduce the
+  published accuracy numbers:
+    1. ``repetition_penalty=1.2`` (NOT ``no_repeat_ngram_size``: the latter hard-caps
+       identical runs at N, TRUNCATING legit manga punctuation like ・・・・・・ / ーー —
+       it cost ~5pt of exact-match agreement on real crops).
+    2. a SYMBOL-AWARE content-run cap: caps runs of identical CONTENT chars
+       (kanji/kana/latin/digit) at 12, while symbols/・・・/ーー repeat freely. This
+       kills pathological single-char loops (啊×63, む×61) that repetition_penalty
+       alone misses — without it, zh lCER regresses from 0.087 to ~0.147.
+
+Usage:
+  python -m baberu.ocr_v2.inference --model <ckpt_dir> --image page_crop.jpg
+  # or, programmatically:
+  ocr = BaberuOCR("<ckpt_dir>"); print(ocr(Image.open("crop.jpg")))
+"""
+from __future__ import annotations
+
+import argparse
+import unicodedata
+from pathlib import Path
+from typing import Union
+
+import torch
+from PIL import Image
+from safetensors.torch import load_file as load_safetensors
+from transformers import LogitsProcessor, LogitsProcessorList
+
+from configuration_baberu import BaberuOCRConfig
+from modeling_baberu import BaberuOCRModel
+from tokenization_baberu import BaberuTokenizer
+
+
+def build_ocr_image_processor(vision_model: str = "facebook/dinov2-base"):
+    """DINOv2 processor with do_center_crop=False — resize the WHOLE crop to 224x224
+    (tall 縦書き crops lose their top/bottom under the default 256→center-crop-224).
+    Matches training/eval preprocessing exactly."""
+    from transformers import AutoImageProcessor
+
+    return AutoImageProcessor.from_pretrained(
+        vision_model, do_center_crop=False, size={"height": 224, "width": 224}, crop_size=None
+    )
+
+
+class CapContentRun(LogitsProcessor):
+    """Cap runs of identical CONTENT chars (Unicode L/N, excl. long-vowel ー～) at
+    ``max_run``; symbols/punctuation/・・・/ーー repeat freely. Kills single-char
+    loops the repetition_penalty misses, without clipping legit symbol runs."""
+
+    def __init__(self, content_ids: set[int], max_run: int = 12):
+        self.content_ids = content_ids
+        self.max_run = max_run
+
+    def __call__(self, input_ids, scores):
+        for b in range(input_ids.shape[0]):
+            t = int(input_ids[b, -1])
+            if t in self.content_ids:
+                run, j = 0, input_ids.shape[1] - 1
+                while j >= 0 and int(input_ids[b, j]) == t:
+                    run += 1
+                    j -= 1
+                    if run >= self.max_run:
+                        break
+                if run >= self.max_run:
+                    scores[b, t] = float("-inf")
+        return scores
+
+
+def _build_content_ids(tok) -> set[int]:
+    def is_content(ch: str) -> bool:
+        if not ch or len(ch) != 1 or ch in "ーｰ〜~":
+            return False
+        return unicodedata.category(ch)[0] in "LN"
+
+    n = getattr(tok, "vocab_size", None) or len(tok)
+    return {i for i in range(n) if is_content(tok.decode([i], skip_special_tokens=False))}
+
+
+class BaberuOCR:
+    """Load a Step-3 checkpoint and OCR a single bubble crop."""
+
+    def __init__(self, model_dir: Union[str, Path], device: str | None = None):
+        model_dir = Path(model_dir)
+        self.device = torch.device(
+            device or ("cuda" if torch.cuda.is_available() else
+                       "mps" if torch.backends.mps.is_available() else "cpu")
+        )
+        dtype = torch.bfloat16 if self.device.type == "cuda" else torch.float32
+
+        config = BaberuOCRConfig.from_pretrained(model_dir)
+        self.model = BaberuOCRModel(config)
+        # strict=False: a non-persistent buffer (causal mask / rope) may be missing;
+        # all trained weights (incl. the vision encoder) load from the checkpoint.
+        self.model.load_state_dict(
+            load_safetensors(model_dir / "model.safetensors"), strict=False
+        )
+        self.model.to(device=self.device, dtype=dtype).eval()
+
+        tok_dir = model_dir / "tokenizer"
+        self.tok = BaberuTokenizer.from_pretrained(
+            tok_dir if tok_dir.exists() else model_dir
+        )
+        # SAME preprocessing as training: no center-crop, whole crop -> 224x224.
+        self.image_processor = build_ocr_image_processor(config.vision_model_name)
+        self._content_ids = _build_content_ids(self.tok)
+
+    @torch.inference_mode()
+    def __call__(
+        self,
+        image: Image.Image,
+        max_new_tokens: int = 128,
+        repetition_penalty: float = 1.2,
+        max_content_run: int = 12,
+    ) -> str:
+        image = image.convert("RGB")
+        pixel_values = self.image_processor(image, return_tensors="pt")["pixel_values"].to(self.device)
+        input_ids = torch.tensor([[self.tok.bos_token_id]], device=self.device)
+        logits_processor = LogitsProcessorList()
+        if max_content_run:
+            logits_processor.append(CapContentRun(self._content_ids, max_content_run))
+        output_ids = self.model.generate(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+            repetition_penalty=repetition_penalty,  # NOT no_repeat_ngram_size — see module docstring
+            logits_processor=logits_processor,       # symbol-aware content-run cap
+            eos_token_id=self.tok.eos_token_id,
+            pad_token_id=self.tok.pad_token_id,
+            bos_token_id=self.tok.bos_token_id,
+            use_cache=True,
+        )
+        pred_ids = output_ids[0, input_ids.shape[1]:].detach().cpu().tolist()
+        return self.tok.decode(pred_ids, skip_special_tokens=True)
+
+
+def main():
+    ap = argparse.ArgumentParser(description="OCR a manga bubble crop with the custom Baberu model.")
+    ap.add_argument("--model", required=True, help="Step-3 checkpoint dir")
+    ap.add_argument("--image", required=True, help="crop image to OCR")
+    ap.add_argument("--device", default=None)
+    ap.add_argument("--max-new-tokens", type=int, default=128)
+    ap.add_argument("--repetition-penalty", type=float, default=1.2)
+    ap.add_argument("--max-content-run", type=int, default=12,
+                    help="cap consecutive identical content chars (0=off)")
+    args = ap.parse_args()
+
+    ocr = BaberuOCR(args.model, device=args.device)
+    text = ocr(Image.open(args.image),
+               max_new_tokens=args.max_new_tokens,
+               repetition_penalty=args.repetition_penalty,
+               max_content_run=args.max_content_run)
+    print(text)
+
+
+if __name__ == "__main__":
+    main()

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4b389c67e3ae84db05270920b55404c41085324d098c431a2f79ca92273288df
+size 451347624

modeling_baberu.py

@@ -0,0 +1,825 @@
+"""baberu-ocr v2: prefix-token VLM for multilingual manga OCR.
+
+Architecture (May 2026 SOTA):
+    Vision (DINOv2) -> MLP projector -> prefix tokens -> [BOS] text [EOS]
+                                              \\
+                                               -> GQA causal LM (6 layers)
+
+Decoder details:
+  - GQA self-attention (Llama-style) with 1D RoPE (text positions)
+  - SwiGLU FFN
+  - RMSNorm sandwich-norm (Gemma2-style: norm before AND after each sub-block)
+  - Tied input/output embeddings
+  - Logit soft-cap (Gemma2-style: tanh(logit/cap)*cap)
+  - Z-loss aux for logit-drift stability
+  - KV cache for fast decode
+
+Two entry points:
+  - ``BaberuCausalLM``  : text-only mode (Step 1 pretraining on FineWeb2)
+  - ``BaberuOCRModel``  : vision + decoder for OCR (Steps 3-1..3-3)
+
+2D RoPE for vision prefix is planned (see config.vision_rope_2d) but currently
+falls back to 1D positions over the combined [vision || text] sequence.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoConfig, AutoModel, PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+
+from configuration_baberu import BaberuOCRConfig
+
+
+# ---------------------------------------------------------------------------
+# Small building blocks
+# ---------------------------------------------------------------------------
+
+
+class BaberuRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_dtype = x.dtype
+        x_f32 = x.to(torch.float32)
+        var = x_f32.pow(2).mean(-1, keepdim=True)
+        x_f32 = x_f32 * torch.rsqrt(var + self.eps)
+        return (self.weight * x_f32).to(input_dtype)
+
+
+class BaberuMLP(nn.Module):
+    """SwiGLU FFN: down(silu(gate) * up)."""
+
+    def __init__(self, config: BaberuOCRConfig):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+# ---------------------------------------------------------------------------
+# Rotary position embedding (1D)
+# ---------------------------------------------------------------------------
+
+
+class BaberuRotaryEmbedding(nn.Module):
+    """Precomputed 1D RoPE table. Llama-style."""
+
+    def __init__(self, dim: int, max_seq_len: int, theta: float = 10000.0):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.max_seq_len_cached = 0
+        self._cos_cached: torch.Tensor | None = None
+        self._sin_cached: torch.Tensor | None = None
+        self._build_cache(max_seq_len, device=torch.device("cpu"))
+
+    def _build_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype = torch.float32):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(seq_len, device=device, dtype=torch.float32)
+        # Recompute inv_freq from theta/dim rather than reading the buffer. It is
+        # registered non-persistent (not in the checkpoint), so under
+        # `from_pretrained` it is created on the meta device then materialised by
+        # `to_empty()` to *uninitialised* memory — no longer `is_meta`, but filled
+        # with garbage that silently breaks RoPE (positions collapse, output turns
+        # to junk). Recomputing here is cheap and always correct.
+        inv_freq = 1.0 / (
+            self.theta ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim)
+        )
+        freqs = torch.einsum("i,j->ij", t, inv_freq)
+        emb = torch.cat([freqs, freqs], dim=-1)
+        self._cos_cached = emb.cos().to(dtype)
+        self._sin_cached = emb.sin().to(dtype)
+
+    def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        # position_ids: [B, T] long
+        seq_max = int(position_ids.max().item()) + 1
+        if (
+            self._cos_cached is None
+            or seq_max > self.max_seq_len_cached
+            or self._cos_cached.device != x.device
+            or self._cos_cached.dtype != x.dtype
+        ):
+            self._build_cache(max(seq_max, self.max_seq_len_cached, 64), device=x.device, dtype=x.dtype)
+        cos = self._cos_cached[position_ids]  # [B, T, dim]
+        sin = self._sin_cached[position_ids]
+        return cos, sin
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def apply_rotary_pos_emb(
+    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    # q, k: [B, n_heads, T, head_dim]; cos, sin: [B, T, head_dim]
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)
+    q_rot = (q * cos) + (_rotate_half(q) * sin)
+    k_rot = (k * cos) + (_rotate_half(k) * sin)
+    return q_rot, k_rot
+
+
+# ---------------------------------------------------------------------------
+# GQA self-attention with KV cache
+# ---------------------------------------------------------------------------
+
+
+def _repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """Expand KV heads to match Q heads for GQA. [B, kv_heads, T, D] -> [B, q_heads, T, D]."""
+    if n_rep == 1:
+        return hidden_states
+    bsz, kv_heads, slen, head_dim = hidden_states.shape
+    hidden_states = hidden_states[:, :, None, :, :].expand(bsz, kv_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(bsz, kv_heads * n_rep, slen, head_dim)
+
+
+def _cache_length(past_key_values) -> int:
+    """Return number of tokens already in the cache, handling DynamicCache and legacy tuples."""
+    if past_key_values is None:
+        return 0
+    if hasattr(past_key_values, "layers"):
+        layers = past_key_values.layers
+        if not layers or not layers[0].is_initialized:
+            return 0
+        return layers[0].keys.shape[-2]
+    if isinstance(past_key_values, (list, tuple)) and len(past_key_values) > 0:
+        return past_key_values[0][0].shape[-2]
+    return 0
+
+
+class BaberuAttention(nn.Module):
+    def __init__(self, config: BaberuOCRConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.head_dim
+        self.num_kv_heads = config.num_key_value_heads
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+        self.attn_softcap = config.attn_logit_softcap
+        self.attention_dropout = config.attention_dropout
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,  # [B, T, H]
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,  # [B, 1, T, S] additive (-inf for masked)
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.shape
+
+        q = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        if past_key_value is not None:
+            past_k, past_v = past_key_value
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+
+        present = (k, v) if use_cache else None
+
+        k_expanded = _repeat_kv(k, self.num_kv_groups)
+        v_expanded = _repeat_kv(v, self.num_kv_groups)
+
+        # Scaled dot-product attention.  Q: [B, H, Tq, D], K/V: [B, H, Tk, D]
+        if self.attn_softcap is not None:
+            # SDPA/flash kernels cannot apply Gemma-style attention logit
+            # soft-capping, so keep the explicit (memory-heavy) path only when
+            # soft-capping is enabled.
+            attn_scores = torch.matmul(q, k_expanded.transpose(-1, -2)) / (self.head_dim ** 0.5)
+            attn_scores = torch.tanh(attn_scores / self.attn_softcap) * self.attn_softcap
+            if attention_mask is not None:
+                attn_scores = attn_scores + attention_mask[:, :, :, : k.shape[-2]]
+            attn_weights = F.softmax(attn_scores.to(torch.float32), dim=-1).to(q.dtype)
+            if self.attention_dropout > 0 and self.training:
+                attn_weights = F.dropout(attn_weights, p=self.attention_dropout)
+            attn_output = torch.matmul(attn_weights, v_expanded)  # [B, H, Tq, D]
+        else:
+            # FlashAttention / mem-efficient SDPA: never materialises the
+            # [B, H, Tq, Tk] score matrix, which is the bulk of the memory for
+            # long sequences. Mathematically identical to the explicit path.
+            attn_mask = attention_mask[:, :, :, : k.shape[-2]] if attention_mask is not None else None
+            attn_output = F.scaled_dot_product_attention(
+                q, k_expanded, v_expanded,
+                attn_mask=attn_mask,
+                is_causal=attn_mask is None,
+                dropout_p=self.attention_dropout if self.training else 0.0,
+                scale=self.head_dim ** -0.5,
+            )
+        attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, -1)
+        attn_output = self.o_proj(attn_output)
+        return attn_output, present
+
+
+# ---------------------------------------------------------------------------
+# Decoder layer (sandwich norm)
+# ---------------------------------------------------------------------------
+
+
+class BaberuDecoderLayer(nn.Module):
+    """One layer of the causal decoder.
+
+    With ``sandwich_norm=True``, both pre- and post-RMSNorm wrap each
+    sub-block (attention and FFN), in the style of Gemma 2.
+    """
+
+    def __init__(self, config: BaberuOCRConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.sandwich = config.sandwich_norm
+        self.input_norm = BaberuRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attn = BaberuAttention(config, layer_idx)
+        self.post_attn_norm = BaberuRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.pre_ffn_norm = BaberuRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp = BaberuMLP(config)
+        self.post_ffn_norm = BaberuRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+        residual = hidden_states
+        x = self.input_norm(hidden_states)
+        x, present = self.attn(
+            hidden_states=x,
+            cos=cos,
+            sin=sin,
+            attention_mask=attention_mask,
+            past_key_value=past_key_value,
+            use_cache=use_cache,
+        )
+        if self.sandwich:
+            x = self.post_attn_norm(x)
+        hidden_states = residual + x
+
+        residual = hidden_states
+        x = self.pre_ffn_norm(hidden_states)
+        x = self.mlp(x)
+        if self.sandwich:
+            x = self.post_ffn_norm(x)
+        hidden_states = residual + x
+        return hidden_states, present
+
+
+# ---------------------------------------------------------------------------
+# Base preTrained model
+# ---------------------------------------------------------------------------
+
+
+class BaberuPreTrainedModel(PreTrainedModel):
+    config_class = BaberuOCRConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["BaberuDecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+# ---------------------------------------------------------------------------
+# Backbone (embeddings + decoder layers + final norm)
+# ---------------------------------------------------------------------------
+
+
+class BaberuModel(BaberuPreTrainedModel):
+    """Decoder backbone, language-only (no LM head, no vision).
+
+    Accepts either ``input_ids`` (token IDs that go through the embedding
+    table) or ``inputs_embeds`` (pre-embedded, used when prepending vision).
+    """
+
+    def __init__(self, config: BaberuOCRConfig):
+        super().__init__(config)
+        config.validate()
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id
+        )
+        self.layers = nn.ModuleList(
+            [BaberuDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
+        )
+        self.norm = BaberuRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = BaberuRotaryEmbedding(
+            config.head_dim, config.max_position_embeddings, theta=config.rope_theta
+        )
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.embed_tokens = value
+
+    @staticmethod
+    def _build_causal_mask(
+        q_len: int, kv_len: int, device: torch.device, dtype: torch.dtype
+    ) -> torch.Tensor:
+        """Lower-triangular additive mask: -inf above the diagonal aligned to kv_len."""
+        # Q has length q_len, K has length kv_len. The last q_len queries see the
+        # last q_len keys causally; earlier keys (from KV cache) are all visible.
+        mask = torch.full((q_len, kv_len), float("-inf"), device=device, dtype=dtype)
+        # Compute offset so that diagonal at (q, q + (kv_len - q_len)) is allowed.
+        offset = kv_len - q_len
+        for i in range(q_len):
+            mask[i, : i + 1 + offset] = 0.0
+        return mask  # [Tq, Tk]
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,  # [B, T] of 1/0
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else True
+
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Pass exactly one of input_ids or inputs_embeds.")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        bsz, q_len, _ = inputs_embeds.shape
+        device = inputs_embeds.device
+
+        # Normalize transformers Cache objects to a list of (k, v) tuples.
+        if past_key_values is not None and hasattr(past_key_values, "layers"):
+            layers = past_key_values.layers
+            past_key_values = (
+                [(l.keys, l.values) for l in layers if l.is_initialized]
+                if layers
+                else None
+            )
+            if past_key_values is not None and len(past_key_values) == 0:
+                past_key_values = None
+
+        past_len = _cache_length(past_key_values)
+        kv_len = past_len + q_len
+
+        if position_ids is None:
+            position_ids = torch.arange(past_len, kv_len, device=device).unsqueeze(0).expand(bsz, -1)
+
+        cos, sin = self.rotary_emb(inputs_embeds, position_ids)
+
+        # Causal mask. Reshape to broadcast across batch and heads.
+        causal_mask = self._build_causal_mask(q_len, kv_len, device, inputs_embeds.dtype)
+        attn_mask = causal_mask[None, None, :, :].expand(bsz, 1, q_len, kv_len).clone()
+        if attention_mask is not None:
+            # attention_mask: [B, kv_len] with 1 for keep, 0 for pad
+            pad_mask = (1 - attention_mask[:, None, None, :]).to(inputs_embeds.dtype) * float("-inf")
+            # Replace NaN from (-inf * 0)
+            pad_mask = torch.nan_to_num(pad_mask, nan=0.0)
+            attn_mask = attn_mask + pad_mask[:, :, :, :kv_len]
+
+        hidden_states = inputs_embeds
+        next_cache = [] if use_cache else None
+        all_hidden_states = [] if output_hidden_states else None
+
+        for i, layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states.append(hidden_states)
+            layer_past = past_key_values[i] if past_key_values is not None else None
+            if self.gradient_checkpointing and self.training:
+                hidden_states, present = torch.utils.checkpoint.checkpoint(
+                    layer.__call__,
+                    hidden_states, cos, sin, attn_mask, layer_past, use_cache,
+                    use_reentrant=False,
+                )
+            else:
+                hidden_states, present = layer(
+                    hidden_states=hidden_states,
+                    cos=cos,
+                    sin=sin,
+                    attention_mask=attn_mask,
+                    past_key_value=layer_past,
+                    use_cache=use_cache,
+                )
+            if use_cache:
+                next_cache.append(present)
+
+        hidden_states = self.norm(hidden_states)
+        if output_hidden_states:
+            all_hidden_states.append(hidden_states)
+
+        if not return_dict:
+            outputs = (hidden_states,)
+            if use_cache:
+                outputs = outputs + (tuple(next_cache),)
+            if output_hidden_states:
+                outputs = outputs + (tuple(all_hidden_states),)
+            return outputs
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=tuple(next_cache) if use_cache else None,
+            hidden_states=tuple(all_hidden_states) if output_hidden_states else None,
+            attentions=None,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Causal LM (text-only, for Step 1 pretraining)
+# ---------------------------------------------------------------------------
+
+
+def _apply_softcap(logits: torch.Tensor, cap: float) -> torch.Tensor:
+    return torch.tanh(logits / cap) * cap
+
+
+def _z_loss(logits: torch.Tensor, ignore_index: int = -100, labels: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """Auxiliary z-loss: -1e-4 * log(Z)^2 averaged over non-ignored positions.
+
+    Encourages logit normalizer log(Z) -> 0, stabilizing training. Apply on
+    pre-softcap logits to keep gradient signal meaningful.
+    """
+    if labels is not None:
+        mask = (labels != ignore_index).float()
+        log_z = torch.logsumexp(logits.float(), dim=-1)  # [B, T]
+        denom = mask.sum().clamp(min=1.0)
+        return ((log_z**2) * mask).sum() / denom
+    log_z = torch.logsumexp(logits.float(), dim=-1)
+    return (log_z**2).mean()
+
+
+@dataclass
+class BaberuCausalLMOutput(CausalLMOutputWithPast):
+    z_loss: Optional[torch.FloatTensor] = None
+
+
+class BaberuCausalLM(BaberuPreTrainedModel, GenerationMixin):
+    """Decoder-only causal LM. Used for Step 1 text pretraining."""
+
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(self, config: BaberuOCRConfig):
+        super().__init__(config)
+        self.model = BaberuModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head
+
+    def set_output_embeddings(self, value: nn.Linear) -> None:
+        self.lm_head = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaberuCausalLMOutput]:
+        return_dict = return_dict if return_dict is not None else True
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+        )
+        hidden_states = outputs.last_hidden_state
+        logits_raw = self.lm_head(hidden_states)
+        logits = (
+            _apply_softcap(logits_raw, self.config.final_logit_softcap)
+            if self.config.final_logit_softcap
+            else logits_raw
+        )
+
+        loss: Optional[torch.Tensor] = None
+        z_loss: Optional[torch.Tensor] = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            ce = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+            if self.config.z_loss_weight and self.config.z_loss_weight > 0:
+                # Use pre-softcap logits for z-loss to retain gradient.
+                shift_raw = logits_raw[:, :-1, :].contiguous()
+                z_loss = _z_loss(shift_raw, ignore_index=-100, labels=shift_labels)
+                loss = ce + self.config.z_loss_weight * z_loss
+            else:
+                loss = ce
+
+        if not return_dict:
+            output = (logits,) + ((outputs.past_key_values,) if outputs.past_key_values is not None else ())
+            return ((loss,) + output) if loss is not None else output
+
+        return BaberuCausalLMOutput(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=None,
+            z_loss=z_loss,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[Tuple] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ) -> dict:
+        # Determine how many positions are already in the cache. An empty
+        # DynamicCache is not None but has past_len=0, so we must inspect.
+        past_len = _cache_length(past_key_values)
+        # When KV cache has content, only the last token needs to be fed.
+        if past_len > 0:
+            input_ids = input_ids[:, -1:]
+        model_inputs = {"input_ids": input_ids, "inputs_embeds": None}
+        if inputs_embeds is not None and past_len == 0:
+            model_inputs = {"input_ids": None, "inputs_embeds": inputs_embeds}
+        model_inputs.update(
+            {
+                "past_key_values": past_key_values,
+                "attention_mask": attention_mask,
+                "use_cache": kwargs.get("use_cache", True),
+            }
+        )
+        return model_inputs
+
+
+# ---------------------------------------------------------------------------
+# Vision projector
+# ---------------------------------------------------------------------------
+
+
+class BaberuVisionProjector(nn.Module):
+    """MLP projector: vision_hidden -> hidden_size."""
+
+    def __init__(self, config: BaberuOCRConfig):
+        super().__init__()
+        self.linear1 = nn.Linear(config.vision_hidden_size, config.hidden_size, bias=True)
+        self.act = nn.GELU()
+        self.linear2 = nn.Linear(config.hidden_size, config.hidden_size, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.linear2(self.act(self.linear1(x)))
+
+
+# ---------------------------------------------------------------------------
+# Full OCR model: vision encoder + projector + causal LM
+# ---------------------------------------------------------------------------
+
+
+class BaberuOCRModel(BaberuPreTrainedModel, GenerationMixin):
+    """Full OCR model. Reads an image, generates text.
+
+    Input layout fed to the decoder is:
+        [vision tokens (projected DINOv2 outputs) | BOS | text tokens | EOS]
+
+    The vision tokens are computed once per image (per generate() call) and
+    cached in past_key_values along with all decoded text steps.
+    """
+
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(self, config: BaberuOCRConfig):
+        super().__init__(config)
+        if config.vision_model_name is None:
+            raise ValueError(
+                "BaberuOCRModel requires config.vision_model_name. For text-only "
+                "pretraining, use BaberuCausalLM instead."
+            )
+        # Build the vision encoder from its *config* (structure only, no weight
+        # download), which is safe under `from_pretrained`'s meta-device init —
+        # calling `AutoModel.from_pretrained` here breaks loading via
+        # `AutoModel.from_pretrained(..., trust_remote_code=True)`. When this
+        # model is built fresh (training, params on a real device) we populate
+        # the pretrained DINOv2 weights; when loaded from a checkpoint the params
+        # start on meta and the vision weights come from the checkpoint, so we
+        # skip the (redundant) download.
+        vision_config = AutoConfig.from_pretrained(config.vision_model_name)
+        self.vision_encoder = AutoModel.from_config(vision_config)
+        if next(self.vision_encoder.parameters()).device.type != "meta":
+            pretrained = AutoModel.from_pretrained(config.vision_model_name)
+            self.vision_encoder.load_state_dict(pretrained.state_dict())
+            del pretrained
+        if config.freeze_vision_encoder:
+            for p in self.vision_encoder.parameters():
+                p.requires_grad = False
+            self.vision_encoder.eval()
+        self.projector = BaberuVisionProjector(config)
+        self.model = BaberuModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head
+
+    def _encode_image(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """[B, 3, H, W] -> [B, N_vision, hidden_size]."""
+        out = self.vision_encoder(pixel_values=pixel_values, return_dict=True)
+        # DINOv2 returns last_hidden_state of shape [B, N+1, vision_hidden_size]
+        # where the +1 is the CLS token. We drop it for OCR.
+        feats = out.last_hidden_state[:, 1:, :]
+        return self.projector(feats)
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        input_ids: Optional[torch.LongTensor] = None,  # text tokens (no vision prefix)
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple] = None,
+        labels: Optional[torch.LongTensor] = None,  # text-only labels; we shift inside
+        use_cache: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaberuCausalLMOutput]:
+        return_dict = return_dict if return_dict is not None else True
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        # On the first forward we must prepend the vision prefix; once it is in
+        # the KV cache we only feed new text. Branch on cache *length*, not
+        # ``is None``: HF ``generate`` passes an empty (but non-None) Cache on
+        # the first step, so an ``is None`` check would skip the vision prefix
+        # entirely and every image would decode the same unconditional text.
+        if _cache_length(past_key_values) == 0:
+            if pixel_values is None:
+                raise ValueError("pixel_values is required on the first forward (no KV cache).")
+            vision_embeds = self._encode_image(pixel_values)  # [B, Nv, H]
+            text_embeds = (
+                self.model.embed_tokens(input_ids) if input_ids is not None else None
+            )
+            if text_embeds is not None:
+                inputs_embeds = torch.cat([vision_embeds, text_embeds], dim=1)
+            else:
+                inputs_embeds = vision_embeds
+            # Build attention mask that covers both.
+            if attention_mask is not None and input_ids is not None:
+                vision_mask = torch.ones(
+                    vision_embeds.shape[:2], device=vision_embeds.device, dtype=attention_mask.dtype
+                )
+                full_mask = torch.cat([vision_mask, attention_mask], dim=1)
+            else:
+                full_mask = None
+            outputs = self.model(
+                inputs_embeds=inputs_embeds,
+                attention_mask=full_mask,
+                past_key_values=None,
+                use_cache=use_cache,
+                output_hidden_states=output_hidden_states,
+                return_dict=True,
+            )
+            n_vision = vision_embeds.shape[1]
+        else:
+            # KV cache already holds the 256 vision tokens + prior text; only the
+            # new text token goes in. HF's attention_mask counts text tokens only
+            # (it is unaware of the vision prefix), so its length won't match
+            # kv_len — drop it. There is no padding in single-sequence greedy
+            # generation, so the internal causal mask is sufficient; position_ids
+            # default to arange(past_len, kv_len), giving the new token its true
+            # post-vision position.
+            outputs = self.model(
+                input_ids=input_ids,
+                attention_mask=None,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                output_hidden_states=output_hidden_states,
+                return_dict=True,
+            )
+            n_vision = 0
+
+        hidden_states = outputs.last_hidden_state
+        logits_raw = self.lm_head(hidden_states)
+        logits = (
+            _apply_softcap(logits_raw, self.config.final_logit_softcap)
+            if self.config.final_logit_softcap
+            else logits_raw
+        )
+
+        loss: Optional[torch.Tensor] = None
+        z_loss: Optional[torch.Tensor] = None
+        if labels is not None:
+            # The logits include vision-prefix positions; slice them out for the loss.
+            text_logits = logits[:, n_vision:, :]
+            text_logits_raw = logits_raw[:, n_vision:, :]
+            shift_logits = text_logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            ce = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+            if self.config.z_loss_weight and self.config.z_loss_weight > 0:
+                shift_raw = text_logits_raw[:, :-1, :].contiguous()
+                z_loss = _z_loss(shift_raw, ignore_index=-100, labels=shift_labels)
+                loss = ce + self.config.z_loss_weight * z_loss
+            else:
+                loss = ce
+
+        if not return_dict:
+            output = (logits,) + ((outputs.past_key_values,) if outputs.past_key_values is not None else ())
+            return ((loss,) + output) if loss is not None else output
+
+        return BaberuCausalLMOutput(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=None,
+            z_loss=z_loss,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[Tuple] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        **kwargs,
+    ) -> dict:
+        past_len = _cache_length(past_key_values)
+        if past_len > 0:
+            input_ids = input_ids[:, -1:]
+            pixel_values = None  # already encoded in the cache
+        return {
+            "input_ids": input_ids,
+            "pixel_values": pixel_values,
+            "past_key_values": past_key_values,
+            "attention_mask": attention_mask,
+            "use_cache": kwargs.get("use_cache", True),
+        }

onnx/decoder_prefill_int8.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f8d28d4e5e8b6c7f3a4510663aa0a9199a7eec06bcb1dc0afb3ca571494e97f
+size 35133596

onnx/decoder_step_int8.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2eaf792ca70b67e55140ef0bcfbfc3c3b1e07d72ca277dcc46aae27c81f27ac8
+size 33929034

onnx/vision_fp16.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ce333804846b23f1983376efe363c164d4624ecc6bd16f99a15a859d8423117
+size 172917304

onnx/vision_int4.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9526bcf99eb8de805df7ea2f66b385cfcbcb5dff0e99a5dcfb02bdb2d4182bde
+size 52293486

onnx_infer.py

@@ -0,0 +1,117 @@
+"""Self-contained ONNX inference for Baberu OCR — no PyTorch, no transformers,
+no custom package. Just `onnxruntime`, `numpy`, `pillow`.
+
+  pip install onnxruntime numpy pillow
+  python onnx_infer.py --onnx-dir onnx --tokenizer tokenizer --image crop.jpg
+
+Runs the three ONNX graphs (vision -> decoder prefill -> decoder step, KV cache)
+in a pure-numpy host loop with the published decode (greedy + repetition_penalty
+1.2 + symbol-aware content-run cap 12), so you get the same text as the PyTorch
+model. Pick your size tier by which vision file you point at:
+
+  vision_fp16.onnx + decoder_*_int8.onnx  = 242 MB  (lossless, exact published nCER)
+  vision_int4.onnx + decoder_*_int8.onnx  = 121 MB  (smallest; ~+0.002 nCER)
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import unicodedata
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+
+# DINOv2 / ImageNet preprocessing (do_center_crop=False, resize whole crop to 224).
+_MEAN = np.array([0.485, 0.456, 0.406], np.float32)
+_STD = np.array([0.229, 0.224, 0.225], np.float32)
+_PAST = [f"past_k{i}" for i in range(6)] + [f"past_v{i}" for i in range(6)]
+
+
+def preprocess(image: Image.Image) -> np.ndarray:
+    img = image.convert("RGB").resize((224, 224), Image.BICUBIC)
+    x = (np.asarray(img, np.float32) / 255.0 - _MEAN) / _STD
+    return x.transpose(2, 0, 1)[None]  # [1,3,224,224]
+
+
+class Vocab:
+    """Char-level vocab. ids 0..3 = <pad>/<bos>/<eos>/<unk>; id>=4 -> charset[id-4]."""
+
+    def __init__(self, vocab_json: Path):
+        charset = json.loads(Path(vocab_json).read_text(encoding="utf-8"))
+        self.id2ch = {0: "", 1: "", 2: "", 3: ""}
+        for i, ch in enumerate(charset):
+            self.id2ch[i + 4] = ch
+        self.bos, self.eos = 1, 2
+        self.vocab_size = len(charset) + 4
+        self.content_ids = {
+            i + 4 for i, ch in enumerate(charset)
+            if len(ch) == 1 and ch not in "ーｰ〜~" and unicodedata.category(ch)[0] in "LN"
+        }
+
+    def decode(self, ids) -> str:
+        return "".join(self.id2ch.get(i, "") for i in ids if i >= 4)
+
+
+class BaberuOnnxOCR:
+    def __init__(self, onnx_dir, tokenizer_dir, vision="vision_fp16.onnx"):
+        d = Path(onnx_dir)
+        so = ort.SessionOptions()
+        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        mk = lambda p: ort.InferenceSession(str(p), so, providers=["CPUExecutionProvider"])
+        self.vis = mk(d / vision)
+        self.pre = mk(d / "decoder_prefill_int8.onnx")
+        self.stp = mk(d / "decoder_step_int8.onnx")
+        self.vocab = Vocab(Path(tokenizer_dir) / "vocab.json")
+
+    def __call__(self, image, max_new_tokens=128, repetition_penalty=1.2, max_content_run=12):
+        v = self.vocab
+        vis = self.vis.run(["vision_embeds"], {"pixel_values": preprocess(image)})[0]
+        out = self.pre.run(None, {"vision_embeds": vis, "input_ids": np.array([[v.bos]], np.int64)})
+        logits, present = out[0][0, -1].astype(np.float64), out[1:]
+        seq, toks, pos = [v.bos], [], vis.shape[1] + 1
+        for _ in range(max_new_tokens):
+            if repetition_penalty != 1.0:
+                for tid in set(seq):
+                    s = logits[tid]
+                    logits[tid] = s * repetition_penalty if s < 0 else s / repetition_penalty
+            if max_content_run and toks and toks[-1] in v.content_ids:
+                last = toks[-1]
+                run = 0
+                for t in reversed(toks):
+                    if t == last:
+                        run += 1
+                    else:
+                        break
+                if run >= max_content_run:
+                    logits[last] = -np.inf
+            nxt = int(np.argmax(logits))
+            if nxt == v.eos:
+                break
+            toks.append(nxt)
+            seq.append(nxt)
+            if len(toks) >= max_new_tokens:
+                break
+            feed = {"input_ids": np.array([[nxt]], np.int64), "position_ids": np.array([[pos]], np.int64)}
+            feed.update({nm: t for nm, t in zip(_PAST, present)})
+            out = self.stp.run(None, feed)
+            logits, present = out[0][0, -1].astype(np.float64), out[1:]
+            pos += 1
+        return v.decode(toks)
+
+
+def main():
+    ap = argparse.ArgumentParser(description="Self-contained ONNX OCR for a manga bubble crop.")
+    ap.add_argument("--onnx-dir", default="onnx")
+    ap.add_argument("--tokenizer", default="tokenizer")
+    ap.add_argument("--vision", default="vision_fp16.onnx",
+                    help="vision_fp16.onnx (lossless) or vision_int4.onnx (smallest)")
+    ap.add_argument("--image", required=True)
+    args = ap.parse_args()
+    ocr = BaberuOnnxOCR(args.onnx_dir, args.tokenizer, args.vision)
+    print(ocr(Image.open(args.image)))
+
+
+if __name__ == "__main__":
+    main()

quantization-results.md

@@ -0,0 +1,61 @@
+# Baberu OCR — ONNX quantization size×accuracy sweep (2026-06-25)
+
+Model: custom 115M (DINOv2-ViT-B/14 86M + 6L GQA decoder 26M), ckpt step295000.
+Eval: Manga109 independent GT n=2000, headline decode (greedy + rep-penalty 1.2 +
+content-run cap 12). Baseline torch fp32 = nCER 0.0871 / lCER 0.0345.
+
+## On-disk sizes (ONNX, MB)
+
+| component | fp32 | fp16 | int8(dyn) | int4(NBits) |
+|---|---|---|---|---|
+| vision (86M) | 349.2 | 172.9 | 91.7 | 52.3 |
+| decoder prefill | 134.0 | — | 35.1 | 45.3 |
+| decoder step | 132.8 | — | 33.9 | 44.1 |
+| decoder total | 266.8 | — | 69.0 | 89.4 |
+
+Notes: decoder ONNX = 2 graphs, each duplicates the tied embedding/lm_head (tie not
+preserved in ONNX) → fp32 267MB. int4 on the decoder is BIGGER than int8 (89 vs 69MB):
+MatMulNBits adds per-block fp scales and leaves embedding (Gather) unquantized → int4
+only helps the vision tower.
+
+## End-to-end accuracy (n=2000)
+
+| config | vision | decoder | bundle MB | nCER | lCER | nAcc |
+|---|---|---|---|---|---|---|
+| A torch fp32 (baseline) | — | — | (451 st) | 0.0871 | 0.0345 | 0.6965 |
+| B pipe fp32 (ONNX) | 349.2 | 266.8 | 616.0 | 0.0871 | 0.0345 | 0.6965 |
+| C vis fp16 + dec fp32 (≈current ship) | 172.9 | 266.8 | 439.7 | 0.0871 | 0.0345 | 0.6965 |
+| D vis int8 + dec fp32 | 91.7 | 266.8 | 358.5 | 0.0887 | 0.0364 | 0.6935 |
+| E vis int4 + dec fp32 | 52.3 | 266.8 | 319.1 | 0.0881 | 0.0350 | 0.6940 |
+| F vis fp32 + dec int8 | 349.2 | 69.0 | 418.2 | 0.0867 | 0.0338 | 0.6960 |
+| G vis int8 + dec int8 (full int8) | 91.7 | 69.0 | 160.7 | 0.0890 | 0.0363 | 0.6905 |
+| H vis int4 + dec int8 (smallest) | 52.3 | 69.0 | 121.3 | 0.0893 | 0.0365 | 0.6920 |
+| Z manga-ocr (teacher) | — | — | ~444 | 0.0954 | 0.0422 | 0.6885 |
+
+vis fp16 ≡ vis fp32 (byte-identical preds), so vis fp16 + dec int8 = F = 0.0867 @ 241.9 MB.
+
+## Findings
+
+1. **decoder int8 is FREE.** F (dec int8) = nCER 0.0867 ≈ baseline 0.0871 (within noise,
+   fractionally better). This OVERTURNS the "keep decoder fp32" rule. That rule came from
+   fp16 *logit parity* drift (max|diff| 5-9), but logit-parity ≠ CER: greedy argmax is
+   robust to absolute logit noise, and int8 per-channel scales preserve the argmax-relevant
+   structure (unlike fp16's limited mantissa on the ±30 soft-capped logits). Decoder 267→69 MB free.
+
+2. **The only cost is vision quantization, and it's tiny.** vision fp16 = lossless (0.0871).
+   vision int4 weight-only = +0.0010 (0.0881) AND smaller than int8-dynamic. vision
+   int8-DYNAMIC = +0.0016 (0.0887) — worse than int4 because dynamic also quantizes
+   activations (adds noise); int4 here is weight-only (activations fp32).
+
+3. **Every quantized config still beats the teacher manga-ocr (0.0954).**
+
+## Recommended deployment tiers
+
+- **Zero-regret (best): vis fp16 + dec int8 = 242 MB, nCER 0.0867** — half the current
+  440MB ship at same/better accuracy. Decoder int8 free + vision fp16 lossless.
+- **Balanced full-int8: 161 MB, nCER 0.0890** — 3.8× smaller than fp32-ONNX.
+- **Smallest: vis int4 + dec int8 = 121 MB, nCER 0.0893** — 5.1× smaller, still beats manga-ocr.
+
+Artifacts written to /Volumes/junkanki/onnx-export/: vision_{int8,int4}.onnx,
+decoder_{prefill,step}_{int8,int4}.onnx. Scripts: experiments/{quantize_for_size,
+eval_onnx_pipeline}.py + eval_ocr_cer.py --vision-onnx.

tokenization_baberu.py

@@ -0,0 +1,185 @@
+"""Character-level tokenizer for baberu-ocr v2.
+
+Loads a charset file (JSON array of single characters), prepends special
+tokens, and exposes a HuggingFace-compatible PreTrainedTokenizer.
+
+Token layout:
+    0: <pad>
+    1: <bos>
+    2: <eos>
+    3: <unk>
+    4..N: characters from charset (e.g., '\\n', ' ', '!', ..., 'あ', '世', ...)
+
+Designed for multilingual OCR (JA + ZH + EN). Pure character-level: every
+codepoint maps to one token. Unicode codepoints above U+FFFF are stored as
+single Python str elements and tokenized as one token (verified for chars
+like '𠝹' U+2057D).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+from transformers import PreTrainedTokenizer
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.json"}
+
+DEFAULT_SPECIAL_TOKENS = {
+    "pad_token": "<pad>",
+    "bos_token": "<bos>",
+    "eos_token": "<eos>",
+    "unk_token": "<unk>",
+}
+
+
+class BaberuTokenizer(PreTrainedTokenizer):
+    """Character-level tokenizer with prepended special tokens.
+
+    Args:
+        vocab_file: Path to a JSON array of characters (e.g., baberu_charset.txt).
+        pad_token / bos_token / eos_token / unk_token: Special token strings.
+
+    Example:
+        >>> tok = BaberuTokenizer(vocab_file="data/baberu_charset.txt")
+        >>> tok("こんにちは Hello").input_ids
+        [1, 1234, 5678, ..., 4, 'H'_id, ...]  # includes <bos> ... (no <eos> by default)
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file: Optional[str] = None,
+        pad_token: str = DEFAULT_SPECIAL_TOKENS["pad_token"],
+        bos_token: str = DEFAULT_SPECIAL_TOKENS["bos_token"],
+        eos_token: str = DEFAULT_SPECIAL_TOKENS["eos_token"],
+        unk_token: str = DEFAULT_SPECIAL_TOKENS["unk_token"],
+        add_bos_token: bool = True,
+        add_eos_token: bool = True,
+        **kwargs,
+    ):
+        if vocab_file is None:
+            raise ValueError("vocab_file must be provided")
+        if not Path(vocab_file).is_file():
+            raise FileNotFoundError(f"vocab_file not found: {vocab_file}")
+
+        charset = json.loads(Path(vocab_file).read_text(encoding="utf-8"))
+        if not isinstance(charset, list):
+            raise ValueError("vocab_file must contain a JSON array of strings")
+
+        # Coerce special tokens to plain strings — from_pretrained may pass
+        # AddedToken instances which would break string-keyed dict lookups.
+        pad_str = str(pad_token)
+        bos_str = str(bos_token)
+        eos_str = str(eos_token)
+        unk_str = str(unk_token)
+
+        # Build vocab with special tokens first, then charset.
+        self._token_to_id: dict[str, int] = {
+            pad_str: 0,
+            bos_str: 1,
+            eos_str: 2,
+            unk_str: 3,
+        }
+        for ch in charset:
+            if ch in self._token_to_id:
+                continue
+            self._token_to_id[ch] = len(self._token_to_id)
+        self._id_to_token: dict[int, str] = {i: t for t, i in self._token_to_id.items()}
+
+        self.vocab_file = vocab_file
+        self.add_bos_token = add_bos_token
+        self.add_eos_token = add_eos_token
+
+        super().__init__(
+            pad_token=pad_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            unk_token=unk_token,
+            add_bos_token=add_bos_token,
+            add_eos_token=add_eos_token,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self._token_to_id)
+
+    def get_vocab(self) -> dict[str, int]:
+        vocab = dict(self._token_to_id)
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def _tokenize(self, text: str, **kwargs) -> List[str]:
+        return list(text)
+
+    def _convert_token_to_id(self, token: str) -> int:
+        token_str = str(token)
+        if token_str in self._token_to_id:
+            return self._token_to_id[token_str]
+        # Fallback to unk; tolerate self.unk_token being an AddedToken.
+        return self._token_to_id.get(str(self.unk_token), 3)
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return self._id_to_token.get(index, str(self.unk_token))
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        return "".join(tokens)
+
+    def build_inputs_with_special_tokens(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None,
+    ) -> List[int]:
+        """Wrap a sequence with [BOS] ... [EOS] tokens.
+
+        For OCR we only ever have a single sequence (token_ids_1 unused).
+        """
+        bos = [self.bos_token_id] if self.add_bos_token else []
+        eos = [self.eos_token_id] if self.add_eos_token else []
+        result = bos + token_ids_0 + eos
+        if token_ids_1 is not None:
+            result = result + token_ids_1 + eos
+        return result
+
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None,
+        already_has_special_tokens: bool = False,
+    ) -> List[int]:
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0,
+                token_ids_1=token_ids_1,
+                already_has_special_tokens=True,
+            )
+        bos = [1] if self.add_bos_token else []
+        eos = [1] if self.add_eos_token else []
+        mask = bos + [0] * len(token_ids_0) + eos
+        if token_ids_1 is not None:
+            mask += [0] * len(token_ids_1) + eos
+        return mask
+
+    def save_vocabulary(
+        self, save_directory: str, filename_prefix: Optional[str] = None
+    ) -> Tuple[str]:
+        save_dir = Path(save_directory)
+        save_dir.mkdir(parents=True, exist_ok=True)
+        out_path = save_dir / (
+            (filename_prefix + "-" if filename_prefix else "")
+            + self.vocab_files_names["vocab_file"]
+        )
+        # Save only the charset (special tokens are recreated on load).
+        special_ids = {0, 1, 2, 3}
+        charset = [
+            self._id_to_token[i]
+            for i in range(len(self._id_to_token))
+            if i not in special_ids
+        ]
+        out_path.write_text(json.dumps(charset, ensure_ascii=False, indent=2))
+        return (str(out_path),)

tokenizer/tokenizer_config.json

@@ -0,0 +1,44 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<bos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<eos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "backend": "custom",
+  "bos_token": "<bos>",
+  "eos_token": "<eos>",
+  "is_local": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "tokenizer_class": "BaberuTokenizer",
+  "unk_token": "<unk>"
+}