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
84 lines
4.1 KiB
Python
84 lines
4.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the official Transformers Cohere ASR reference, without changing dependencies."""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import soundfile as sf
|
|
import torch
|
|
from transformers import AutoProcessor, CohereAsrForConditionalGeneration
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model", type=Path, required=True)
|
|
parser.add_argument("--cases", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--backend", choices=("cpu", "cuda"), default="cuda")
|
|
parser.add_argument("--dtype", choices=("float32", "bfloat16"), default="float32")
|
|
parser.add_argument("--threads", type=int, default=8)
|
|
parser.add_argument("--repeats", type=int, default=3)
|
|
parser.add_argument("--log", action="store_true", required=True)
|
|
args = parser.parse_args()
|
|
logging.basicConfig(level=logging.INFO)
|
|
torch.set_num_threads(args.threads)
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
|
torch.backends.cudnn.allow_tf32 = False
|
|
processor = AutoProcessor.from_pretrained(args.model)
|
|
model = CohereAsrForConditionalGeneration.from_pretrained(
|
|
args.model, dtype=getattr(torch, args.dtype), attn_implementation="sdpa"
|
|
).to(args.backend).eval()
|
|
records = []
|
|
with torch.inference_mode():
|
|
for case in json.loads(args.cases.read_text()):
|
|
waveform, sr = sf.read(case["audio"], dtype="float32")
|
|
if waveform.ndim != 1 or sr != 16000:
|
|
raise ValueError("Reference fixtures must be mono 16 kHz")
|
|
durations = []
|
|
processor.feature_extractor.max_audio_clip_s = case.get("audio_chunk_duration_sec", 35)
|
|
processor.feature_extractor.overlap_chunk_second = min(
|
|
5, processor.feature_extractor.max_audio_clip_s / 2
|
|
)
|
|
source = waveform
|
|
if case.get("audio_chunk_mode") == "fixed":
|
|
chunk_samples = int(processor.feature_extractor.max_audio_clip_s * sr)
|
|
source = [waveform[i:i + chunk_samples] for i in range(0, len(waveform), chunk_samples)]
|
|
for repeat in range(args.repeats + 1):
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize()
|
|
torch.cuda.reset_peak_memory_stats()
|
|
start = time.perf_counter()
|
|
inputs = processor(
|
|
source.copy(), sampling_rate=sr, return_tensors="pt",
|
|
language=case.get("language", "en"), punctuation=case.get("pnc", True),
|
|
)
|
|
chunk_index = inputs.pop("audio_chunk_index", None)
|
|
inputs = inputs.to(args.backend, dtype=model.dtype)
|
|
outputs = model.generate(**inputs, max_new_tokens=case.get("max_tokens", 256), do_sample=False)
|
|
text = processor.decode(
|
|
outputs, skip_special_tokens=True, audio_chunk_index=chunk_index,
|
|
language=case.get("language", "en"),
|
|
)
|
|
if case.get("audio_chunk_mode") == "fixed":
|
|
separator = "" if case.get("language", "en") in ("ja", "zh") else " "
|
|
text = [separator.join(piece.strip() for piece in text if piece.strip())]
|
|
if args.backend == "cuda":
|
|
torch.cuda.synchronize()
|
|
elapsed = time.perf_counter() - start
|
|
if repeat:
|
|
durations.append(elapsed)
|
|
logging.info("case=%s repeat=%d elapsed=%.6f text=%s", case["id"], repeat, elapsed, text)
|
|
records.append({
|
|
"id": case["id"], "text": text, "tokens": outputs.tolist(),
|
|
"seconds": durations, "audio_seconds": len(waveform) / sr,
|
|
"peak_allocated_bytes": torch.cuda.max_memory_allocated() if args.backend == "cuda" else None,
|
|
"backend": args.backend, "dtype": args.dtype, "torch": torch.__version__,
|
|
})
|
|
args.output.write_text(json.dumps(records, ensure_ascii=False, indent=2) + "\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|