multilingual-embedders-onnx-coreml

Multilingual Embedding Models — ONNX + CoreML

Three multilingual text embedding models, each exported to ONNX and converted to CoreML (.mlpackage, fp16 + int8-weight) for iOS/macOS.

Full writeup: see COMPARISON.md (or the published artifact, if you asked for one).

ONNX precision matches each checkpoint's actual training precision

Checked directly against each model's safetensors header rather than assumed: multilingual-e5-small is genuinely trained/stored in fp32; f2llm-v2-80m and granite-embedding-97m-multilingual-r2 are both stored in bf16. Shipping an "fp32" file for the latter two would just be their bf16 numbers sitting in a container twice the needed size — not extra precision. So:

The bf16 files are a genuine, direct bf16 export (no fp32→fp16 graph surgery) and are validated at cosine ≥ 0.9998 against the original PyTorch checkpoint. Getting ONNX Runtime's CPU provider to actually run them required patching ~15 missing bf16 CPU kernels (Mul, MatMul, Add, Where, Sigmoid, Softmax, Clip, LayerNormalization, and more) — see scripts/export_onnx.py's _autopatch_bf16_kernel_gaps if reconverting after an ONNX Runtime upgrade.

The bf16 files use a different output tensor name than the other variants — sentence_embedding instead of embedding — because the literal name "embedding" collides with the auto-generated name of the internal aten::embedding (token lookup) op under PyTorch's dynamo ONNX exporter, producing an invalid graph. Only model_pooled_bf16.onnx files are affected.

Directory layout

models/<model-name>/
  onnx/
    model_pooled_fp32.onnx   # multilingual-e5-small only (its native precision)
    model_pooled_bf16.onnx   # f2llm-v2-80m + granite only (their native precision)
                              # output tensor name is "sentence_embedding", not "embedding"
    model_pooled_fp16.onnx   # e5-small + granite (post-hoc converted; not shipped for f2llm)
    model_pooled_int8.onnx   # dynamic weight quantization, all 3 models
    refs.npz                 # reference input_ids/attention_mask/embeddings for parity checks
  coreml/
    <name>_seq{64,128,256}_fp16.mlpackage
    <name>_seq{64,128,256}_w8.mlpackage
    refs.npz                 # reference inputs (padded to 128) + PyTorch embeddings
  tokenizer/                 # HF tokenizer files (tokenizer.json, vocab, etc.)

What's baked into every exported graph

Each model is wrapped so the exported graph takes (input_ids, attention_mask) and returns a single L2-normalized embedding vector — pooling and normalization are inside the graph, not left to the caller:

• multilingual-e5-small: mean pooling over the attention mask. Prepend "query: " or "passage: " to input text before tokenizing (not baked in).
• f2llm-v2-80m: last-token pooling (hidden state at the final non-pad position). The tokenizer auto-appends EOS. For queries, prepend "Instruct: <task description>\nQuery: "; documents get no prefix.
• granite-embedding-97m-multilingual-r2: CLS token pooling. No prefix needed.

ONNX usage

import onnxruntime as ort
from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("models/multilingual-e5-small/tokenizer")
sess = ort.InferenceSession("models/multilingual-e5-small/onnx/model_pooled_int8.onnx",
                             providers=["CPUExecutionProvider"])
enc = tok(["query: what is the capital of france?"], padding=True, return_tensors="np")
emb = sess.run(["embedding"], {"input_ids": enc["input_ids"], "attention_mask": enc["attention_mask"]})[0]

For the bf16 files (f2llm-v2-80m, granite-embedding-97m-multilingual-r2), use "sentence_embedding" instead of "embedding" as the output name:

sess = ort.InferenceSession("models/f2llm-v2-80m/onnx/model_pooled_bf16.onnx",
                             providers=["CPUExecutionProvider"])
emb = sess.run(["sentence_embedding"], {"input_ids": ..., "attention_mask": ...})[0]

CoreML usage (Swift, iOS/macOS)

The graphs use fixed sequence lengths (64/128/256 tokens) rather than one flexible-shape model. Pick the shortest length that fits your input, pad/truncate to it, and load the matching .mlpackage.

let model = try f2llm_v2_80m_seq128_fp16(configuration: .init())
let input = f2llm_v2_80m_seq128_fp16Input(input_ids: idsArray, attention_mask: maskArray)
let embedding = try model.prediction(input: input).embedding

Verify on a Mac before shipping — this pipeline was built and validated on Windows/WSL, which cannot run CoreML inference. Run:

pip install coremltools numpy
python scripts/verify_coreml_on_mac.py models/<name>/coreml/<name>_seq128_fp16.mlpackage models/<name>/coreml/refs.npz

Expect cosine similarity ≥ 0.999 (fp16) / ≥ 0.99 (int8-weight) against the bundled PyTorch reference embeddings.

Precision note for f2llm-v2-80m and granite-embedding-97m-multilingual-r2: their CoreML packages use FLOAT16 compute precision, but both checkpoints are natively bf16 — bf16 has a wider exponent range than fp16, so converting to fp16 compute can in principle overflow for some values (a RuntimeWarning: overflow encountered in cast was observed during conversion). This has not been numerically verified — that requires running the actual .mlpackage on Apple hardware, which this pipeline can't do. If you hit unexpected accuracy issues with these two models specifically, this is the first thing to check; a FLOAT32-compute CoreML variant would sidestep it.

Rebuilding

Everything here was produced by scripts in scripts/, run inside WSL (scripts/wsl_env.sh sets up the environment):

bash scripts/run_wsl_pipeline.sh all      # ONNX export, all 3 models
wsl bash scripts/do_coreml_convert.sh <model-dir-name> <model-dir-name>   # per model
wsl bash -c "source scripts/wsl_env.sh; \$VENV/python scripts/save_tokenizers.py"
wsl bash -c "source scripts/wsl_env.sh; \$VENV/python scripts/benchmark_onnx.py"

---

MobileCLIP-S0 — ONNX + ORT for Android

Apple's MobileCLIP-S0 (CVPR 2024) — a dual-tower image+text CLIP model — exported to ONNX and converted to .ort format for Android. Full plan/history: MOBILECLIP_PLAN.md.

License: Apple's own apple-amlr (research license, not MIT/Apache) — keep any shared copy of these weights private/team-only, not public.

Precision: fp32 + fp16 only, both essentially exact

The checkpoint is genuinely fp16-native (every leaf tensor in the raw state dict is torch.float16). Both fp32 and fp16 ONNX/ORT variants are validated at cosine ≥ 0.999999 against the original PyTorch model — for practical purposes, exact. fp16 was produced via post-hoc conversion of the fp32 graph (onnxruntime.transformers.onnx_model.OnnxModel.convert_float_to_float16, keep_io_types=True) rather than a "native" fp16 export — this repo's proven, lower-risk path (already used successfully for two of the three text models above), and it automatically keeps numerically-sensitive ops (LayerNorm, Softmax) in fp32 via Cast insertion — which happens to match this checkpoint's own LayerNormFP32 modules' design intent exactly.

No int8 file is shipped. Dynamic quantization degraded the image encoder to cosine ≈ 0.03 (uncorrelated) and the text encoder to ≈ 0.74 — both well below this project's bar, and both discarded rather than shipped. Dynamic quantization suits MatMul/Gemm-dominant graphs; the image encoder in particular is Conv-dominant, which generally needs calibration-based static quantization to retain accuracy, not the dynamic kind used here.

Directory layout

models/mobileclip-s0/
  pytorch/
    mobileclip_s0_fp16.pt          # original checkpoint
  onnx/
    image_encoder_fp32.onnx        # + .onnx.data (external weights)
    image_encoder_fp16.onnx
    text_encoder_fp32.onnx
    text_encoder_fp16.onnx
    refs.npz                      # fixture inputs + PyTorch reference embeddings (both towers)
  ort/
    Fixed/     image_encoder_{fp32,fp16}.ort, text_encoder_{fp32,fp16}.ort   # RECOMMENDED, see below
    Runtime/   same 4, suffixed .with_runtime_opt.ort   # for NNAPI/CoreML-style compiling EPs
  tokenizer/
    bpe_simple_vocab_16e6.txt.gz   # CLIP BPE vocab (same file the model was tokenized with)
    preprocessing_spec.md         # exact image preprocessing + tokenization spec — READ THIS FIRST
  fixtures/
    test_images/*.png             # synthetic parity-test images (solid colors/gradient/checkerboard/noise)

Read tokenizer/preprocessing_spec.md before integrating — it has the exact (empirically-confirmed, not assumed) image preprocessing and tokenization recipe. Getting either wrong silently produces plausible-looking but wrong embeddings, no error thrown.

Use the CPU execution provider, not NNAPI

onnxruntime.tools.check_onnx_model_mobile_usability was run against both towers. Conclusion for both: NNAPI and CoreML EPs cover only a small fraction of nodes and are predicted to perform worse than plain CPU execution due to partition fragmentation, even after reshaping to fixed input shapes. Use the CPU EP — that's what the ort/Fixed/ files are optimized for. ort/Runtime/ variants are shipped too (in case your specific device differs) but target compiling EPs like NNAPI, which this model doesn't benefit from here.

ONNX usage (Python, for validation)

import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession("models/mobileclip-s0/onnx/image_encoder_fp16.onnx",
                             providers=["CPUExecutionProvider"])
pixel_values = ...  # (1,3,256,256) float32, see preprocessing_spec.md
emb = sess.run(["image_embedding"], {"pixel_values": pixel_values})[0]

Android usage (.ort)

val env = OrtEnvironment.getEnvironment()
val session = env.createSession("image_encoder_fp16.ort", OrtSession.SessionOptions())
val input = OnnxTensor.createTensor(env, pixelValuesFloatBuffer, longArrayOf(1, 3, 256, 256))
val output = session.run(mapOf("pixel_values" to input))
val embedding = (output[0].value as Array<FloatArray>)[0]  // 512-d, L2-normalized

For the text tower, tokenize per tokenizer/preprocessing_spec.md (either port the CLIP BPE algorithm using the bundled vocab file, or bake tokenization into the ONNX graph itself via onnxruntime-extensions' CLIP tokenizer custom op — recommended if you want to avoid a hand-ported BPE implementation entirely).