import torch
from datasets import load_dataset, Audio
from transformers import WhisperProcessor, WhisperForConditionalGeneration
import jiwer

MODEL_ID = "openai/whisper-medium"
NUM_SAMPLES = 50  # how many to eval, None = full set

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32
print(f"[*] device: {device}")

processor = WhisperProcessor.from_pretrained(MODEL_ID)
model = WhisperForConditionalGeneration.from_pretrained(MODEL_ID, torch_dtype=dtype).to(device)
model.eval()
print(f"[+] base model loaded. params: {sum(p.numel() for p in model.parameters()):,}")

print("[*] loading common voice ru test split...")
ds = load_dataset("fsicoli/common_voice_17_0", "ru", split="test", trust_remote_code=True)
ds = ds.cast_column("audio", Audio(sampling_rate=16000))
if NUM_SAMPLES:
    ds = ds.select(range(NUM_SAMPLES))

references, hypotheses = [], []

forced_decoder_ids = processor.get_decoder_prompt_ids(language="russian", task="transcribe")

for i, sample in enumerate(ds):
    audio = sample["audio"]
    inputs = processor(
        audio["array"],
        sampling_rate=audio["sampling_rate"],
        return_tensors="pt"
    )
    input_features = inputs.input_features.to(device, dtype=dtype)

    with torch.no_grad():
        predicted_ids = model.generate(
            input_features,
            forced_decoder_ids=forced_decoder_ids,
        )

    text = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0].strip()
    ref = sample["sentence"].strip()
    references.append(ref)
    hypotheses.append(text)

    if i < 5:  # print first few so you see something's happening
        print(f"  [{i}] ref : {ref!r}")
        print(f"  [{i}] hyp : {text!r}")

wer = jiwer.wer(references, hypotheses)
print(f"\n[+] base model WER on {len(ds)} samples: {wer:.4f} ({wer*100:.2f}%)")