evaluate.py

"""F0.5 auf Falko-MERLIN (dev/test) und dem eigenen Hard-Benchmark (bench-dev: 1000, bench: 4000) + Korrektur-CLI.

python evaluate.py --data data --split dev                           # Scorer-Check: Oracle und Identität
python evaluate.py --data data --split bench --ckpt ckpts/ckpt-*.pt # mehrere Checkpoints -> gemittelt
python evaluate.py --data data --ckpt best.pt --correct "ich weis nicht das es so ist"
"""
import argparse
import difflib
import json
import os

import sentencepiece as spm
import torch

from model import EOS, GEC
from prepare import _pair_split, detokenize, split_sentences, tokenize


def load_model(paths, device):
    states = [torch.load(p, map_location="cpu", weights_only=False) for p in paths]
    avg = {k: sum(s["model"][k].float() for s in states) / len(states) for k in states[0]["model"]}
    model = GEC(**states[0]["cfg"])
    model.load_state_dict(avg)
    return model.to(device).eval()


def correct_tokens(model, sp, sents, beam=5):
    """sents: tokenisierte Sätze -> korrigierte tokenisierte Sätze."""
    dev = next(model.parameters()).device
    out = []
    with torch.autocast(dev.type, dtype=torch.float16, enabled=dev.type == "cuda"):
        for s in sents:
            ids = sp.encode(s)[: model.cfg["max_len"] - 1] + [EOS]
            out.append(sp.decode(model.beam_search(torch.tensor(ids, device=dev), beam=beam)))
    return out


def parse_m2(path):
    """-> [(Quelltokens, {(start, ende, korrektur)})], nur Annotator 0."""
    data = []
    for block in open(path, encoding="utf-8").read().strip().split("\n\n"):
        lines = block.split("\n")
        edits = set()
        for a in lines[1:]:
            span, typ, cor, *_, annot = a[2:].split("|||")
            if typ != "noop" and annot == "0":
                s, e = map(int, span.split())
                edits.add((s, e, cor))
        data.append((lines[0][2:], edits))
    return data


def _token_edits(a, b, off):
    """Levenshtein über Tokens -> Einzel-Token-Edits wie in den Gold-M2 (ähnliche Wörter werden bevorzugt ersetzt)."""
    D = [[i + j if not i * j else 0 for j in range(len(b) + 1)] for i in range(len(a) + 1)]
    sub = lambda x, y: 2 - difflib.SequenceMatcher(None, x.lower(), y.lower()).ratio()
    for i in range(1, len(a) + 1):
        for j in range(1, len(b) + 1):
            D[i][j] = min(D[i - 1][j] + 1, D[i][j - 1] + 1, D[i - 1][j - 1] + sub(a[i - 1], b[j - 1]))
    edits, i, j = set(), len(a), len(b)
    while i or j:
        if i and j and D[i][j] == D[i - 1][j - 1] + sub(a[i - 1], b[j - 1]):
            edits.add((off + i - 1, off + i, b[j - 1])); i -= 1; j -= 1
        elif j and D[i][j] == D[i][j - 1] + 1:
            edits.add((off + i, off + i, b[j - 1])); j -= 1
        else:
            edits.add((off + i - 1, off + i, "")); i -= 1
    return edits


def hyp_edits(src, hyp, gold=frozenset()):
    """Wie der M2-Scorer (MaxMatch): pro Änderungsblock die Zerlegung wählen, die am besten zu Gold passt."""
    # ponytail: nur zwei Zerlegungen (ganzer Block / pro Token) statt voller MaxMatch-Lattice -> Werte annähernd literaturvergleichbar
    a, b = src.split(), hyp.split()
    edits = set()
    for tag, i1, i2, j1, j2 in difflib.SequenceMatcher(None, a, b, autojunk=False).get_opcodes():
        if tag != "equal":
            merged = {(i1, i2, " ".join(b[j1:j2]))}
            split = _token_edits(a[i1:i2], b[j1:j2], i1)
            edits |= max((merged, split), key=lambda e: (len(e & gold), -len(e)))
    return edits


def ref_edits(src, trg):
    """Gold-Edits aus einem Paar ohne M2-Annotation (Hard-Benchmark), pro Token geschnitten."""
    a, b = src.split(), trg.split()
    ops = difflib.SequenceMatcher(None, a, b, autojunk=False).get_opcodes()
    return set().union(*(_token_edits(a[i1:i2], b[j1:j2], i1) for tag, i1, i2, j1, j2 in ops if tag != "equal"))


def _bench_items(rows):
    items = []
    for r in rows:
        src = tokenize(r["input"])
        items.append((src, ref_edits(src, tokenize(r["target"])), [tokenize(x) for x in _pair_split(r["input"])]))
    return items, [tokenize(r["target"]) for r in rows]


def load_split(data, split):
    """-> [(Quelle tokenisiert, Gold-Edits, zu korrigierende Sätze)], [Zieltexte fürs Oracle].
    dev-a/dev-b = erste/zweite Hälfte von Falko-dev, hh-dev = real/hh-dev.jsonl (Schema wie bench.jsonl)."""
    if split in ("dev", "test", "dev-a", "dev-b"):
        name = "dev" if split.startswith("dev") else split
        m2 = parse_m2(f"{data}/real/fm-{name}.m2")
        trg = open(f"{data}/real/fm-{name}.trg", encoding="utf-8").read().split("\n")
        items, trg = [(s, g, [s]) for s, g in m2], trg[: len(m2)]
        h = len(items) // 2
        part = {"dev-a": slice(None, h), "dev-b": slice(h, None)}.get(split, slice(None))
        return items[part], trg[part]
    path = f"{data}/real/" + ("hh-dev.jsonl" if split == "hh-dev" else "bench.jsonl")
    if not os.path.exists(path):
        raise FileNotFoundError(f"Split '{split}': {path} fehlt")
    rows = [json.loads(l) for l in open(path, encoding="utf-8")]
    if split != "hh-dev":
        rows = rows[:1000] if split == "bench-dev" else rows[1000:]
    return _bench_items(rows)


def group_sentences(sents, sp, limit):
    """Sätze zu Gruppen bis `limit` Tokens bündeln (Kontext!). Ein Satz, der allein nicht passt, wird an Wortgrenzen
    geteilt: correct_tokens würde ihn sonst nach `limit` Tokens abschneiden und den Rest still verlieren."""
    n = lambda words: len(sp.encode(" ".join(words)))
    pieces = []
    for s in sents:
        if n([s]) <= limit:
            pieces.append(s)
            continue
        cur = []
        for w in s.split():
            if cur and n(cur + [w]) > limit:
                pieces.append(" ".join(cur))
                cur = []
            cur.append(w)
        pieces.append(" ".join(cur))
    groups, cur = [], []
    for p in pieces:
        if cur and n(cur + [p]) > limit:
            groups.append(" ".join(cur))
            cur = []
        cur.append(p)
    return groups + [" ".join(cur)]


def correct_text(model, sp, sents, beam=5):
    """Tokenisierte Sätze -> in Gruppen korrigiert (so viel Kontext, wie ins Modell passt) und zusammengesetzt."""
    return " ".join(correct_tokens(model, sp, group_sentences(sents, sp, model.cfg["max_len"] - 1), beam))


def correct_items(model, sp, items, beam=5):
    return [correct_text(model, sp, sents, beam) for _, _, sents in items]


def f05(items, hyps):
    tp = fp = fn = 0
    for (src, gold, *_), hyp in zip(items, hyps):
        h = hyp_edits(src, hyp, gold)
        tp, fp, fn = tp + len(h & gold), fp + len(h - gold), fn + len(gold - h)
    p = tp / (tp + fp) if tp + fp else 1.0
    r = tp / (tp + fn) if tp + fn else 1.0
    f = 1.25 * p * r / (0.25 * p + r) if p + r else 0.0
    return p, r, f


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="data")
    ap.add_argument("--split", default="dev", choices=["dev", "test", "bench-dev", "bench"])
    ap.add_argument("--ckpt", nargs="*")
    ap.add_argument("--beam", type=int, default=5)
    ap.add_argument("--n", type=int, help="nur die ersten n Beispiele")
    ap.add_argument("--correct", help="beliebigen Text korrigieren")
    a = ap.parse_args()
    items, trg = load_split(a.data, a.split)
    items, trg = items[: a.n], trg[: a.n]

    if not a.ckpt:  # Scorer-Selbsttest: Zieltexte müssen nahe 1.0 liegen, Identität bei 0
        p, r, f = f05(items, trg)
        print(f"Oracle    P={p:.3f} R={r:.3f} F0.5={f:.3f}")
        assert f > 0.85, "Scorer weicht zu stark von den Gold-Edits ab"
        assert all(" ".join(sents) == src for src, _, sents in items), "Satzzerlegung verändert Tokens"
        p, r, f = f05(items, [s for s, *_ in items])
        print(f"Identität P={p:.3f} R={r:.3f} F0.5={f:.3f}")
        assert r == 0
        raise SystemExit

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = load_model(a.ckpt, device)
    sp = spm.SentencePieceProcessor(model_file=f"{a.data}/spm.model")
    if a.correct:
        sents = [tokenize(s) for s in split_sentences(a.correct) if s]
        print(detokenize(correct_text(model, sp, sents, a.beam)))
        raise SystemExit
    hyps = correct_items(model, sp, items, a.beam)
    with open(f"hyp-{a.split}.txt", "w", encoding="utf-8") as fh:
        fh.write("\n".join(hyps) + "\n")
    p, r, f = f05(items, hyps)
    print(f"{a.split}: P={p:.3f} R={r:.3f} F0.5={f:.3f}  (Hypothesen in hyp-{a.split}.txt)")