| |
| """Transcribes FLEURS Vietnamese test clips with the exported ONNX files and reports WER. |
| |
| Runs the same loop a mobile app does: encoder once, then the merged decoder step by step with its |
| KV cache (use_cache_branch false on the first step, true after), greedy, prompt |
| <|startoftranscript|><|vi|><|transcribe|><|notimestamps|>. Writes out/<size>/wer.json (read by publish.py) |
| and fails when WER exceeds --max-wer. |
| |
| Usage: |
| .venv/bin/python -I validate.py tiny [--count 30] [--max-wer 0.4] |
| """ |
|
|
| import argparse |
| import csv |
| import json |
| import re |
| import sys |
| import time |
| import unicodedata |
| from pathlib import Path |
|
|
| import jiwer |
| import numpy as np |
| import onnxruntime as ort |
| import soundfile as sf |
| from transformers import WhisperFeatureExtractor, WhisperTokenizer |
|
|
| HERE = Path(__file__).resolve().parent |
| FLEURS = HERE / "data" |
| MAX_TOKENS = 224 |
|
|
|
|
| def clips(count: int) -> list[tuple[Path, str]]: |
| rows = [] |
| with open(FLEURS / "fleurs/data/vi_vn/test.tsv", encoding="utf-8") as f: |
| for row in csv.reader(f, delimiter="\t", quoting=csv.QUOTE_NONE): |
| wav = next((FLEURS / "vi").glob(f"*/{row[1]}"), None) |
| if wav is not None: |
| rows.append((wav, row[2])) |
| if len(rows) == count: |
| break |
| return rows |
|
|
|
|
| def normalize(text: str) -> str: |
| text = unicodedata.normalize("NFC", text.lower()) |
| return re.sub(r"\s+", " ", re.sub(r"[^\w\s]", " ", text)).strip() |
|
|
|
|
| class Transcriber: |
| def __init__(self, repo: Path, quantized: bool): |
| suffix = "_quantized" if quantized else "" |
| opts = ort.SessionOptions() |
| opts.intra_op_num_threads = 8 |
| self.encoder = ort.InferenceSession(str(repo / f"onnx/encoder_model{suffix}.onnx"), opts) |
| self.decoder = ort.InferenceSession(str(repo / f"onnx/decoder_model_merged{suffix}.onnx"), opts) |
| self.features = WhisperFeatureExtractor.from_pretrained(repo) |
| self.tokenizer = WhisperTokenizer.from_pretrained(repo) |
| names = ["<|startoftranscript|>", "<|vi|>", "<|transcribe|>", "<|notimestamps|>"] |
| self.prompt = self.tokenizer.convert_tokens_to_ids(names) |
| self.eot = self.tokenizer.eos_token_id |
| self.past = [i for i in self.decoder.get_inputs() if i.name.startswith("past_key_values.")] |
|
|
| def __call__(self, audio: np.ndarray) -> str: |
| feats = self.features(audio, sampling_rate=16000, return_tensors="np").input_features |
| hidden = self.encoder.run(None, {"input_features": feats.astype(np.float32)})[0] |
| |
| cache = {i.name: np.zeros([1, *[d if isinstance(d, int) else 0 for d in i.shape[1:]]], np.float32) |
| for i in self.past} |
| for i in self.past: |
| if not isinstance(i.shape[1], int) or not isinstance(i.shape[3], int): |
| sys.exit(f"error: cannot size the empty cache for {i.name} {i.shape}") |
| tokens, ids = list(self.prompt), np.array([self.prompt], np.int64) |
| outputs = [o.name for o in self.decoder.get_outputs()] |
| for step in range(MAX_TOKENS): |
| feed = {"input_ids": ids, "encoder_hidden_states": hidden, |
| "use_cache_branch": np.array([step > 0]), **cache} |
| result = dict(zip(outputs, self.decoder.run(None, feed))) |
| nxt = int(result["logits"][0, -1].argmax()) |
| if nxt == self.eot: |
| break |
| tokens.append(nxt) |
| for name in cache: |
| present = result[name.replace("past_key_values.", "present.")] |
| |
| if step == 0 or ".decoder." in name: |
| cache[name] = present |
| ids = np.array([[nxt]], np.int64) |
| return self.tokenizer.decode(tokens, skip_special_tokens=True).strip() |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("size") |
| parser.add_argument("--count", type=int, default=30) |
| parser.add_argument("--max-wer", type=float, default=0.4) |
| args = parser.parse_args() |
| repo = HERE / "out" / args.size / "repo" |
| samples = clips(args.count) |
| failed, results = False, {"clips": len(samples)} |
| for quantized in (False, True): |
| run = Transcriber(repo, quantized) |
| hyps, refs, start = [], [], time.time() |
| for wav, ref in samples: |
| audio, rate = sf.read(wav, dtype="float32") |
| assert rate == 16000, wav |
| hyps.append(normalize(run(audio))) |
| refs.append(normalize(ref)) |
| wer = jiwer.wer(refs, hyps) |
| label = "int8" if quantized else "fp32" |
| print(f"{args.size} {label}: WER {wer:.3f} on {len(samples)} clips in {time.time() - start:.0f} s") |
| print(f" ref: {refs[0]}\n hyp: {hyps[0]}") |
| results[label] = round(wer, 4) |
| failed |= wer > args.max_wer |
| (repo.parent / "wer.json").write_text(json.dumps(results) + "\n") |
| if failed: |
| sys.exit(f"error: WER above {args.max_wer}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
