kba-cli 0.3.0rc4__py3-none-any.whl

kba/eval.py

@@ -0,0 +1,332 @@
+"""Labeled-corpus evaluation for extraction and analysis accuracy.
+
+This harness measures behavior on a labeled corpus; synthetic fixtures only prove
+regression/mechanics, not real-world accuracy. No data leaves the machine.
+"""
+from __future__ import annotations
+
+import itertools
+import tempfile
+from collections import defaultdict
+from pathlib import Path
+
+from . import analyze, fingerprint
+
+SCHEMA = "kba.eval/1"
+
+
+def _pair(a: str, b: str) -> tuple[str, str]:
+    return (a, b) if a < b else (b, a)
+
+
+def _group_pairs(values) -> set[tuple[str, str]]:
+    items = [str(v) for v in (values or []) if str(v)]
+    return {_pair(a, b) for a, b in itertools.combinations(sorted(set(items)), 2)}
+
+
+def _pr(tp: int, fp: int, fn: int) -> dict:
+    precision = tp / (tp + fp) if tp + fp else None
+    recall = tp / (tp + fn) if tp + fn else None
+    if precision is None or recall is None:
+        f1 = None
+    elif precision + recall == 0:
+        f1 = 0.0
+    else:
+        f1 = 2 * precision * recall / (precision + recall)
+    return {
+        "tp": tp, "fp": fp, "fn": fn,
+        "precision": round(precision, 4) if precision is not None else None,
+        "recall": round(recall, 4) if recall is not None else None,
+        "f1": round(f1, 4) if f1 is not None else None,
+    }
+
+
+def _score_sets(pred: set, truth: set) -> dict:
+    return _pr(len(pred & truth), len(pred - truth), len(truth - pred))
+
+
+class _Resolver:
+    def __init__(self, records: list[dict], warnings: list[str]):
+        self.by_rel = {r.get("rel"): r for r in records if r.get("rel")}
+        self.by_p = {r.get("p"): r for r in records if r.get("p")}
+        self.by_name: dict[str, list[dict]] = defaultdict(list)
+        for r in records:
+            self.by_name[r.get("name", "")].append(r)
+        self.warnings = warnings
+
+    def get(self, value, warn: bool = True) -> dict | None:
+        if not isinstance(value, str) or not value:
+            return None
+        if value in self.by_rel:
+            return self.by_rel[value]
+        if value in self.by_p:
+            return self.by_p[value]
+        matches = self.by_name.get(value) or []
+        if len(matches) == 1:
+            return matches[0]
+        if warn and len(matches) > 1:
+            self.warnings.append(f"label 文件名不唯一，需用相对路径：{value}")
+        elif warn:
+            self.warnings.append(f"label 未匹配到文件：{value}")
+        return None
+
+
+def _record_pairs(records: list[dict], groups) -> tuple[set[tuple[str, str]], list[str]]:
+    warnings: list[str] = []
+    resolver = _Resolver(records, warnings)
+    pairs: set[tuple[str, str]] = set()
+    for group in groups or []:
+        if not isinstance(group, (list, tuple)) or len(group) < 2:
+            continue
+        keys = [resolver.get(v) for v in group]
+        if any(k is None for k in keys):
+            continue
+        pairs |= _group_pairs([k["p"] for k in keys])
+    return pairs, warnings
+
+
+def _page_pairs(overlaps: list[dict]) -> set[tuple[str, str, int, int]]:
+    result = set()
+    for o in overlaps:
+        a, b = o.get("a", ""), o.get("b", "")
+        pa, pb = int(o.get("a_page", 0) or 0), int(o.get("b_page", 0) or 0)
+        if not a or not b:
+            continue
+        if a > b:
+            a, b, pa, pb = b, a, pb, pa
+        result.add((a, b, pa, pb))
+    return result
+
+
+def _in_scope(p: str, scope: set[str]) -> bool:
+    return not scope or p in scope
+
+
+def _eval_page(records: list[dict], analysis: dict, labels: dict,
+               warnings: list[str], scope: set[str]) -> dict:
+    resolver = _Resolver(records, warnings)
+    truth = set()
+    for item in labels.get("page_pairs") or []:
+        if not isinstance(item, dict):
+            continue
+        a = resolver.get(item.get("a"))
+        b = resolver.get(item.get("b"))
+        if not a or not b or not _in_scope(a["p"], scope) or not _in_scope(b["p"], scope):
+            continue
+        pa = int(item.get("a_page", 1) or 1)
+        pb = int(item.get("b_page", 1) or 1)
+        if a["p"] > b["p"]:
+            a, b, pa, pb = b, a, pb, pa
+        truth.add((a["p"], b["p"], pa, pb))
+    pred = {p for p in _page_pairs(analysis.get("page_overlaps") or [])
+            if _in_scope(p[0], scope) and _in_scope(p[1], scope)}
+    return _score_sets(pred, truth)
+
+
+def _eval_version(records: list[dict], analysis: dict, labels: dict,
+                  warnings: list[str], scope: set[str]) -> dict:
+    resolver = _Resolver(records, warnings)
+    truth = set()
+    for item in labels.get("version_pairs") or []:
+        if isinstance(item, dict):
+            f, d = resolver.get(item.get("final")), resolver.get(item.get("draft"))
+        elif isinstance(item, (list, tuple)) and len(item) >= 2:
+            f, d = resolver.get(item[0]), resolver.get(item[1])
+        else:
+            continue
+        if f and d and _in_scope(f["p"], scope) and _in_scope(d["p"], scope):
+            truth.add((f["p"], d["p"]))
+
+    pred = set()
+    keep_final = 0
+    for c in analysis.get("clusters") or []:
+        finals = [m["p"] for m in c.get("members", [])
+                  if m.get("version_stage") == "final" and _in_scope(m["p"], scope)]
+        drafts = [m["p"] for m in c.get("members", [])
+                  if m.get("version_stage") == "draft" and _in_scope(m["p"], scope)]
+        for f in finals:
+            for d in drafts:
+                pred.add((f, d))
+                if c.get("recommend", {}).get("keep") == f:
+                    keep_final += 1
+    score = _score_sets(pred, truth)
+    score["keep_final_rate"] = round(keep_final / len(pred), 4) if pred else None
+    return score
+
+
+def _eval_pii(records: list[dict], analysis: dict, labels: dict,
+              warnings: list[str], scope: set[str]) -> dict:
+    resolver = _Resolver(records, warnings)
+    truth = set()
+    for value in labels.get("pii_files") or []:
+        r = resolver.get(value)
+        if r and _in_scope(r["p"], scope):
+            truth.add(r["p"])
+    pii = analysis.get("pii") or {}
+    if pii.get("files", 0) > len(pii.get("examples") or []):
+        warnings.append("PII 文件数超过分析结果保留的示例上限（50），召回可能偏低")
+    pred = set()
+    for ex in pii.get("examples") or []:
+        r = resolver.get(ex.get("p") or ex.get("name", ""), warn=False)
+        if r and _in_scope(r["p"], scope):
+            pred.add(r["p"])
+    return _score_sets(pred, truth)
+
+
+def _eval_tags(records: list[dict], labels: dict,
+               warnings: list[str], scope: set[str]) -> dict:
+    resolver = _Resolver(records, warnings)
+    label_tags = labels.get("tags") or {}
+    if not isinstance(label_tags, dict):
+        return {"files": 0, "exact": 0, "exact_rate": None,
+                "jaccard": None, "precision": None, "recall": None}
+    exact = 0
+    jaccards = []
+    tp = fp = fn = 0
+    for value, expected in label_tags.items():
+        r = resolver.get(value)
+        if not r or not _in_scope(r["p"], scope):
+            continue
+        expected_set = {str(t).strip() for t in (expected or []) if str(t).strip()}
+        predicted = {str(t).strip() for t in (r.get("tags") or []) if str(t).strip()}
+        if predicted == expected_set:
+            exact += 1
+        union = predicted | expected_set
+        jaccards.append(len(predicted & expected_set) / len(union) if union else 1.0)
+        tp += len(predicted & expected_set)
+        fp += len(predicted - expected_set)
+        fn += len(expected_set - predicted)
+    precision = tp / (tp + fp) if tp + fp else None
+    recall = tp / (tp + fn) if tp + fn else None
+    return {
+        "files": len(jaccards),
+        "exact": exact,
+        "exact_rate": round(exact / len(jaccards), 4) if jaccards else None,
+        "jaccard": round(sum(jaccards) / len(jaccards), 4) if jaccards else None,
+        "precision": round(precision, 4) if precision is not None else None,
+        "recall": round(recall, 4) if recall is not None else None,
+    }
+
+
+def _eval_extraction(records: list[dict], expected: dict | None) -> dict:
+    by_ext: dict[str, dict] = {}
+    for r in records:
+        ext = r.get("ext", "")
+        if not ext:
+            continue
+        bucket = by_ext.setdefault(ext, {
+            "files": 0, "extracted": 0, "no_text": 0, "failed": 0,
+            "not_enabled": 0, "unsupported": 0, "other": 0,
+        })
+        bucket["files"] += 1
+        status = r.get("text_status") or "other"
+        bucket[status if status in bucket else "other"] += 1
+    attempted = sum(b["extracted"] + b["no_text"] + b["failed"] for b in by_ext.values())
+    extracted = sum(b["extracted"] for b in by_ext.values())
+    expected = expected or {}
+    match = all(int(expected.get(ext, -1)) == by_ext.get(ext, {}).get("files", 0)
+                for ext in expected)
+    return {
+        "by_ext": by_ext,
+        "expected": expected,
+        "match": match if expected else None,
+        "attempted": attempted,
+        "success_rate": round(extracted / attempted, 4) if attempted else None,
+    }
+
+
+def eval_corpus(root: Path, labels: dict | None = None,
+                extract_dir: Path | None = None) -> dict:
+    root = Path(root).resolve()
+    if not root.is_dir():
+        raise NotADirectoryError(str(root))
+    labels = labels or {}
+    warnings: list[str] = []
+    with tempfile.TemporaryDirectory(prefix="kba-eval-job-") as td:
+        job = Path(td)
+        fp = fingerprint.scan(root, job, extract_dir=extract_dir)
+        result = analyze.analyze(job)
+        records = fp["files"]
+
+        scope: set[str] = set()
+        if labels.get("scope_files"):
+            scope_resolver = _Resolver(records, warnings)
+            for value in labels.get("scope_files") or []:
+                r = scope_resolver.get(value)
+                if r:
+                    scope.add(r["p"])
+
+        near_truth, near_warn = _record_pairs(records, labels.get("near_pairs") or [])
+        warnings += near_warn
+        near_truth = {p for p in near_truth if _in_scope(p[0], scope) and _in_scope(p[1], scope)}
+        pred_near = set()
+        for c in result.get("clusters") or []:
+            if c.get("type") in ("exact", "near"):
+                members = [m["p"] for m in c.get("members", []) if _in_scope(m["p"], scope)]
+                pred_near |= {_pair(a, b) for a, b in itertools.combinations(members, 2)}
+
+        same_truth, same_warn = _record_pairs(records, labels.get("same_name_groups") or [])
+        warnings += same_warn
+        same_truth = {p for p in same_truth if _in_scope(p[0], scope) and _in_scope(p[1], scope)}
+        pred_same = set()
+        for c in result.get("clusters") or []:
+            if c.get("type") == "same-name":
+                members = [m["p"] for m in c.get("members", []) if _in_scope(m["p"], scope)]
+                pred_same |= {_pair(a, b) for a, b in itertools.combinations(members, 2)}
+
+        metrics = {
+            "near": _score_sets(pred_near, near_truth),
+            "same_name": _score_sets(pred_same, same_truth),
+            "page": _eval_page(records, result, labels, warnings, scope),
+            "version": _eval_version(records, result, labels, warnings, scope),
+            "pii": _eval_pii(records, result, labels, warnings, scope),
+            "tags": _eval_tags(records, labels, warnings, scope),
+        }
+
+        return {
+            "schema": SCHEMA,
+            "root": {"name": root.name, "files": len(records)},
+            "labels": {
+                "near_pairs": len(near_truth),
+                "same_name_pairs": len(same_truth),
+                "page_pairs": len(labels.get("page_pairs") or []),
+                "version_pairs": len(labels.get("version_pairs") or []),
+                "pii_files": len(labels.get("pii_files") or []),
+                "tag_files": len(labels.get("tags") or {}),
+                "scope_files": len(scope) if labels.get("scope_files") else None,
+            },
+            "metrics": metrics,
+            "extraction": _eval_extraction(records, labels.get("extract_expected")),
+            "warnings": warnings,
+        }
+
+
+def format_text(result: dict) -> str:
+    lines = [f"KBA eval: {result['root']['files']} files / schema {result['schema']}"]
+    names = {
+        "near": "近重复", "same_name": "同名异内容", "page": "页级重合",
+        "version": "终稿/底稿", "pii": "PII 文件",
+    }
+    for key, label in names.items():
+        m = result["metrics"][key]
+        lines.append(
+            f"{label:<10} P={m['precision']} R={m['recall']} F1={m['f1']} "
+            f"(tp={m['tp']} fp={m['fp']} fn={m['fn']})"
+        )
+    tags = result["metrics"]["tags"]
+    lines.append(
+        f"{'标签':<10} exact={tags['exact_rate']} jaccard={tags['jaccard']} "
+        f"P={tags['precision']} R={tags['recall']} (files={tags['files']})"
+    )
+    ex = result["extraction"]
+    lines.append(
+        f"抽取        success={ex['success_rate']} expected_match={ex['match']} "
+        f"attempted={ex['attempted']}"
+    )
+    for ext, b in sorted(ex["by_ext"].items()):
+        lines.append(f"  {ext}: files={b['files']} extracted={b['extracted']} "
+                     f"no_text={b['no_text']} failed={b['failed']} "
+                     f"not_enabled={b['not_enabled']} unsupported={b['unsupported']}")
+    for w in result.get("warnings") or []:
+        lines.append(f"WARN {w}")
+    return "\n".join(lines)