File size: 5,132 Bytes
5b2d29f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | #!/usr/bin/env python3
"""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]
# Shape of an empty cache: [1, heads, 0, head_dim]; dims come from the graph's symbolic shapes.
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.")]
# The encoder (cross-attention) cache only comes out of the first step.
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()
|