Loading…
SOTA OCR with Core ML and dots.ocr
Hugging FaceChristopher Fleetwood, Pedro Cuenca
Summary
Running competitive models like RedNote's 3B parameter dots.ocr on-device provides zero-cost, network-free optical character recognition while leveraging power-efficient hardware like Apple's Neural Engine. Because the Neural Engine requires Core ML, converting the 1.2B parameter NaViT vision encoder from PyTorch requires resolving multiple graph tracing incompatibilities. The initial conversion pipeline targets FLOAT32 execution on GPU by capturing the execution graph with torch.jit.trace and compiling it via coremltools. Simplifying the model for single-image inference eliminates dynamic sequence masking, dynamic tensor iteration, and multi-attention complexity. Although the resulting Core ML model accurately matches PyTorch baseline numerical precision, the unoptimized artifact spans over 5GB and takes over one second per forward pass.
Context
Executing OCR models on-device avoids API costs and network requirements while taking advantage of power-efficient hardware like Apple's Neural Engine. However, the Neural Engine is only accessible via Core ML, which introduces graph conversion hurdles when translating PyTorch implementations with unsupported operators, dynamic shapes, and strict data type constraints.
Approach / What changed
The authors trace the 1.2B NaViT vision encoder using torch.jit.trace and coremltools, targeting FLOAT32 on GPU with static shapes. They simplify the graph for single-image inference by standardizing on scaled_dot_product_attention, casting torch.arange outputs to match target dtypes, replacing dynamic sequence masking with constant float masks, and eliminating dynamic loop iteration over tensor shapes.
Takeaways
- coremltools ignores the dtype argument on torch.arange and outputs int32 by default, requiring an explicit type cast to prevent matmul dtype mismatch errors.
- Apple's Neural Engine does not support boolean tensors, necessitating the replacement of boolean attention masks with float masks containing zero values.
- An initial unoptimized Core ML FLOAT32 conversion of the 1.2B parameter vision encoder produced accurate outputs but exceeded 5GB in size and required over one second per forward pass.
Related reading
Introducing AnyLanguageModel: One API for Local and Remote LLMs on Apple Platforms
Apple developers frequently face integration friction when trying to support a mix of local and cloud language models across disparate APIs. To resolve this fragmentation, AnyLanguageModel introduces a unified Swift package that acts as a drop-in replacement for Apple's Foundation Models framework. The package standardizes interactions across local backends such as Core ML, MLX, and llama.cpp alongside cloud providers like Anthropic and OpenAI. To avoid dependency bloat from multi-backend support, the library employs Swift 6.1 package traits so projects only import the specific runtime engines they require. Furthermore, the library extends beyond current Foundation Models limitations by adding prompt image support for vision-language models.
MatttGrab ·
How we built a custom vision LLM to improve document processing at Grab