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

MODEL_ID = "openai/whisper-medium"
BASE_DIR = "whisper-medium-ru-lora"
LORA_PATH = f"{BASE_DIR}/checkpoint-4000"
NUM_SAMPLES = 500

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)
base_model = WhisperForConditionalGeneration.from_pretrained(MODEL_ID, torch_dtype=dtype).to(device)

model = PeftModel.from_pretrained(base_model, LORA_PATH)
model.eval()

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"[+] LoRA loaded. trainable: {trainable:,} / {total:,} ({100*trainable/total:.3f}%)")

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)

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

    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(f"  [{i}] ref : {ref!r}")
        print(f"  [{i}] hyp : {text!r}")

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