bangla-zipformer-pt

Bangla Zipformer - PyTorch

Bangla speech recognition: a 65.5M-parameter Zipformer transducer, fine-tuned from an English LibriSpeech model on about 1,210 hours of Bangla speech. CPU-friendly ONNX versions (fp32 / fp16 / int8): SayedShaun/bangla-zipformer-onnx.

Quick start

The model code comes from icefall, so you need to install it and a few packages once:

git clone https://github.com/k2-fsa/icefall
git -C icefall checkout 3f848bb6d0acc970c9b294a30ca0a04a7c9c78d1      # the version this model was trained with
export ICEFALL_ROOT=$PWD/icefall

pip install torch lhotse sentencepiece soundfile soxr numpy kaldialign pypinyin tensorboard dill huggingface_hub
pip install k2        # needs a wheel matching your torch / CUDA: https://k2-fsa.github.io/k2/installation/index.html

Then, using the repo name:

import sys
from huggingface_hub import snapshot_download
sys.path.insert(0, snapshot_download("SayedShaun/bangla-zipformer-pt"))
from asr import BanglaASR

asr = BanglaASR.from_pretrained("SayedShaun/bangla-zipformer-pt")           # uses the GPU if available
print(asr.transcribe("audio.wav"))                                  # beam search (default)
print(asr.transcribe("audio.wav", decoding="greedy"))               # faster

Accepts a file path (wav / flac / ogg ...) or (samples, sample_rate). Don't want to install icefall and k2? Use the ONNX version, which needs only onnxruntime.

Using the model in PyTorch

asr.py is a thin wrapper. To use the weights directly, set up icefall as in the quick start (ICEFALL_ROOT), then:

import os, sys, argparse
import torch, sentencepiece as spm, soundfile as sf

ICEFALL = os.environ["ICEFALL_ROOT"]
sys.path[:0] = [ICEFALL, f"{ICEFALL}/egs/librispeech/ASR/zipformer"]    # the `icefall` package + the stock Zipformer recipe

from train import add_model_arguments, get_model, get_params           # recipe modules (not pip packages)
from beam_search import greedy_search_batch, modified_beam_search
from lhotse import Fbank, FbankConfig

sp = spm.SentencePieceProcessor(); sp.load("bpe.model")

# 1. build the network: default (medium, non-streaming) architecture - do not change these arguments
parser = argparse.ArgumentParser(); add_model_arguments(parser)
params = get_params(); params.update(vars(parser.parse_args([])))
params.update(context_size=2, use_transducer=True, use_ctc=False, use_attention_decoder=False,
              blank_id=sp.piece_to_id("<blk>"), unk_id=sp.piece_to_id("<unk>"), vocab_size=sp.get_piece_size())
model = get_model(params)

# 2. load the weights
model.load_state_dict(torch.load("model.pt", map_location="cpu")["model"], strict=True)
model.eval()

# 3. features -> encoder -> search -> text
wave, sr = sf.read("audio.wav", dtype="float32")                        # 16 kHz mono, up to ~20 s
feats = torch.from_numpy(Fbank(FbankConfig(num_mel_bins=80)).extract(wave, sr))[None]     # (1, T, 80)
with torch.no_grad():
    enc, enc_lens = model.forward_encoder(feats, torch.tensor([feats.shape[1]]))          # (1, T', 512)
    tokens = modified_beam_search(model=model, encoder_out=enc, encoder_out_lens=enc_lens, beam=4)   # or greedy_search_batch(...)
print(sp.decode(tokens)[0])

Results

Full dev (5,282 utterances, 8.1 h) and test (10,485 utterances, 16.4 h) sets.

WER counts punctuation attached to words; CER ignores punctuation. With punctuation stripped, WER is 13.82 (beam) / 14.11 (greedy) on test and 13.69 / 13.94 on dev. The averaging window (epochs 23-25) and beam size were chosen on dev only.

Beam search or greedy?

Beam search (beam 4) is the default. Both decoders use the same model; they differ in accuracy and speed.

• Use beam search when accuracy matters and latency does not: offline or batch transcription, subtitles, labelling data. It is about 0.3 WER points (about 2% relative) better, consistently on dev and test.
• Use greedy for low latency or high throughput, especially on a GPU. On the GPU greedy is about 3.8x faster than beam search; on the CPU the gap is smaller (beam search is about 27% slower). The loss is small: about 0.3 WER points.
• Why the GPU does not help beam search much: it speeds up the encoder (about 14 ms vs 108 ms on CPU for a 5 s clip), but the beam search runs frame by frame and still takes roughly 90 ms for that clip, against about 15 ms for greedy on the GPU.
• Beam size: the default is 4 (from_pretrained(..., beam=4)). Beam 8 was no better on dev (14.21 vs 14.17 WER) and is slower.
• Need fast CPU inference? The ONNX version runs about 5x faster than this loader on CPU (int8, beam search) at nearly the same accuracy, and has both decoders too.

asr.transcribe("audio.wav")                      # beam search (default)
asr.transcribe("audio.wav", decoding="greedy")   # greedy

Timings: median of 5 repeats per clip on 60 dev clips of 1-15 s (20 each of short, medium and long), beam and greedy run back to back on the same clip, audio passed as an array; 32-core CPU with 4 threads and one GPU (RTX 5080, peak model memory about 315 MB). On this CPU the PyTorch loader was no faster with 4 threads than with 1 (I did not investigate why).

Training data

772,605 utterances, 1,209.6 hours of Bangla speech from nine corpora. Dev and test are held-out utterances from the same sources.

*PyTorch model, beam search. Kathbath, OpenSLR and IndicVoices are easiest (10-12%); Vaani (~25%) and Shrutilipi (~20%) are harder. BEN10 regional-dialect speech is handled poorly (~78%) - it is tiny (6 training hours, 16 test utterances), so treat that number as indicative only.

How it was trained

• Base: encoder initialised from the English LibriSpeech Zipformer; decoder, joiner and tokenizer trained from scratch.
• Code: icefall Zipformer recipe, commit 3f848bb, with k2, lhotse and PyTorch 2.11.
• Recipe: pruned RNN-T loss, ScaledAdam + Eden schedule, 25 epochs (~3 h each) on one GPU in bf16, about 400 s of audio per batch.
• Final weights: average of epochs 23-25. Decoding: greedy or beam search (beam 4).

Limitations

• Scores are on held-out data from the training sources; expect worse results on other domains, noisy phone audio and regional dialects.
• Up to ~20 s per call (the loader splits longer audio automatically); 16 kHz mono; not a streaming model.
• Output follows the training text: no {...} tags, and it may contain punctuation.

License

• Model weights, tokenizer and this card: CC BY-SA 4.0 (LICENSE). Attribution to the data sources is required.
• Loader code (asr.py): Apache-2.0 (LICENSE-CODE-Apache-2.0).

References

Model and code 1. Z. Yao et al. Zipformer: A faster and better encoder for automatic speech recognition. ICLR 2024. arXiv:2310.11230 2. F. Kuang et al. Pruned RNN-T for fast, memory-efficient ASR training. 2022. arXiv:2206.13236 3. icefall (commit 3f848bb6d0acc970c9b294a30ca0a04a7c9c78d1), k2 4. P. Zelasko et al. Lhotse: a speech data representation library for the modern deep learning ecosystem. NeurIPS 2021 DCAI workshop. arXiv:2110.12561 5. V. Panayotov et al. LibriSpeech: an ASR corpus based on public domain audio books. ICASSP 2015 (the base model's training data)

Training data 6. IndicVoices: T. Javed et al. 2024. arXiv:2403.01926 7. Kathbath: T. Javed et al. IndicSUPERB. 2022. arXiv:2208.11761 8. Shrutilipi: K. S. Bhogale et al. 2022. arXiv:2208.12666 9. Vaani: S. Pulikodan et al. VAANI: Capturing the language landscape for an inclusive digital India. 2026. arXiv:2603.28714 10. FLEURS: A. Conneau et al. 2022. arXiv:2205.12446 11. Common Voice: R. Ardila et al. LREC 2020. arXiv:1912.06670 12. OpenSLR-53: O. Kjartansson et al. SLTU 2018. doi:10.21437/SLTU.2018-11 13. OpenSLR-37: K. Sodimana et al. SLTU 2018 (the citation its page requests). doi:10.21437/SLTU.2018-14 14. BEN10: no citation information found.