63 lines
2.1 KiB
Python
63 lines
2.1 KiB
Python
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}%)") |