53 lines
1.8 KiB
Python
53 lines
1.8 KiB
Python
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}%)") |