Instructions to use rajb3/kea-t with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- NeMo
How to use rajb3/kea-t with NeMo:
import nemo.collections.asr as nemo_asr asr_model = nemo_asr.models.ASRModel.from_pretrained("rajb3/kea-t") transcriptions = asr_model.transcribe(["file.wav"]) - Notebooks
- Google Colab
- Kaggle
Kea-T: Parakeet-TDT-0.6B-v2 with ternary encoder weights {-1, 0, +1}
Kea-T is NVIDIA's English speech recognizer nvidia/parakeet-tdt-0.6b-v2 with every large matrix in its 24-layer FastConformer encoder (the attention and feed-forward projections and the pointwise convolutions, 98% of the model's parameters) constrained to three values (-1, 0, +1) with one FP32 scale per output row. The accuracy loss is largely recovered by quantization-aware training (QAT) with distillation from the original model. Its 1-bit sibling Kea-B is 105.3 MB with a 7.05% test mean.
At a glance
| Kea-T | Original | |
|---|---|---|
| Weight file | 180.8 MB | 2,472 MB |
| Test WER, mean of 7 sets | 6.84% | 6.45% |
| Latency, 10 s clip | 4.46 ms (23.2% lower) | 5.81 ms |
| GPU energy, 10 s clip | 1.57 J (17.1% less) | 1.90 J |
| Process VRAM, 10 s clip | 1.89 GB (13.5% lower) | 2.18 GB |
The original's size is NVIDIA's .nemo file, and its WER is the FP32 model as we scored it.
The speed rows are warm batch-one transcription on one RTX 5090 at a 400 W power limit, Kea-T
in expanded mode against our own optimized BF16 build of the original (not stock NeMo). That speed
needs our custom CUDA runtime
on GitHub, not load_ternary.py. The runtime's compact mode uses 1.27 GB of process VRAM, 41.8%
less than that build. All configurations, the 3 s and 30.04 s clips and the method are in
RUNTIME.md. Our research runtime with two-part INT8 activations is faster still: at
10 s, 4.30 ms (25.9% lower than the BF16 build) and about 1.33 J (about 30% less), measured in a
later session, so the energy comparison is approximate
(details).
Results
Test WER (%), greedy decoding:
| Model | LS clean | LS other | AMI | Earnings-22 | GigaSpeech | SPGISpeech | VoxPopuli | Mean of 7 | Common Voice |
|---|---|---|---|---|---|---|---|---|---|
| Original, FP32 (our scoring) | 1.70 | 3.19 | 11.15 | 11.24 | 9.78 | 2.14 | 5.94 | 6.45 | 8.50 |
| Kea-T (this model) | 2.05 | 4.20 | 10.42 | 11.72 | 10.35 | 2.94 | 6.19 | 6.84 | 12.56 |
We scored both models with the same code on every utterance of the public Open ASR Leaderboard
test bundle. TED-LIUM isn't in that bundle, so the mean covers the other seven sets, and Common
Voice is shown on its own. Decoding is greedy in strict FP32, and both the references and the
outputs go through the Whisper English text normalizer. Our scoring of the FP32 original
reproduces NVIDIA's published numbers within 0.09 points per set. Kea-T was scored once, on the
model rebuilt from this exact export.safetensors; on our development sets, both GPU runtimes give
the same mean WER as FP32 decoding (6.09%). NVIDIA's published row, the difference and ratio rows
and the scoring details are in EVALUATION.md.
How it was trained
- We started from NVIDIA's pretrained FP32 checkpoint. In each of the 264 ternary matrices, a
weight W becomes
clamp(round(W / s), -1, 1)times its row's scale s, the row's mean |W|. Over the first 25% of steps (62,500) the weights were blended from float to fully ternary; after that they stayed fully ternary. - The frozen original model was the teacher. Kea-T learned the original's transcripts of each training batch and matched its encoder outputs, and the prediction and joint networks kept NVIDIA's weights.
- Training ran 250,000 steps of batches of up to 600 s of audio, with a peak learning rate of 5e-4: 33,248 hours of audio from YODAS-Granary, LibriSpeech, People's Speech, VoxPopuli and AMI. It took 69.3 hours on one RTX 5090.
- The remaining learned weights are stored in FP16. We checked the export against the trained model: the codes, the scales and the greedy transcripts of four check clips were identical.
The full recipe, the data mix, the pilot and the file format are in TRAINING.md.
Use
Needs nemo_toolkit[asr] (with its torch), safetensors and soundfile.
pip install "nemo_toolkit[asr]"
hf download rajb3/kea-t --local-dir kea-t
python kea-t/load_ternary.py clip.wav # 16 kHz audio; prints the transcript
import sys; sys.path.insert(0, "kea-t")
from load_ternary import load, transcribe
model = load("kea-t", device="cuda") # or "cpu"
print(transcribe(model, ["clip.wav", "clip2.flac"])) # paths or 1-D 16 kHz float arrays
load() checks the weights and tokenizer files against the SHA-256 hashes in manifest.json,
unpacks the ternary codes and rebuilds a standard NeMo model in FP32; transcribe() decodes
exactly as the evaluation did (greedy TDT, FP32, dither 0). Output is punctuated and
capitalized, like the original. The rebuilt model needs the same memory and compute as the
original, so with load_ternary.py the small file is a storage and download saving only.
Each input is decoded in one pass with full attention, with no chunking. The model was trained on 1 to 30 s utterances; split long recordings into segments of about 30 s or less.
Files
export.safetensors(180,797,564 bytes): the ternary codes packed four per byte, one FP32 scale per output row, and every other learned tensor in FP16.manifest.json(hashes, configuration, accounting, provenance),reconstruction.json(the check of the export against the trained model) andtokenizer/(unchanged).load_ternary.py: the standalone loader and a small transcription CLI (MIT License).LICENSE(CC-BY-4.0 legal code) andNOTICE(attribution and changes).- RUNTIME.md, EVALUATION.md and TRAINING.md.
Limitations
- English only, like the base model.
- The gap to the original is largest on Common Voice (+4.06 points), which is outside the training mixture, followed by LibriSpeech test-other (+1.01) and SPGISpeech (+0.81). Expect larger degradation on domains far from the training data.
- The runtime results were measured on one RTX 5090 at a 400 W limit, warm, at batch one, with one clip per length.
- Only greedy decoding was evaluated. Tested with NeMo 3.0.0 and torch 2.11.0+cu128; loading on the CPU was tested, decoding on the CPU was not. More in EVALUATION.md.
Used in
- LocalFlow, an offline macOS dictation app (beta; Apple Silicon, macOS 26 or later) forked from FreeFlow. It runs a Core ML conversion of these weights (ternary lookup tables, FP16 values) on the Mac's Neural Engine, and audio never leaves the device.
License and attribution
The model weights, this model card and its detail documents (RUNTIME.md, EVALUATION.md,
TRAINING.md) are licensed under CC-BY-4.0
(see LICENSE), the license of the base model. The model is a modified version of
nvidia/parakeet-tdt-0.6b-v2 by NVIDIA:
the encoder's projection and pointwise-convolution weights were retrained with ternary
quantization-aware training and are stored as ternary codes with per-row FP32 scales; the
remaining learned weights are stored in FP16; the weights are distributed as safetensors with a
standalone loader instead of a .nemo file. NOTICE lists the changes and
the training data, and TRAINING.md their licenses.
load_ternary.py is MIT-licensed (notice at the top of the file), and so is the
training and evaluation code.
- Downloads last month
- 139
Model tree for rajb3/kea-t
Base model
nvidia/parakeet-tdt-0.6b-v2Datasets used to train rajb3/kea-t
espnet/yodas-granary
MLCommons/peoples_speech
Evaluation results
- Test WER on LibriSpeech test-clean (ESB test bundle)self-reported2.050
- Test WER on LibriSpeech test-other (ESB test bundle)self-reported4.200
- Test WER on AMI test (ESB test bundle)test set self-reported10.420
- Test WER on Earnings-22 test (ESB test bundle)test set self-reported11.720
- Test WER on GigaSpeech test (ESB test bundle)test set self-reported10.350
- Test WER on SPGISpeech test (ESB test bundle)test set self-reported2.940
- Test WER on VoxPopuli test (ESB test bundle)test set self-reported6.190
- Test WER on Common Voice test (ESB test bundle)test set self-reported12.560