litert-community/japanese-zipformer-base-LiteRT

japanese-zipformer-base — LiteRT (GPU)

Japanese speech recognition with reazon-research/japanese-zipformer-base-k2-rs35kh-bpe (96.5 M params, trained on ReazonSpeech, avg CER 11.46 %) running fully on the LiteRT CompiledModel GPU (ML Drift) — with zero FFT anywhere: the raw 16 kHz waveform goes straight into a wav2vec2-style 1D-conv frontend (no mel/fbank even on the host), then a Zipformer encoder (6 multi-rate stacks) and a CTC head, all in one GPU graph.

Real on-device output: CTC token onsets for a reading of the Preamble of the Constitution of Japan (Wikimedia Commons, CC0). A 16 s window transcribes in 621 ms on a Pixel 8a GPU (RTF ≈ 0.04). Greedy CTC without an LM transcribes phonetically exactly; occasional kanji homophone swaps (選挙→占拠) are the expected no-LM behavior.

File	Size	Input	Output	API
ja_zipformer_ctc_fp16.tflite	197 MB	waveform [1,256000] + 4 mask biases	CTC logits [1,799,3004]	CompiledModel GPU

Pipeline

`16 kHz mono PCM in [-1,1] → [GPU] conv frontend (stride 320 → 50 Hz) + Zipformer2 (6 stacks)

• CTC linear → host greedy-CTC + BPE detokenize`

• Fixed 16 s window: pad 0.5 s of zeros on both sides of the audio (the upstream model-card convention), then zero-pad to 256000 samples.
• Mask inputs: additive attention biases (0 real / -1000 pad) at the internal frame rates: [1,799], [1,400], [1,200], [1,100]. Build the 50 Hz bias with valid = conv_frames(n_samples + 16000) and take [::2], [::4], [::8] slices (conv_frames: fold L=(L-k)//s+1 over (10,5)(3,2)(3,2)(3,2)(3,2)(2,2)(2,2)).
• Output: raw CTC logits at 50 Hz. Blank id = 0 (labeled <unk> in vocab.json but it behaves as the CTC blank — icefall convention), BPE vocab 3004 (vocab.json).

Minimal usage — Python

import json, numpy as np, torch, torchaudio
from ai_edge_litert.interpreter import Interpreter

wave, _ = torchaudio.load("speech_16k.wav")          # 16 kHz mono, [-1,1]
a = wave[0]; lead = 8000
x = torch.zeros(1, 256000)
n = min(a.shape[0], 256000 - 2 * lead)
x[0, lead:lead + n] = a[:n]

L = n + 2 * lead
for k, s in [(10,5),(3,2),(3,2),(3,2),(3,2),(2,2),(2,2)]:
    L = (L - k) // s + 1                             # valid 50 Hz frames
b = np.full((1, 799), -1000.0, np.float32); b[0, :L] = 0.0
biases = {799: b, 400: b[:, ::2], 200: b[:, ::4], 100: b[:, ::8]}

it = Interpreter(model_path="ja_zipformer_ctc_fp16.tflite"); it.allocate_tensors()
for d in it.get_input_details():
    s = list(d["shape"])
    it.set_tensor(d["index"], x.numpy() if s[1] == 256000
                  else np.ascontiguousarray(biases[s[1]]))
it.invoke()
logits = it.get_tensor(it.get_output_details()[0]["index"])[0]   # [799, 3004]

vocab = {i: t for t, i in json.load(open("vocab.json")).items()}
out, prev = [], -1
for i in logits[:L].argmax(-1):
    if i != prev and i != 0: out.append(vocab[int(i)])
    prev = i
print("".join(out).replace("▁", " ").strip())

Minimal usage — Kotlin (Android)

val model = CompiledModel.create(modelPath, CompiledModel.Options(Accelerator.GPU), null)
val inputs = model.createInputBuffers()
val outputs = model.createOutputBuffers()

// resolve slots by capacity: waveform 256000, biases 799/400/200/100 floats
val waveSlot = inputs.indexOfFirst { it.readFloat().size == 256000 }
inputs[waveSlot].writeFloat(pcm)                   // [-1,1], 0.5 s lead pad, zero-padded
for (len in intArrayOf(799, 400, 200, 100)) {      // additive masks: 0 real / -1000 pad
    val slot = inputs.indexOfFirst { it.readFloat().size == len }
    val ds = (799 + len - 1) / len
    inputs[slot].writeFloat(FloatArray(len) { i -> if (i * ds < valid50) 0f else -1000f })
}

model.run(inputs, outputs)
val logits = outputs[0].readFloat()                 // [799 * 3004], readback syncs the GPU
// greedy CTC: per-frame argmax, drop blanks (id 0) + repeats, BPE detok ('▁' -> space)

On-device performance (Pixel 8a, CompiledModel GPU)

• GPU compile 2.4 s (first load); 16 s window: 621 ms run+readback (24 ms enqueue), RTF ≈ 0.04.
• Device vs desktop float reference: per-frame argmax agreement 98.5–100 %, transcripts identical on the test sweep.

Conversion notes

Converted from the upstream PyTorch (safetensors) with litert-torch; all rewrites numerically exact (tflite vs PyTorch: corr 1.000000). Same re-authoring set as the LibriSpeech Zipformer CR-CTC conversion (see litert-community/Zipformer-medium-CR-CTC-LiteRT): guard-free Swoosh softplus, pad+reshape+slice rel-shift, per-rate additive-bias masks, concat-repeat up/downsample with baked weight softmax, plus the wav2vec2-frontend recipe (tanh-GELU, 4D group-norm).

Sources & license

• Model: reazon-research/japanese-zipformer-base-k2-rs35kh-bpe — Apache-2.0 (ReazonSpeech corpus).
• Architecture: Zipformer (k2-fsa/icefall).
• Hero audio: Preamble of the Constitution of Japan, Wikimedia Commons — CC0.