Reza2kn's picture
Add native Swift CoreML runtime and assets
d82e1f6 verified
|
Raw
History Blame Contribute Delete
5.49 kB
metadata
base_model: datalab-to/surya-ocr-2
base_model_relation: quantized
library_name: coremltools
license: openrail
pipeline_tag: image-text-to-text
tags:
  - coreml
  - coremltools
  - apple-silicon
  - ios
  - macos
  - ocr
  - document-ai
  - surya
  - quantized
  - int8
  - vision-language

Surya OCR 2 CoreML Runtime

This repository contains an early CoreML runtime bundle derived from datalab-to/surya-ocr-2 at source commit 3b3d4cdf88d6928b0acdc75181b13206ea67c4a3.

It is intended for native Apple OCR experiments and future iOS/macOS demo-app work. This revision includes a SwiftPM runtime package that runs the fixed-shape OCR canary path without Python.

Included packages

File Purpose Precision / quantization Shape contract
surya_vision_fp16.mlpackage Vision tower CoreML FP16 compute pixel_values [1024,1536] -> image_embeds [256,1024]
surya_vision_int8.mlpackage Quantized vision tower CoreML linear INT8 weight compression pixel_values [1024,1536] -> image_embeds [256,1024]
surya_prefill_fp16_seq300_cache512.mlpackage Language prefill, logits, and initial cache CoreML FP16 compute fixed prefill length 300, cache length 512
surya_decode_step_fp16_cache512.mlpackage One-token cached decode step CoreML FP16 compute one token at a time, cache length 512

Processor/tokenizer assets are included under processor/. Native constants and embedding assets are under native_assets/. Validation JSONs are under validation/.

Native Swift package

The SwiftPM package lives in native/SuryaCoreMLRuntime.

import SuryaCoreMLRuntime

let runtime = try SuryaCoreMLRuntime(modelDirectory: modelDirectoryURL)
let result = try runtime.generate(image: cgImage, maxNewTokens: 128, useInt8Vision: true)
print(result.text ?? "")

The package currently handles fixed 512x512 Qwen image preprocessing, INT8/FP16 vision selection, fixed OCR prompt embeddings, image placeholder insertion, precomputed RoPE tables, generated-token embedding lookup, KV cache insertion, CoreML prefill/decode, and basic token decoding.

Current validation

All validation below was run on a Mac Studio with CoreML package prediction, using the canary prompt/image:

OCR this image to HTML.

The canary image contains:

Invoice 123
Total $42.00
Gate Result
Prefill parity before CoreML export native/custom first token 1039; logits max diff 2.6702880859375e-05
Prefill CoreML smoke Torch/CoreML first token 1039; logits max diff 0.3057253360748291; mean diff 0.03853870555758476
Decode CoreML iterative smoke 9/9 tokens match native; text <p>Invoice
CoreML prefill -> CoreML decode 9/9 tokens match native; text <p>Invoice
CoreML FP16 vision -> CoreML prefill -> CoreML decode 9/9 tokens match native; text <p>Invoice
CoreML INT8 vision -> CoreML prefill -> CoreML decode 9/9 tokens match native; text <p>Invoice
Native Swift image -> CoreML INT8 vision -> CoreML prefill -> CoreML decode 9/9 tokens match native; text <p>Invoice

The INT8 vision package has mean absolute diff 0.021211756393313408 vs the PyTorch vision tower on the canary.

What still lives in host code

The included Swift package is a fixed-shape native runtime for the current OCR path. The host app still owns:

  • providing a CGImage or already-preprocessed pixel_values
  • stopping criteria and max-token policy
  • UI and post-processing of generated OCR HTML
  • expanding beyond the current fixed 512x512 image / 300-token prefill shape

The included scripts/export_surya_coreml_runtime.py shows the Python reference glue used for validation. The included native/SuryaCoreMLRuntime package is the native implementation.

Example: run the current validation harness

pip install coremltools torch transformers pillow qwen-vl-utils huggingface_hub
python scripts/export_surya_coreml_runtime.py vision-combined-runtime-smoke \
  --model-id datalab-to/surya-ocr-2 \
  --vision-package surya_vision_int8.mlpackage \
  --prefill-package surya_prefill_fp16_seq300_cache512.mlpackage \
  --decode-package surya_decode_step_fp16_cache512.mlpackage \
  --output validation/local_vision_int8_prefill_decode_smoke.json \
  --max-cache-length 512 \
  --steps 8

Expected canary output today:

{
  "all_tokens_match": true,
  "coreml_text": "<p>Invoice ",
  "native_text": "<p>Invoice "
}

Example: run the native Swift smoke

cd native/SuryaCoreMLRuntime
swift run surya-coreml-smoke \
  --model-dir ../.. \
  --image /path/to/512x512-document.png \
  --max-tokens 8 \
  --vision int8

Validated canary output:

1039 2009 2046 2054 2047 2041 2035 2037 1979
<p>Invoice 

Important limitations

  • Fixed-shape canary export: the current packages are specialized to the traced 512x512 sample preprocessing path, with pixel_values [1024,1536], prefill length 300, and full-attention cache length 512.
  • The language prefill/decode packages are FP16 CoreML, not INT8/4-bit yet.
  • Only the vision tower has an INT8 package in this release.
  • This has not yet passed full allenai/olmOCR-bench.
  • The Swift decoder is intentionally small and should be hardened against the full tokenizer before production use.

Provenance

Generated non-destructively from datalab-to/surya-ocr-2. No fine-tuning was performed.