hayai-ocr-v2-onnx

Hayai OCR v2.1 — ONNX Runtime

This repository provides optimized ONNX Runtime exports (FP32, FP16, and INT8 Dynamic) for Hayai OCR v2.1—a lightweight (~150M parameter) vision-to-text model for dense Japanese, Chinese, Korean, and English text.

Looking for PyTorch weights, training details, or datasets?
Visit the main repository: JustANormalTinkerer/hayai-ocr-v2.

---

File Structure

.
├── config.json
├── configuration_hayai.py
├── modeling_hayai.py
├── tokenizer.json
├── tokenizer_config.json
├── position_base.npy                # 16x16x768 base grid for NaFlex position interpolation
├── export_onnx.py                   # PyTorch -> ONNX export script
├── inference_onnx.py                # Standalone ONNX Runtime inference script
└── onnx/
    ├── hayai_encoder.onnx               # FP32 Encoder (0.9 MB + 327 MB .data)
    ├── hayai_encoder_fp16.onnx          # FP16 Encoder (~165 MB)
    ├── hayai_encoder_dynamic_quant.onnx # INT8 Dynamic Quantized Encoder (~83 MB)
    ├── hayai_decoder.onnx               # FP32 Decoder (1.4 MB + 208 MB .data)
    ├── hayai_decoder_fp16.onnx          # FP16 Decoder (~106 MB)
    ├── hayai_decoder_dynamic_quant.onnx # INT8 Dynamic Quantized Decoder (~54 MB)
    └── position_base.npy

---

Quick Start (ONNX Runtime)

Inference runs purely on ONNX Runtime without requiring PyTorch.

1. Installation

pip install onnxruntime tokenizers Pillow numpy

(For GPU acceleration, install onnxruntime-gpu instead).

---

2. CLI Inference

# Run FP32 (default)
python inference_onnx.py --image example.png

# Run FP16 or INT8 Dynamic
python inference_onnx.py --image example.png --precision fp16
python inference_onnx.py --image example.png --precision quant

# Batch process a directory
python inference_onnx.py --images ./crops --out results.jsonl

---

3. Python API

import numpy as np
import onnxruntime as ort
import inference_onnx as hayai_onnx

# 1. Load tokenizer and position base
base_grid = np.load("onnx/position_base.npy")
tok, bos_id, eos_id, pad_id = hayai_onnx.load_tokenizer("tokenizer.json")

# 2. Start ONNX inference sessions (CPU or CUDA)
providers = ["CPUExecutionProvider"]
enc_session = ort.InferenceSession("onnx/hayai_encoder.onnx", providers=providers)
dec_session = ort.InferenceSession("onnx/hayai_decoder.onnx", providers=providers)

# 3. Perform OCR
text, _ = hayai_onnx.ocr(
    "example.png",
    encoder_session=enc_session,
    decoder_session=dec_session,
    position_base=base_grid,
    tokenizer=tok,
    bos_id=bos_id,
    eos_id=eos_id,
    pad_id=pad_id,
)

print("Recognized Text:", text)

---

Technical Details

• Host-Computed Positional Embeddings: SigLIP2 NaFlex grid interpolation and 2D mRoPE frequencies are computed host-side in NumPy (cached with an LRU cache in inference_onnx.py, ~0.5 ms). This keeps the ONNX graphs clean and dynamic across batch size ($B$), visual sequence length ($L$), and token generation steps ($N$).
• Dynamic Axes Validated: Encoder handles dynamic batch; Decoder handles dynamic batch, sequence length $N$, and context length $L$.

Re-exporting ONNX Models

To regenerate or re-quantize the ONNX graphs from local PyTorch weights:

pip install torch transformers safetensors onnx onnxruntime

# Export FP32 + FP16
python export_onnx.py --precision fp32 --out_dir onnx --opset 18

# Export with INT8 Dynamic Quantization
python export_onnx.py --quant dynamic --out_dir onnx

---

Text Normalization

For consistent post-processing and CJK evaluation:

import re
import unicodedata

def normalize_text(text: str) -> str:
    text = unicodedata.normalize("NFKC", str(text))
    text = re.sub(r'[\r\n\t]+', ' ', text)
    cjk_pattern = r'([\u4e00-\u9fff\u3040-\u30ff\u3400-\u4dbf\uac00-\ud7af])'
    text = re.sub(f'{cjk_pattern}\\s+{cjk_pattern}', r'\1\2', text)
    return re.sub(r'\s+', ' ', text).strip()