Add OmnilingualASR — 1,672 languages on CoreML + MLX, with shared SentencePiece/SDPA extract

[ivan-digital]

Summary

Single PR consolidating Meta's Omnilingual ASR port for Apple Silicon: CoreML 300M, MLX 300M/1B/3B/7B, plus two orthogonal refactors that extract duplicated building blocks into shared modules. Reference: arXiv 2511.09690, Apache 2.0.

Closes #195. Supersedes #199 and #200 (both will be closed with a pointer here).

What ships

1. OmnilingualASR CoreML (300M) — OmnilingualASRModel

• Sources/OmnilingualASR/: Configuration, SentencePieceVocabulary, CTCGreedyDecoder, CoreML model loader with 5 s and 10 s fixed-window variants (aufklarer/Omnilingual-ASR-CTC-300M-CoreML-INT8{,-10s})
• Utterance-level layer_norm preprocessing matching fairseq2's apply_audio_normalization — normalises the raw waveform before zero-padding so sub-window inputs match the reference pipeline statistics
• 40 s hard cap matching Meta's MAX_ALLOWED_AUDIO_SEC; clear error pointing to SpeechVAD for longer audio

2. OmnilingualASR MLX (300M / 1B / 3B / 7B) — OmnilingualASRMLXModel

• Sources/OmnilingualASR/MLX/: fairseq2-compatible wav2vec2 port from scratch
  • Wav2Vec2Frontend — 7-layer CNN feature extractor (kernels [10,3,3,3,3,2,2], strides [5,2,2,2,2,2,2] → 50 Hz frames), per-channel LayerNorm + GELU, post_extract_layer_norm, model_dim_proj, weight-normed grouped Conv1d positional encoder (kernel 128, groups 16) with even-kernel trim and residual
  • Wav2Vec2EncoderLayer — pre-norm self-attention + 2-linear FFN, quantised projections (QuantizedLinear group size 64, 4 or 8 bits)
  • Wav2Vec2Encoder — stack of N layers + final encoder.layer_norm
  • CTCHead — quantised final_proj → 10288
  • OmnilingualMLXWeightLoader — loads model.safetensors, fuses PyTorch weight_norm(dim=2) for the position encoder (W = g · v / ‖v‖_per_kernel), applies all tensors via MLXCommon helpers
  • OmnilingualMLXConfig — variant table (300M 24L/1024d/16h, 1B 48L/1280d/20h, 3B 60L/2048d/32h, 7B 128L/2048d/32h), auto-detect from HF model id

3. Shared extractions

AudioCommon/SentencePieceModel.swift — protobuf reader for .model files, exposing pieces: [Piece] with (text, score, type) and a PieceType enum. Eliminates two independent near-identical protobuf readers:

• OmnilingualASR/SentencePieceVocabulary.swift now wraps it (175 → 62 lines)
• PersonaPlex/SentencePieceDecoder.swift now wraps it (179 → 85 lines) and keeps the greedy unigram encoder

MLXCommon/SDPA.swift — multiHead / attendAndMerge / mergeHeads helpers. All reshape calls use -1 for the batch dimension so the helpers compose with MLX.compile(shapeless:) graphs that vary batch at runtime.

12 attention classes migrated to use SDPA.* (Qwen3ASR audio + float text + quantised text, Qwen3TTS Talker + CodePredictor + DecoderTransformer, Qwen3Chat GatedAttentionLayer, CosyVoice LLM, PersonaPlex Mimi + Depformer + Temporal, OmnilingualASR Wav2Vec2). Eliminates ~80 lines of duplicated reshape+SDPA+reshape boilerplate.

Caught and fixed during refactor

The first mergeHeads implementation used attn.dim(0) to size the batch dim, which baked it as a Swift constant and broke Qwen3-TTS compile(shapeless:) graphs that vary batch at runtime. Qwen3-TTS Talker testBatchWithDefaultInstruct fails with Cannot reshape array of size 4096 into shape (1,1,2048). Fix: use -1 for batch throughout. Verified end-to-end against Talker batch + CustomVoice + CosyVoice + Omnilingual MLX after fix.

CLI

# CoreML (default, ANE)
audio transcribe recording.wav --engine omnilingual               # 10 s window
audio transcribe recording.wav --engine omnilingual --window 5      # 5 s window

# MLX (GPU/Metal, any input length up to 40 s)
audio transcribe recording.wav --engine omnilingual --backend mlx                              # 300M @ 4-bit
audio transcribe recording.wav --engine omnilingual --backend mlx --variant 1B                  # 1B @ 4-bit
audio transcribe recording.wav --engine omnilingual --backend mlx --variant 3B --bits 8         # 3B @ 8-bit
audio transcribe recording.wav --engine omnilingual --backend mlx --variant 7B                  # 7B @ 4-bit

Docs

• docs/models/omnilingual-asr.md — architecture, all 10 published variants (CoreML 5/10 s + MLX 300M/1B/3B/7B × 4/8 bit), and the full 1,672-language list with English names + country hints (resolved via CLDR / pycountry, grouped by 32 scripts, in a collapsible <details> block)
• docs/inference/omnilingual-asr-inference.md — pipeline, layer-norm rationale, 40 s cap rationale, CoreML + MLX quick-starts, CLI examples, multilingual notes, perf numbers
• docs/shared-protocols.md — added OmnilingualASR (both backends) to the SpeechRecognitionModel conformers + not-thread-safe list; added the MLXCommon module section (SDPA, QuantizedMLP, PreQuantizedEmbedding); added OmnilingualASR module tree; corrected stale entries about QuantizedMLP/PreQuantizedEmbedding location (MLXCommon, not AudioCommon)
• CLAUDE.md / AGENTS.md — added OmnilingualASR + ParakeetStreamingASR to project structure; noted SDPA + SentencePieceModel in their respective common modules
• README.md + all 9 translations (zh/ja/ko/es/de/fr/hi/pt/ru) — feature bullet, SPM products, 4 model rows (300M/1B/3B/7B in 4 and 8 bit, CoreML 5 s and 10 s), memory table, and a Supported Languages row linking the canonical lang_ids.py

Verified output

CoreML 300M (reference-correct after layer-norm fix):

Lang	Transcript
EN	can you guarantee that the replacement part will be shipped tomorrow
AR	كما أثنى الزملاء المصارعون على لونا (reference match)
HI	लूना को साथी पहलवानों ने भी सरधांजलीदी
FR	pense à létineraire desqui comme un étineraire de rent donner similaire

MLX 300M-4bit (matches CoreML modulo one-char differences):

Lang	Transcript
EN	can you guarantee that the replacement part will be shiped tomorrow
AR	كما أثنى الزملاء المصارع على لونا
HI	लूना को साथी पहलवानों ने भी सरधांजली दी
FR	pensez à létineraire deski comme unétineraire de rendonner similaire

MLX 1B-4bit sanity check: can you guarantee that the replacement par o be shipe tomorrow (slightly more quant degradation, expected).

Tests

New unit tests (~26 added):

• OmnilingualASRTests/OmnilingualASRTests — 17 unit (config, layer-norm, CTC decoder, vocabulary)
• OmnilingualASRTests/OmnilingualMLXTests — 9 unit (variant table, frontend output length math)
• AudioCommonTests/SentencePieceModelTests — 7 unit (protobuf decode, type enum, helpers, subscript bounds, empty/unknown fields)
• AudioCLITests/TranscribeCommandTests — 11 new CLI parser tests for --engine omnilingual, --window, --backend, --variant, --bits

New E2E tests (E2E… prefix, skipped in CI):

• E2EOmnilingualASRTests — 9 (English real audio, warmup, unload, chunking, 40 s cap, FLEURS en/hi/fr/ar)
• E2EOmnilingualMLXTests — 7 (load, warmup, English, 40 s cap, FLEURS en/hi/fr/ar) — covers Wav2Vec2 frontend + encoder + CTC head + weight loader on real 300M-4bit weights
• E2EOmnilingualMLX1BTests — 1 (loads real 549 MB 1B-4bit weights, exercises the larger encoder code path)

Test resources: 4 small FLEURS fixtures (fleurs_{en,hi,fr,ar}.wav ~230-300 KB each).

Test plan

• swift build — green (all targets, all dependencies)
• swift test --skip E2E — 770/770 unit tests pass (+17 vs baseline for the new SentencePiece + Wav2Vec2 shape tests)
• swift test --filter OmnilingualASRTests — 26/26 pass (17 unit + 9 E2E on real 300M CoreML weights, ~5 s wall clock)
• swift test --filter E2EOmnilingualMLXTests — 7/7 pass (on real 300M-4bit MLX weights, ~5 s)
• swift test --filter E2EOmnilingualMLX1BTests — 1/1 pass (on real 1B-4bit MLX weights, ~56 s including 549 MB download)
• swift test --filter TranscribeCommandTests — 22/22 pass (11 new parser tests)
• Full swift test sweep — 949 executed, 917 passed, 31 conditionally skipped (XCTSkip for missing models / macOS-only), 1 pre-existing unrelated failure (E2EQwen35CoreMLConversionTests.testPyTorchDecoderProducesHello hard-codes scripts/convert_qwen35_chat_coreml.py which was never committed to the repo; failing long before this branch, tracked separately)
• Multilingual transcripts verified on real weights for EN / HI / FR / AR on both CoreML and MLX backends
• Qwen3-TTS Talker batch test passes after the -1 batch-dim fix in SDPA.mergeHeads
• Docs sync'd: English README + 9 translations, docs/models/, docs/inference/, docs/shared-protocols.md, CLAUDE.md / AGENTS.md

Out of scope (follow-up tickets)

• 3B and 7B MLX variants spot-checked via the same code path as 300M and 1B, but no dedicated E2E tests (weights are 1.7 GB and 3.55 GB — manual verification only until CI has the disk for it)
• Streaming / chunked dictation flow on top of the 5 s CoreML window (reference pipeline does not support streaming)
• The pre-existing testPyTorchDecoderProducesHello failure — tracked separately; unrelated to this PR