feat: update model

9dd7a6b 3 months ago

3.47 kB

language:
  - en
license: other
tags:
  - nemo
  - parakeet
  - tdt
  - automatic-speech-recognition
  - air-traffic-control
  - atc
  - singapore
  - military
base_model: nvidia/parakeet-tdt-0.6b-v2
pipeline_tag: automatic-speech-recognition
metrics:
  - wer
model-index:
  - name: parakeet-tdt-0.6b-v2-atc-singapore
    results:
      - task:
          type: automatic-speech-recognition
        metrics:
          - name: Validation WER
            type: wer
            value: 0.72

Parakeet-TDT 0.6B v2 - Singapore Military ATC

Fine-tuned NVIDIA Parakeet-TDT 0.6B v2 for Singapore Air Force air traffic control speech recognition.

Performance

Checkpoint	Validation WER	Notes
`model.ckpt`	0.72%	Best checkpoint from epoch 76
`epoch=100-val_wer=0.0073-last.ckpt`	0.73%	Final checkpoint, not published here

Model Details

Key	Value
Base model	`nvidia/parakeet-tdt-0.6b-v2`
Framework	NeMo
Model class	`EncDecRNNTBPEModel`
Format	Raw `.ckpt` checkpoint + tokenizer artifacts
Checkpoint size	7.0 GB
Domain	Singapore military ATC (Tengah WSAT, Paya Lebar WSAP)
Training data	Clean originals plus mild ATC radio, speed, and stress augmentations

Included Files

model.ckpt - best fine-tuned checkpoint
artifacts/705f11d22dc04b169effc35ce5cd1361_tokenizer.model
artifacts/a4715c7f6b2d4c2bb709306073d0c0a4_tokenizer.vocab
artifacts/4cf78c8ca4ca44fca36c3754478fb188_vocab.txt

Usage

This repo currently publishes the fine-tuned model as a raw NeMo checkpoint rather than a packaged .nemo archive. The tokenizer artifact paths therefore need to be pointed at the local artifacts/ folder before restore.

from pathlib import Path

import torch
from omegaconf import OmegaConf
from nemo.collections.asr.models import EncDecRNNTBPEModel

model_dir = Path("ASR/parakeet")
ckpt_path = model_dir / "model.ckpt"
artifacts_dir = model_dir / "artifacts"

bundle = torch.load(ckpt_path, map_location="cpu", weights_only=False)
cfg = bundle["hyper_parameters"]["cfg"]
cfg = OmegaConf.create(OmegaConf.to_container(cfg, resolve=False))
cfg.tokenizer.model_path = str(
    artifacts_dir / "705f11d22dc04b169effc35ce5cd1361_tokenizer.model"
)
cfg.tokenizer.vocab_path = str(
    artifacts_dir / "4cf78c8ca4ca44fca36c3754478fb188_vocab.txt"
)
cfg.tokenizer.spe_tokenizer_vocab = str(
    artifacts_dir / "a4715c7f6b2d4c2bb709306073d0c0a4_tokenizer.vocab"
)

model = EncDecRNNTBPEModel.load_from_checkpoint(str(ckpt_path), cfg=cfg)
model.eval().cuda()

hypotheses = model.transcribe(
    ["audio.wav"],
    return_hypotheses=True,
    timestamps=True,
)

if isinstance(hypotheses, tuple):
    hypotheses = hypotheses[0]

hyp = hypotheses[0]
print(hyp.text)
print(hyp.timestamp["word"])

Tested in ASTRA with nemo_toolkit[asr]==2.7.3.

Output Format

The model outputs normalized spoken text intended for downstream ATC formatting:

Input audio says	Model outputs
"CAMEL climb flight level zero nine zero"	`camel climb flight level zero nine zero`
"Contact Tengah Approach one three zero decimal zero"	`contact tengah approach one three zero decimal zero`
"Squawk four five two one"	`squawk four five two one`

ASTRA then applies a deterministic formatter to convert normalized speech into display text such as CAMEL climb FL090.