| """Transcribe audio files with a pruned omniASR CTC model (greedy CTC, no language model). |
| |
| python inference.py audio1.wav [audio2.flac ...] [--model-dir .] [--device cuda] |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import soundfile as sf |
| import torch |
| import torchaudio.functional as AF |
| from safetensors.torch import load_file |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| from omniprune.model import ArchConfig, PrunableCTC, normalize_audio_ |
| from omniprune.text import CharVocab |
|
|
|
|
| def load(model_dir: str | Path = ".", device: str = "cpu") -> tuple[PrunableCTC, CharVocab]: |
| d = Path(model_dir) |
| model = PrunableCTC(ArchConfig.from_dict(json.loads((d / "config.json").read_text()))) |
| model.load_state_dict(load_file(str(d / "model.safetensors"))) |
| return model.eval().to(device), CharVocab.load(d / "vocab.json") |
|
|
|
|
| def read_audio(path: str) -> torch.Tensor: |
| wav, sr = sf.read(path, dtype="float32", always_2d=True) |
| wav = torch.from_numpy(wav).mean(1) |
| return AF.resample(wav, sr, 16000) if sr != 16000 else wav |
|
|
|
|
| @torch.no_grad() |
| def transcribe(model: PrunableCTC, vocab: CharVocab, paths: list[str]) -> list[str]: |
| device = next(model.parameters()).device |
| out = [] |
| for p in paths: |
| wav = read_audio(p)[None].to(device) |
| lens = torch.tensor([wav.shape[1]], device=device) |
| normalize_audio_(wav, lens) |
| with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): |
| res = model(wav, lens) |
| ids = res["logits"][0, : int(res["frame_lens"][0])].argmax(-1) |
| out.append(vocab.greedy(ids.tolist())) |
| return out |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("audio", nargs="+") |
| ap.add_argument("--model-dir", default=str(Path(__file__).parent)) |
| ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| args = ap.parse_args() |
| model, vocab = load(args.model_dir, args.device) |
| for p, text in zip(args.audio, transcribe(model, vocab, args.audio)): |
| print(f"{p}\t{text}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
