src/cleanup/eval/metrics.py

# 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": <transcript_a>, "clean": <transcript_c> }

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"(?<!\w){re.escape(f)}(?!\w)", lower))
    return count


def _punct_positions(text: str) -> 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))