IndexErrorThe/unirec-0.1b-coreml

UniRec-0.1B — Core ML

Core ML conversion of topdu/unirec-0.1b, a unified recogniser for printed text, mathematical formulas and tables. Packaged for on-device use by MathOCR, a macOS app that turns PDF pages without a text layer into Markdown with LaTeX.

Nothing here is a new model. The weights are the original authors'; this repository only changes the format.

Contents

File	Size	Input	Output
UniRecEncoder.mlpackage	82 MB	pixel_values [1,3,H,W]	cross_k, cross_v [6,1,6,S,128]
UniRecDecoder.mlpackage	191 MB	one token + cross-attention K/V	logits [1,56371]

Both float16, minimum_deployment_target = macOS 15.

The split follows the maintainers' own ONNX export: the encoder also computes the decoder's cross-attention key and value projections, which depend only on the image and so are evaluated once per region rather than once per token.

The decoder is stateful — its self-attention KV cache lives on the Core ML side rather than being passed in and out each step.

Requirements and limits

Run with compute units restricted to CPU and GPU.

let configuration = MLModelConfiguration()
configuration.computeUnits = .cpuAndGPU

This is not a performance preference. With the Neural Engine enabled, the encoder is executed incorrectly: at a 192×640 input its output deviates from the reference by 56%, against 0.19% on CPU+GPU. FocalSVTR's focal modulation uses depthwise convolutions with kernels up to 15×15, which appears to be the cause. The failure is silent — the model returns plausible but wrong text rather than an error.

Other limits, fixed at conversion time:

• 512 encoder positions (cross_len marks how many are real; pad the rest). That is a region of roughly 960×544 px. Larger regions must be split.
• 512 generated tokens — the self-attention cache size.

Preprocessing

Identical to the original model:

1. RGB.
2. Fit inside 960×1408 preserving aspect ratio. Images already smaller are not enlarged.
3. Round both sides down to a multiple of 64, minimum 64.
4. Bicubic resample.
5. Scale to [0,1], then (x - 0.5) / 0.5.
6. NCHW, float32.

Decoding

Greedy. Start from bos = 0, stop at eos = 2. The position index follows M2M100's convention: position = pad_token_id + 1 + step, i.e. 2 + step.

Token IDs map to strings through unirec_tokenizer_mapping.json in topdu/unirec_0_1b_onnx. Formulas are emitted as \( … \) and \[ … \]; tables as HTML.

Verification

Checked against the maintainers' ONNX export on six rendered crops — an inline formula, three display equations, a nested radical and a multi-line paragraph. All six produce token-for-token identical greedy sequences.

On an M1: encoder 29–217 ms per region, decoder ~10 ms per token.

The conversion and verification scripts are in the MathOCR repository under tools/unirec_coreml/.

Licence and attribution

Apache-2.0, inherited from the original model.

• Original model and weights: topdu/unirec-0.1b
• Source implementation: Topdu/OpenOCR

Changes made in this redistribution, as Apache-2.0 §4(b) requires:

1. Exported from PyTorch to Core ML (.mlpackage, ML Program), float16.
2. The decoding step was reimplemented with a fixed-size KV cache updated by a masked write, because HuggingFace's Cache class does not trace to a static graph. Same weights, same arithmetic — verified against the original.
3. The encoder additionally returns the decoder's cross-attention K/V projections, matching the upstream ONNX export's split.