File size: 1,749 Bytes
088e8c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | #!/usr/bin/env python3
"""Transcribe the FLEURS clips with native transformers Voxtral (fp32) -> native.json."""
from __future__ import annotations
import json
import time
from pathlib import Path
import soundfile as sf
import torch
REPO = "mistralai/Voxtral-Mini-3B-2507"
HERE = Path(__file__).parent
CLIPS = HERE / "clips"
LANGS = {"en_1": "en", "en_2": "en", "pt_1": "pt", "pt_2": "pt"}
def main() -> None:
from transformers import AutoProcessor, VoxtralForConditionalGeneration
processor = AutoProcessor.from_pretrained(REPO)
model = VoxtralForConditionalGeneration.from_pretrained(REPO, dtype=torch.float32).eval()
results = {}
for path in sorted(CLIPS.glob("*.wav")):
name = path.stem
audio, rate = sf.read(path, dtype="float32")
assert rate == 16_000, rate
inputs = processor.apply_transcription_request(
audio=str(path), model_id=REPO, language=LANGS[name], return_tensors="pt"
)
start = time.perf_counter()
with torch.no_grad():
output = model.generate(**inputs, max_new_tokens=256, do_sample=False)
elapsed = time.perf_counter() - start
text = processor.tokenizer.decode(
output[0, inputs["input_ids"].shape[1] :].tolist(), skip_special_tokens=True
)
results[name] = {
"native": text.strip(),
"duration": len(audio) / rate,
"elapsed": elapsed,
"rtf": elapsed / (len(audio) / rate),
}
print(name, json.dumps(results[name], ensure_ascii=False), flush=True)
with (HERE / "native.json").open("wt", encoding="utf-8") as f:
json.dump(results, f, ensure_ascii=False, indent=2)
if __name__ == "__main__":
main()
|