Luigi/moss-transcribe-diarize-litert

MOSS-Transcribe-Diarize — LiteRT (TFLite)

LiteRT port of OpenMOSS-Team/MOSS-Transcribe-Diarize (0.9B, Apache-2.0): joint zh/en transcription + speaker diarization + utterance timestamps ([ss.ss][Sxx]text[ss.ss]…).

Converted from the official bf16 safetensors checkpoint (f32-promoted) with litert-torch 0.9.1 (the renamed ai-edge-torch generative API). Conversion scripts, host runner and parity harness: https://github.com/vieenrose/LiteRT/tree/moss-td-port/litert/samples/asr/moss_td

Files

The network is split into three flatbuffers (each < 2 GB at f32; the split mirrors the reference C++ implementation, which also keeps embedding lookup and lm_head outside the decoder graph):

file	role	signatures
moss_td_encoder_*	Whisper-medium encoder + 4x time merge + VQAdaptor	mel (1,80,3000) → audio embeds (1,375,1024)
moss_td_embedder_*	tied token embedding / lm_head	embed_1, embed_128, logits
moss_td_decoder_*_ekv{N}	Qwen3-0.6B blocks + final norm	prefill_128, prefill_1024, decode; inputs_embeds + input_pos + mask + external KV cache (BTNH, 28×2 × (1,N,8,128))

Production set (what the live Space runs): moss_td_encoder_q8.tflite + moss_td_embedder_q8.tflite + moss_td_decoder_v2_q4b32_ekv2560.tflite + tokenizer/.

Also kept: the f32 parity reference (moss_td_encoder_f32, moss_td_embedder_f32, moss_td_decoder_f32_ekv6144) and a higher-fidelity moss_td_decoder_q8_ekv2048 (dynamic-range int8, sized for 90 s windows — no silence QAT, see below). Other variants (base int4, other ekv sizes) were removed 2026-07-23; they remain in this repo's git history and can be regenerated with the port's export.py.

fp16 set — for Android GPU-delegate inference (moss_td_encoder_fp16 + moss_td_embedder_fp16 + moss_td_decoder_fp16_ekv2048, ~1.8 GB total, restored 2026-07-24 byte-identical from the pre-cleanup revision): on the GPU delegate fp16 executes natively, making this the right precision for phone-GPU inference. Do not use it on the CPU/XNNPACK path — there fp16 is upconverted to f32 at load and OOMs 8 GB devices; use q8/int4 on CPU. Note the fp16 decoder is the pre-QAT base model (no silence-robust training) at ekv2048 (90 s windows).

v2 decoder (silence-robust QAT) — use this one

moss_td_decoder_v2_q4b32_ekv2560.tflite is exported from the project's q4mix-v2 QAT checkpoint (moss_q4mix_v4_qat/step_50) instead of the base weights. It fixes the base int4 decoder's silent-audio failure: on a silent tail the base int4-b32 decoder emitted a looping English hallucination ("the audio contains a sound effect…", 1185 chars) while v2 emits nothing (immediate EOS — the correct no-speech result). Measured deltas vs the now-removed base int4-b32 ekv2560 decoder on the same q8 encoder/embedder (host, compiled engine, 2026-07-23):

gate	base int4-b32	v2 int4-b32
jfk text	100% (golden)	100% — identical text, 2 end-timestamp digits shift
silent tail	1185-char hallucination loop	empty (immediate EOS)
zh90s text-only vs pinned f32 ref	98.728%	98.094% (−0.63 pt, within tolerance)

Encoder/adaptor and embedder are identical to base in this checkpoint, so moss_td_encoder_q8.tflite / moss_td_embedder_q8.tflite are used unchanged.

IMPORTANT: set the CPU thread count explicitly

CompiledModel.from_file(path) without options lets the runtime pick the thread count — and it picks very low (measured 8.5× slower in one production host case). Always pass:

from ai_edge_litert.compiled_model import CompiledModel
from ai_edge_litert.cpu_options import CpuOptions
from ai_edge_litert.options import Options
from ai_edge_litert.hardware_accelerator import HardwareAccelerator

opts = Options(hardware_accelerators=HardwareAccelerator.CPU,  # single enum, despite the plural name
               cpu_options=CpuOptions(num_threads=N))
cm = CompiledModel.from_file(path, options=opts)  # `options` XOR `hardware_accel` kwarg — not both

On x86 use the usable core count; on ARM big.LITTLE use the big-core count only (Exynos 1280: 398 ms/token at 4 big-core threads vs 849 ms at 8 — little-core contention doubles the cost). The classic Interpreter(num_threads=…) path never had this problem.

Parity (greedy decode, full pipeline)

Against the project's pinned f32 reference (official PyTorch pipeline, byte-matched independently by the C++ GGUF implementation on CPU and CUDA), the LiteRT f32 pipeline is byte-identical on all three validation clips (jfk 11 s, en 300 s, zh 318 s; difflib SequenceMatcher, autojunk=False, full text incl. timestamps and speaker tags).

Historical note: an earlier revision of this port interleaved time markers every 2 s instead of the processor's actual 5 s (time_marker_every_seconds = 5 in processor_config.json), which produced systematically finer segmentation. Audio, mel features and weights were verified bit-identical; the corrected prompt restores byte-exact parity.

Quantized (vs pinned f32 reference): q8 zh90s 99.80% full / 100% text-only; q8 jfk 98.76% / 100%; int4-b32 zh90s 95.45% / 98.99%; int4 jfk 98.76% / 100%. For scale, the C++ q4mix build scores 93.67% / 96.68% (zh90s) and 96.89% / 100% (jfk) on the same metric.

Runners

• runner.py --engine interpreter: reference implementation.
• runner.py --engine compiled and engine_cpp/moss_td_engine.cc: the KV cache lives in LiteRT TensorBuffers aliased as input AND output of every prefill/decode call (no host KV traffic). Host x86 zh90s q8: 72 s / 7.9 GB (interpreter) -> 43.9 s / 4.0 GB (Python compiled) -> C++ engine 2.08 GB; int4-b32 C++ engine 1.51 GB peak — equal peak RSS to the C++ GGUF q4mix build (1.50 GB) with higher transcript fidelity.
• Raspberry Pi 4: the prebuilt LiteRT-Next runtime requires ARMv8 AES (SIGILL on Cortex-A72); only the classic Interpreter path runs there.
• fp16 weights OOM 8 GB Android devices on the CPU/XNNPACK path (f32 upconversion at load); use q8/int4 there. On the GPU delegate fp16 runs natively — see the fp16 set above.

Provenance

Weights: OpenMOSS-Team/MOSS-Transcribe-Diarize @ e6d68cd (Apache-2.0). All non-v2 files: no fine-tuning; numerics-preserving conversion only. The v2 decoder is the one trained artifact: silence-robust QAT (q4mix-v2, checkpoint step_50), decoder weights only.