source/src/vimeml/benchmarks/evaluate_ajimee.py

"""Evaluate frozen AzooKey candidates, preserving official references and order."""
import argparse
import json
import math
import time
from pathlib import Path

from vimeml.benchmarks.ajimee import DEVELOPMENT_FORMAT, FORMAT, ROOT, convert_items, json_bytes, sha


def read_json(path):
    return json.loads(path.read_text(encoding="utf-8-sig"))


def load_export(directory):
    manifest = read_json(directory / "manifest.json")
    if manifest.get("format") not in (FORMAT, DEVELOPMENT_FORMAT) or manifest.get("status") != "complete":
        raise ValueError("Expected a complete prepared AJIMEE manifest.")
    for name, digest in manifest["files_sha256"].items():
        if sha((directory / name).read_bytes()) != digest:
            raise ValueError(f"Prepared file changed: {name}.")
    source = directory / "evaluation_items.json"
    if sha(source.read_bytes()) != manifest["source_sha256"]:
        raise ValueError("Source snapshot hash mismatch.")
    inputs, mapping, stats = convert_items(read_json(source))
    if inputs != read_json(directory / "ajimee-input.json") or mapping != read_json(directory / "case-map.json") or stats != manifest["stats"]:
        raise ValueError("Prepared input/mapping does not match source snapshot.")
    exported = directory / "ajimee-results"
    if inputs != read_json(exported / "ajimee-input.json"):
        raise ValueError("Mac input differs from frozen local input.")
    raw = read_json(exported / "azookey-candidates.json")
    items, n_best = raw.get("items"), raw.get("n_best")
    if type(n_best) is not int or n_best < 1 or not isinstance(items, list) or len(items) != len(inputs):
        raise ValueError("Invalid n_best or exported sample count.")
    rows = []
    for position, (expected, item, case) in enumerate(zip(inputs, items, mapping)):
        if any(item.get(key) != expected[value] for key, value in
               (("query", "query"), ("answers", "answer"), ("left_context", "left_context"))):
            raise ValueError(f"Row {position}: query/context/references misaligned.")
        if item.get("right_context") not in (None, ""):
            raise ValueError(f"Row {position}: unexpected right context.")
        candidates = item.get("outputs")
        if not isinstance(candidates, list) or len(candidates) > n_best:
            raise ValueError(f"Row {position}: invalid candidate count.")
        texts = []
        for candidate in candidates:
            text, score = candidate.get("text"), candidate.get("score")
            if not isinstance(text, str) or not text or type(score) not in (int, float) or not math.isfinite(score):
                raise ValueError(f"Row {position}: invalid candidate text/score.")
            texts.append(text)
        if len(set(texts)) != len(texts):
            raise ValueError(f"Row {position}: duplicate candidate text.")
        rank = next((i for i, text in enumerate(texts) if text in case["answers"]), -1)
        if item.get("max_rank") != rank:
            raise ValueError(f"Row {position}: CLI rank inconsistent with candidates.")
        rows.append({**case, "candidates": candidates, "orders": {"azookey": texts}, "fallbacks": {}})
    # Preserve provenance verbatim; the CLI output does not embed every flag.
    version_files = ["converter-version.txt", "dictionary-versions.txt", "swift-version.txt"]
    provenance = {name: (exported / name).read_text(encoding="utf-8-sig") for name in version_files}
    for line in provenance["dictionary-versions.txt"].splitlines():
        if not line.startswith(" "):
            raise ValueError("Dictionary submodule is uninitialized, conflicted or differs from its pinned commit.")
    provenance["files_sha256"] = {name: sha((exported / name).read_bytes()) for name in
                                  ["ajimee-input.json", "azookey-candidates.json", *version_files]}
    provenance["n_best"] = n_best
    provenance["cli_execution_seconds"] = raw.get("execution_time")
    provenance["flag_provenance"] = "Export requested with typo mode off and no Zenzai; CLI JSON does not embed these flags."
    return manifest, provenance, rows


def edit_distance(reference, hypothesis):
    previous = list(range(len(hypothesis) + 1))
    for i, left in enumerate(reference, 1):
        current = [i]
        for j, right in enumerate(hypothesis, 1):
            current.append(min(current[-1] + 1, previous[j] + 1, previous[j - 1] + (left != right)))
        previous = current
    return previous[-1]


def min_cer(answers, hypothesis):
    return min(edit_distance(answer, hypothesis) / len(answer) for answer in answers)


def metrics(rows, method):
    count = len(rows)
    if not count:
        return {"cases": 0}
    top1 = top5 = covered = corrected = regressed = 0
    cer = reciprocal = 0.0
    for row in rows:
        order, answers = row["orders"][method], row["answers"]
        ranks = [i + 1 for i, text in enumerate(order) if text in answers]
        hit = bool(ranks) and min(ranks) == 1
        original = bool(row["orders"]["azookey"]) and row["orders"]["azookey"][0] in answers
        top1 += hit
        top5 += bool(ranks) and min(ranks) <= 5
        covered += bool(ranks)
        reciprocal += 1 / min(ranks) if ranks else 0
        cer += min_cer(answers, order[0] if order else "")
        corrected += hit and not original
        regressed += original and not hit
    return {"cases": count, "top1_correct": top1, "top1_accuracy": top1 / count,
            "top5_correct": top5, "top5_accuracy": top5 / count, "candidate_pool_covered": covered,
            "candidate_pool_recall": covered / count, "covered_top1_accuracy": top1 / covered if covered else None,
            "mean_min_cer": cer / count, "mrr": reciprocal / count,
            "corrected_vs_azookey": corrected, "regressed_vs_azookey": regressed,
            "fallback_cases": sum(method in row["fallbacks"] for row in rows)}


def eligibility(lm, context, texts):
    if not texts:
        return "empty_candidates"
    context_ids = lm.processor.encode(context, out_type=int)
    if lm.processor.decode(context_ids) != context:
        return "context_roundtrip_mismatch"
    if len(context_ids) + 1 > lm.model.config.context_length:
        return "context_length"
    for text in texts:
        ids = lm.processor.encode(context + text, out_type=int)
        if lm.processor.decode(ids) != context + text:
            return "candidate_roundtrip_mismatch"
        if not ids or len(ids) > lm.model.config.context_length:
            return "context_length"
    return None


def rerank(lm, rows):
    from vimeml.training.evaluate_ime import score_candidates
    for position, row in enumerate(rows):
        texts = row["orders"]["azookey"]
        row["lm_scores"] = {}
        for field, context, methods in (
            ("contextual", row["left_context"], ("lm_context_sum", "lm_context_mean_secondary")),
            ("context_free", "", ("lm_no_context_sum",)),
        ):
            reason = eligibility(lm, context, texts)
            if reason:
                for method in methods:
                    row["orders"][method] = list(texts)
                    row["fallbacks"][method] = reason
                continue
            scored = score_candidates(lm, context, texts)
            if any(not math.isfinite(c[key]) for c in scored["candidates"] for key in
                   ("log_probability_sum", "log_probability_mean")):
                raise ValueError(f"Nonfinite LM score: {row['id']}.")
            row["lm_scores"][field] = scored
            for method in methods:
                key = "log_probability_mean" if method.endswith("secondary") else "log_probability_sum"
                # Stable sort: ties preserve the actual AzooKey candidate order.
                row["orders"][method] = [c["text"] for c in sorted(scored["candidates"], key=lambda c: c[key], reverse=True)]
        if (position + 1) % 25 == 0 or position == len(rows) - 1:
            print(f"Evaluated {position + 1}/{len(rows)}", flush=True)


def summarize(rows):
    groups = {"all": rows, "with_context": [r for r in rows if r["left_context"]],
              "without_context": [r for r in rows if not r["left_context"]]}
    if "lm_context_sum" in rows[0]["orders"]:
        groups["contextually_scorable"] = [r for r in rows if "lm_context_sum" not in r["fallbacks"]]
    return {name: {method: metrics(items, method) for method in rows[0]["orders"]} for name, items in groups.items()}


def write_summary(path, report, rows):
    lines = ["# AJIMEE / AzooKey + Tiny LM", "",
             "完整输入、全部可接受答案、真实候选池；不注入答案，不正规化文字，不加EOS。",
             "上下文使用官方给定左文，可能跨句；联合分词的公共前缀后logP是候选评分代理，不是精确字符串条件概率。",
             "主指标为sum，mean仅作预先声明的次要诊断；没有在本评测集调组合系数。并列保持AzooKey原顺序。",
             "任何候选超出128-token评分窗口时，该策略整条回退原排序；空候选计失败，仍保留完整分母。",
             "正确答案未进入候选池的样本也保留；训练语料与公开评测文本重叠未知，分句DataLoader的test split没有用于本轮评测。", "",
             "| 样本组 | 策略 | 样本数 | Top-1 | Top-5 | 候选池覆盖 | MinCER | 纠正 / 改坏 | 回退 |",
             "| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |"]
    for group, methods in report["metrics"].items():
        for method, item in methods.items():
            if not item["cases"]:
                continue
            lines.append(f"| {group} | {method} | {item['cases']} | {item['top1_correct']} ({item['top1_accuracy']:.1%}) | "
                         f"{item['top5_correct']} ({item['top5_accuracy']:.1%}) | {item['candidate_pool_covered']} ({item['candidate_pool_recall']:.1%}) | "
                         f"{item['mean_min_cer']:.4f} | {item['corrected_vs_azookey']} / {item['regressed_vs_azookey']} | {item['fallback_cases']} |")
    lines += ["", "候选池覆盖是这次冻结候选下精确命中的上限，不等于理论上能达到的实际模型性能。开发机耗时不是iOS延迟。",
              "候选原分数与全部LM token分数见scores.jsonl；changed-cases.json记录纠正/改坏样本；metrics.json记录输入、候选、版本和模型hash。", ""]
    path.write_text("\n".join(lines), encoding="utf-8")


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--benchmark", type=Path, default=ROOT / "artifacts/benchmarks/ajimee-jwtd-v2-v1")
    parser.add_argument("--checkpoint", type=Path, default=ROOT / "artifacts/models/tiny-ja-v1/best.pt")
    parser.add_argument("--tokenizer", type=Path, default=ROOT / "artifacts/tokenizers/ja-unigram-16k-v1")
    parser.add_argument("--output", type=Path, default=ROOT / "outputs/ime-eval/tiny-ja-v1-ajimee")
    parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
    parser.add_argument("--threads", type=int, default=4)
    parser.add_argument("--baseline-only", action="store_true")
    args = parser.parse_args(argv)
    if args.threads < 1:
        parser.error("threads must be positive.")
    if args.output.exists() and any(args.output.iterdir()):
        parser.error("Output is not empty; choose a new --output directory to preserve previous results.")
    started = time.perf_counter()
    manifest, provenance, rows = load_export(args.benchmark)
    model = None
    if not args.baseline_only:
        import torch
        from vimeml.training.infer import JapaneseLM
        torch.set_num_threads(args.threads)
        lm = JapaneseLM(args.checkpoint, args.tokenizer, args.device)
        model = lm.metadata
        rerank(lm, rows)
    report = {"format": "ajimee_ime_metrics_v1", "status": "complete", "model": model,
              "benchmark_manifest": manifest, "export_provenance": provenance, "metrics": summarize(rows),
              "policy": "Original queries/references/candidates; supplied context; suffix logP sum primary; mean secondary; stable ties; whole-case fallback; no EOS or truncation.",
              "training_overlap": "Unknown; no corpus overlap audit performed.", "test_split_used": False,
              "empty_candidate_cases": sum(not r["candidates"] for r in rows),
              "elapsed_seconds": time.perf_counter() - started}
    changed = []
    for row in rows:
        original = row["orders"]["azookey"]
        before = bool(original) and original[0] in row["answers"]
        for method, order in row["orders"].items():
            after = bool(order) and order[0] in row["answers"]
            if before != after:
                changed.append({"id": row["id"], "method": method, "change": "corrected" if after else "regressed",
                                "query": row["query"], "context": row["left_context"], "answers": row["answers"],
                                "before": original[0], "after": order[0]})
    args.output.mkdir(parents=True, exist_ok=True)
    (args.output / "scores.jsonl").write_text("".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows), encoding="utf-8")
    (args.output / "changed-cases.json").write_bytes(json_bytes(changed))
    report["files_sha256"] = {name: sha((args.output / name).read_bytes()) for name in ("scores.jsonl", "changed-cases.json")}
    (args.output / "metrics.json").write_bytes(json_bytes(report))
    write_summary(args.output / "results.md", report, rows)
    print(json.dumps(report["metrics"]["all"], ensure_ascii=False, indent=2))
    print(f"Report: {args.output.resolve() / 'results.md'}")


if __name__ == "__main__":
    main()