Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
66 lines
3 KiB
Python
66 lines
3 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the original NeMo Canary model on the same fixed chunks as the native ASR path."""
|
|
|
|
import argparse
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import soundfile as sf
|
|
import torch
|
|
from nemo.collections.asr.models import EncDecMultiTaskModel
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--model", type=Path, required=True)
|
|
parser.add_argument("--audio", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--chunk-seconds", type=float, default=40)
|
|
parser.add_argument("--language", default="en")
|
|
parser.add_argument("--target-language", default=None)
|
|
parser.add_argument("--pnc", choices=("yes", "no"), default="yes")
|
|
parser.add_argument("--device", default="cuda")
|
|
args = parser.parse_args()
|
|
torch.set_num_threads(8)
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
|
torch.backends.cudnn.allow_tf32 = False
|
|
signal, sample_rate = sf.read(args.audio, dtype="float32")
|
|
if sample_rate != 16000 or signal.ndim != 1:
|
|
raise ValueError("Parity input must be mono 16 kHz")
|
|
if not 0.02 <= args.chunk_seconds <= 40 or len(signal) < 320:
|
|
raise ValueError("Canary requires at least 20 ms of audio and chunks between 20 ms and 40 s")
|
|
width = int(args.chunk_seconds * sample_rate)
|
|
boundaries = list(range(0, len(signal), width)) + [len(signal)]
|
|
if len(boundaries) > 2 and boundaries[-1] - boundaries[-2] < 320:
|
|
boundaries[-2] = boundaries[-1] - 320
|
|
if boundaries[-2] - boundaries[-3] < 320:
|
|
raise ValueError("Chunk duration is too short to retain a 20 ms final chunk")
|
|
model = EncDecMultiTaskModel.restore_from(str(args.model), map_location="cpu").eval().to(args.device)
|
|
model.cfg.decoding.beam.beam_size = 1
|
|
model.change_decoding_strategy(model.cfg.decoding)
|
|
chunks = []
|
|
with torch.inference_mode():
|
|
for start, end in zip(boundaries, boundaries[1:]):
|
|
samples = signal[start:end]
|
|
if args.device == "cuda":
|
|
torch.cuda.synchronize()
|
|
begin = time.perf_counter()
|
|
result = model.transcribe([samples], batch_size=1, verbose=False,
|
|
source_lang=args.language, target_lang=args.target_language or args.language,
|
|
pnc=args.pnc, timestamps=False)[0]
|
|
if args.device == "cuda":
|
|
torch.cuda.synchronize()
|
|
seconds = time.perf_counter() - begin
|
|
item = {"start_sample": start, "samples": len(samples), "text": result.text,
|
|
"seconds": seconds, "tokens": result.y_sequence.tolist()}
|
|
chunks.append(item)
|
|
print(json.dumps(item), flush=True)
|
|
args.output.write_text(json.dumps({"chunks": chunks,
|
|
"text": " ".join(item["text"].strip() for item in chunks),
|
|
"audio_seconds": len(signal) / sample_rate,
|
|
"seconds": sum(item["seconds"] for item in chunks)}, indent=2) + "\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|