#!/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()