File size: 9,500 Bytes
57e91f2 | 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 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 | #!/usr/bin/env python3
"""Checks the exported models and the reference pipeline, and measures what the refiner does to real
PhoWhisper output.
1. Tokenizer: qwen_refine.Tokenizer gives the same ids as Hugging Face's tokenizer on every FLEURS sentence,
and the prompt equals the model's own chat template.
2. Refiner, on data/asr_hyps.tsv (asr_hyps.py): word error rate against the reference before and after
(case and punctuation ignored), punctuation F1 on the words both share, how often the guard falls back, and
the time per transcript.
3. Fixtures for the Kotlin and Swift checks in out/<size>/fixtures: cases.tsv (input, expected output, both as
JSON strings), tokens.tsv (text as a JSON string, ids) and generations.tsv (prompt ids, generated ids).
Usage:
.venv/bin/python -I validate.py [--size 1.5B 3B] [--count 300]
"""
import argparse
import csv
import json
import sys
import time
import unicodedata
from difflib import SequenceMatcher
from pathlib import Path
import jiwer
import regex
from tokenizers import Tokenizer as HfTokenizer
sys.path.insert(0, str(Path(__file__).resolve().parent)) # -I drops the script's own directory
from common import HERE, SIZES, fleurs_sentences, out_dir, src_dir, strip # noqa: E402
from qwen_refine import EXAMPLES, SYSTEM, Refiner, Tokenizer, prepare, prompt_prefix, prompt_suffix # noqa: E402
MARKS = ".,?!:"
WER_TOLERANCE = 0.005
# Spoken-style transcripts for the bullet and fix checks; FLEURS has none of these.
CASES = [
"ờ hôm nay chúng ta sẽ họp về ba vấn đề một là ngân sách quý tới hai là tuyển thêm nhân viên và ba là kế hoạch du lịch công ty",
"thủ tướng phạm minh chính đã tiếp bộ trưởng nông nghiệp mỹ tại hà lội vào sáng nay",
"bạn có biết mấy giờ rồi không",
"nhớ mua giúp mình mấy thứ nhé trứng sữa tươi bánh mì và một ít táo",
"các bước cài đặt như sau đầu tiên tải ứng dụng sau đó đăng nhập bằng số điện thoại cuối cùng là bật thông báo",
"ừm thì là cuộc họp ngày mai được dời sang hai giờ chiều nha mọi người",
"tôi tôi nghĩ là chúng ta nên nên kiểm tra lại số liệu trước khi gửi báo cáo cho khách hàng",
"em ơi cho anh hỏi đường đến chợ bến thành đi như thế nào",
"hôm qua trời mưa to quá nên tôi ở nhà đọc sách cả ngày",
"dự án cần hoàn thành trước ngày mười lăm tháng mười hai năm hai nghìn không trăm hai mươi sáu",
]
def check_tokenizer(size: str) -> dict:
ours = Tokenizer(out_dir(size))
hf = HfTokenizer.from_file(str(src_dir(size) / "tokenizer.json"))
texts = [t for _, t in fleurs_sentences()] + [s for s, _ in fleurs_sentences()] + CASES
texts += [SYSTEM] + [t for pair in EXAMPLES for t in pair]
texts += [" hai khoảng\ttrắng\n\nxuống dòng ", "It's 5:30 AM — ok?", "emoji 🙂 và ký tự lạ ½ fi"]
mismatches = [t for t in texts if ours.encode(t) != hf.encode(t, add_special_tokens=False).ids]
for t in mismatches[:5]:
print(" tokenizer mismatch:", repr(t))
# The prompt must be exactly what the chat template gives (the model was tuned on that format).
from transformers import AutoTokenizer
auto = AutoTokenizer.from_pretrained(src_dir(size))
messages = [{"role": "system", "content": SYSTEM}]
for transcript, edited in EXAMPLES:
messages += [{"role": "user", "content": transcript}, {"role": "assistant", "content": edited}]
messages.append({"role": "user", "content": CASES[0]})
text = auto.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
template = auto(text, add_special_tokens=False).input_ids
ours_prompt = prompt_prefix(ours) + prompt_suffix(ours, CASES[0])
roundtrip = sum(ours.decode(ours.encode(t)) != unicodedata.normalize("NFC", t) for t in texts)
return {
"texts": len(texts),
"mismatches": len(mismatches),
"chat_template_equal": template == ours_prompt,
"decode_roundtrip_failures": roundtrip,
}
def marks_by_word(text: str) -> list[tuple[str, str]]:
"""(word, mark after it) for each word; a bullet line break counts as a full stop."""
out = []
for token in regex.findall(r"[\p{L}\p{N}]+|[.,?!:;\n]", unicodedata.normalize("NFC", text.lower())):
if token in MARKS + ";\n":
mark = "." if token in ".!;\n" else token
if out and out[-1][1] == "":
out[-1] = (out[-1][0], mark)
else:
out.append((token, ""))
return out
def punctuation_counts(hyp: str, ref: str) -> tuple[int, int, int]:
"""True positives, predicted and reference marks on the words both texts share (aligned by difflib)."""
h, r = marks_by_word(hyp), marks_by_word(ref)
tp = pred = gold = 0
matcher = SequenceMatcher(a=[w for w, _ in h], b=[w for w, _ in r], autojunk=False)
for block in matcher.get_matching_blocks():
for k in range(block.size):
hm, rm = h[block.a + k][1], r[block.b + k][1]
pred += hm != ""
gold += rm != ""
tp += hm != "" and hm == rm
return tp, pred, gold
def read_hyps(count: int) -> list[tuple[str, str, str]]:
with open(HERE / "data" / "asr_hyps.tsv", encoding="utf-8") as f:
rows = list(csv.reader(f, delimiter="\t", quoting=csv.QUOTE_NONE, escapechar="\\"))
return [(size, ref, hyp) for size, _, ref, hyp in rows[:count]]
def evaluate(size: str, count: int) -> dict:
refiner = Refiner(out_dir(size), threads=8)
rows = read_hyps(count)
refs, before, after, outputs = [], [], [], []
tp = pred = gold = fallbacks = 0
start = time.time()
for _, ref, hyp in rows:
edited = refiner.refine(hyp)
outputs.append(edited)
fallbacks += edited == " ".join(unicodedata.normalize("NFC", hyp).split())
refs.append(strip(ref))
before.append(strip(hyp))
after.append(strip(edited))
t, p, g = punctuation_counts(edited, ref)
tp, pred, gold = tp + t, pred + p, gold + g
seconds = (time.time() - start) / len(rows)
precision, recall = tp / max(pred, 1), tp / max(gold, 1)
by_size = {}
for asr in sorted({s for s, _, _ in rows}):
idx = [i for i, (s, _, _) in enumerate(rows) if s == asr]
by_size[asr] = {
"wer_before": round(jiwer.wer([refs[i] for i in idx], [before[i] for i in idx]), 4),
"wer_after": round(jiwer.wer([refs[i] for i in idx], [after[i] for i in idx]), 4),
}
cases = [(c, refiner.refine(c)) for c in CASES]
fixtures = out_dir(size).parent / "fixtures"
fixtures.mkdir(exist_ok=True)
with open(fixtures / "cases.tsv", "w", encoding="utf-8") as f:
for text, edited in cases + list(zip([h for _, _, h in rows[:20]], outputs[:20])):
f.write(json.dumps(text, ensure_ascii=False) + "\t" + json.dumps(edited, ensure_ascii=False) + "\n")
# For the Swift check's ONNX Runtime stand-in: the whole prompt and the ids the model generated after it.
with open(fixtures / "generations.tsv", "w", encoding="utf-8") as f:
for text in CASES + [h for _, _, h in rows[:20]]:
prepared = prepare(text)
prompt = refiner.prefix + prompt_suffix(refiner.tok, prepared)
f.write(",".join(map(str, prompt)) + "\t" + ",".join(map(str, refiner.generate(prepared))) + "\n")
with open(fixtures / "tokens.tsv", "w", encoding="utf-8") as f:
for text in CASES + [t for _, t in fleurs_sentences()[:200]] + [SYSTEM]:
ids = refiner.tok.encode(text)
f.write(json.dumps(text, ensure_ascii=False) + "\t" + ",".join(map(str, ids)) + "\n")
for text, edited in cases:
print(f" {text}\n -> {edited!r}")
return {
"transcripts": len(rows),
"wer": by_size,
"punctuation": {"precision": round(precision, 4), "recall": round(recall, 4),
"f1": round(2 * precision * recall / max(precision + recall, 1e-9), 4)},
"guard_fallbacks": fallbacks,
"seconds_per_transcript": round(seconds, 2),
"examples": [{"input": c, "output": e} for c, e in cases],
}
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--size", choices=SIZES, nargs="+", default=list(SIZES))
parser.add_argument("--count", type=int, default=300)
args = parser.parse_args()
failed = False
for size in args.size:
result = {"tokenizer": check_tokenizer(size)}
print(size, "tokenizer:", result["tokenizer"])
result.update(evaluate(size, args.count))
print(size, {k: v for k, v in result.items() if k not in ("examples", "tokenizer")})
(out_dir(size).parent / "validation.json").write_text(json.dumps(result, ensure_ascii=False, indent=1) + "\n")
tok = result["tokenizer"]
failed |= tok["mismatches"] > 0 or not tok["chat_template_equal"] or tok["decode_roundtrip_failures"] > 0
# The merge undoes rewording, so the word error rate should not go up beyond noise (one word in 200).
failed |= any(w["wer_after"] > w["wer_before"] + WER_TOLERANCE for w in result["wer"].values())
if failed:
sys.exit("error: tokenizer mismatch, prompt differs from the chat template, or the refiner adds word errors")
if __name__ == "__main__":
main()
|