# the five quality metrics plus a 3-model comparison driver (raw input vs # qwen base zero-shot vs fine-tuned). all metrics are computed on the held-out # test split (real disfluencyspeech pairs the model never sees). # # the canonical reference is data/pairs/test.json which has rows like # { "raw": , "clean": } import json import re import string from pathlib import Path from typing import Callable, Optional import Levenshtein import torch from tqdm import tqdm # common dictation fillers we track for the removal-rate metric. FILLER_TOKENS = {"um", "uh", "er", "ah", "like", "you know", "i mean", "so", "well"} PUNCT_CHARS = set(".,;:!?-") # ---------- text utilities ---------- def _words(text: str) -> list[str]: return text.lower().strip().split() def _content_words(text: str) -> list[str]: # drop punctuation, lowercase, split. used by faithfulness metric. stripped = "".join(c for c in text if c not in string.punctuation) return stripped.lower().split() def _count_fillers(text: str) -> int: # count occurrences of each filler form as whole words. lower = " " + text.lower() + " " count = 0 for f in FILLER_TOKENS: count += len(re.findall(rf"(? list[tuple[int, str]]: # return (content_word_index, punct_char) for every sentence punctuation # mark, anchored to the index of the preceding content word. content = _content_words(text) raw_tokens = text.split() positions: list[tuple[int, str]] = [] content_idx = -1 for tok in raw_tokens: # strip trailing punctuation from the token to find the content word body = "".join(c for c in tok if c not in string.punctuation) if body: content_idx += 1 tail_punct = "".join(c for c in tok if c in PUNCT_CHARS) for p in tail_punct: positions.append((content_idx, p)) return positions # ---------- per-example metrics ---------- def disfluency_removal_rate(raw: str, out: str) -> Optional[float]: raw_count = _count_fillers(raw) if raw_count == 0: return None survived = _count_fillers(out) removed = max(0, raw_count - survived) return removed / raw_count def punctuation_f1(out: str, ref: str) -> tuple[int, int, int]: # returns (true_positives, predicted, gold) for a corpus-level micro-f1 out_positions = set(_punct_positions(out)) ref_positions = set(_punct_positions(ref)) tp = len(out_positions & ref_positions) return tp, len(out_positions), len(ref_positions) def faithfulness(out: str, ref: str) -> float: out_words = _content_words(out) ref_words = _content_words(ref) if not ref_words: return 1.0 # token-level levenshtein on the lowercased content-only word lists. dist = Levenshtein.distance(" ".join(out_words), " ".join(ref_words)) ref_len = len(" ".join(ref_words)) if ref_len == 0: return 1.0 return max(0.0, 1.0 - dist / ref_len) def length_ratio(out: str, ref: str) -> float: out_words = _content_words(out) ref_words = _content_words(ref) if not ref_words: return 0.0 return len(out_words) / len(ref_words) # ---------- aggregation ---------- def aggregate(rows: list[dict]) -> dict: # given a list of {"raw": ..., "out": ..., "clean": ...} rows, compute # corpus-level metrics. returns the dict shape consumed by the report. disfl = [ d for d in (disfluency_removal_rate(r["raw"], r["out"]) for r in rows) if d is not None ] tp_sum = pred_sum = gold_sum = 0 faithful_vals: list[float] = [] length_vals: list[float] = [] pass_count = 0 pass_thresholds = { "disfluency": 0.95, "punct_f1": 0.85, "faithfulness": 0.98, "length_min": 0.85, "length_max": 1.05, } per_example = [] for r in rows: tp, pred, gold = punctuation_f1(r["out"], r["clean"]) tp_sum += tp pred_sum += pred gold_sum += gold fa = faithfulness(r["out"], r["clean"]) faithful_vals.append(fa) lr = length_ratio(r["out"], r["clean"]) length_vals.append(lr) d = disfluency_removal_rate(r["raw"], r["out"]) ok = ( (d is None or d >= pass_thresholds["disfluency"]) and fa >= pass_thresholds["faithfulness"] and pass_thresholds["length_min"] <= lr <= pass_thresholds["length_max"] ) per_example.append( { "raw": r["raw"], "out": r["out"], "clean": r["clean"], "disfluency_removal": d, "faithfulness": fa, "length_ratio": lr, } ) if ok: pass_count += 1 precision = tp_sum / pred_sum if pred_sum > 0 else 0.0 recall = tp_sum / gold_sum if gold_sum > 0 else 0.0 punct_f1 = ( 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0 ) return { "num_rows": len(rows), "disfluency_removal_rate": sum(disfl) / len(disfl) if disfl else None, "punctuation_precision": precision, "punctuation_recall": recall, "punctuation_f1": punct_f1, "faithfulness_mean": sum(faithful_vals) / len(faithful_vals) if faithful_vals else 0.0, "length_ratio_mean": sum(length_vals) / len(length_vals) if length_vals else 0.0, "pass_rate": pass_count / len(rows) if rows else 0.0, "per_example": per_example[:50], # keep a sample for the report } # ---------- model generators ---------- def make_raw_generator() -> Callable[[str], str]: # baseline 1: no cleanup at all. lets us measure what shipping nothing # looks like on the same test set. def gen(raw: str) -> str: return raw return gen def make_qwen_generator(model_id_or_path: str, adapter_path: Optional[str] = None) -> Callable[[str], str]: # baseline 2 (no adapter) or row 3 (with adapter): qwen 0.5b. greedy # decode, max_new_tokens capped so the model cannot balloon into a chat # reply. from transformers import AutoModelForCausalLM, AutoTokenizer from cleanup.prompts import build_messages tokenizer = AutoTokenizer.from_pretrained(model_id_or_path, use_fast=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( model_id_or_path, torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, device_map="auto" if torch.cuda.is_available() else None, ) if adapter_path: from peft import PeftModel model = PeftModel.from_pretrained(model, adapter_path) model.eval() def gen(raw: str) -> str: messages = build_messages(raw) prompt = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) inputs = tokenizer(prompt, return_tensors="pt").to(model.device) raw_token_count = len(tokenizer.encode(raw)) max_new = min(256, max(8, int(raw_token_count * 1.6))) with torch.no_grad(): out_ids = model.generate( **inputs, do_sample=False, max_new_tokens=max_new, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id, ) new_tokens = out_ids[0][inputs.input_ids.shape[1]:] text = tokenizer.decode(new_tokens, skip_special_tokens=True) return text.strip() return gen def evaluate_one(test_rows: list[dict], generator: Callable[[str], str]) -> dict: out_rows = [] for row in tqdm(test_rows, desc="generating"): out = generator(row["raw"]) out_rows.append({"raw": row["raw"], "clean": row["clean"], "out": out}) return aggregate(out_rows) def write_eval(report: dict, run_dir: Path) -> None: run_dir = Path(run_dir) (run_dir / "eval.json").write_text(json.dumps(report, indent=2, ensure_ascii=False))