fithia2/apps/qa/evaluate_gold_set.py

"""QA Tool: evaluate parser output against a gold set of manually labeled documents."""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

from libs.common.config import get_settings
from libs.common.logging import configure_logging, get_logger
from libs.parser.rule_parser import RuleBasedParser
from libs.parser.text_normalizer import normalize_text

logger = get_logger(__name__)

_GOLD_SET_DIR = Path("data/gold_set")

_DIRECTION_FIELDS = ("event_type", "event_direction")
_ONEOFF_FIELD = "oneoff_item"


def _load_gold_set(gold_dir: Path) -> list[dict[str, Any]]:
    """Load all gold set JSON files from directory."""
    entries = []
    for path in sorted(gold_dir.glob("*.json")):
        data = json.loads(path.read_text())
        if isinstance(data, list):
            entries.extend(data)
        else:
            entries.append(data)
    return entries


def _precision_recall_f1(tp: int, fp: int, fn: int) -> dict[str, float]:
    precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
    f1 = (
        2 * precision * recall / (precision + recall)
        if (precision + recall) > 0
        else 0.0
    )
    return {"precision": precision, "recall": recall, "f1": f1}


def evaluate(gold_entries: list[dict[str, Any]], parser: RuleBasedParser) -> dict[str, Any]:
    """Run parser on gold set and compute accuracy metrics.

    Each gold entry must have:
    - document_id: str
    - form_type: str
    - text: str (raw exhibit text)
    - gold_event_type: str
    - gold_event_direction: str
    - gold_guidance_status: str (optional)
    - gold_oneoff: bool (optional)
    - gold_quality_assessment: str (optional)
    """
    total = len(gold_entries)
    if total == 0:
        return {"error": "Empty gold set"}

    # Accuracy trackers
    event_type_correct = 0
    direction_correct = 0
    guidance_correct = 0
    guidance_total = 0

    # Oneoff precision/recall
    oneoff_tp = oneoff_fp = oneoff_fn = 0

    # Evidence presence
    evidence_present = 0

    per_doc: list[dict[str, Any]] = []

    for entry in gold_entries:
        text = normalize_text(entry.get("text", ""))
        meta = {
            "filing_date": entry.get("filing_date", "1900-01-01"),
            "form_type": entry.get("form_type", "8-K"),
        }
        output = parser.parse(
            document_id=entry.get("document_id", "unknown"),
            form_type=entry.get("form_type", "8-K"),
            text=text,
            metadata=meta,
        )
        out = output.model_dump()

        gold_et = entry.get("gold_event_type", "")
        gold_dir = entry.get("gold_event_direction", "")
        gold_guidance = entry.get("gold_guidance_status", None)
        gold_oneoff = entry.get("gold_oneoff", None)

        et_match = out["event_type"] == gold_et
        dir_match = out["event_direction"] == gold_dir
        event_type_correct += int(et_match)
        direction_correct += int(dir_match)

        if gold_guidance is not None:
            guidance_correct += int(out["guidance"]["status"] == gold_guidance)
            guidance_total += 1

        if gold_oneoff is not None:
            pred_oneoff = out["risk_flags"]["oneoff_item"]
            if pred_oneoff and gold_oneoff:
                oneoff_tp += 1
            elif pred_oneoff and not gold_oneoff:
                oneoff_fp += 1
            elif not pred_oneoff and gold_oneoff:
                oneoff_fn += 1

        has_evidence = len(out.get("evidence", [])) > 0
        evidence_present += int(has_evidence)

        per_doc.append({
            "document_id": entry.get("document_id"),
            "et_match": et_match,
            "dir_match": dir_match,
            "pred_event_type": out["event_type"],
            "gold_event_type": gold_et,
            "pred_direction": out["event_direction"],
            "gold_direction": gold_dir,
            "confidence": out["confidence"]["overall"],
        })

    oneoff_metrics = _precision_recall_f1(oneoff_tp, oneoff_fp, oneoff_fn)
    summary: dict[str, Any] = {
        "total": total,
        "event_type_accuracy": event_type_correct / total,
        "direction_accuracy": direction_correct / total,
        "guidance_accuracy": guidance_correct / guidance_total if guidance_total > 0 else None,
        "oneoff_precision": oneoff_metrics["precision"],
        "oneoff_recall": oneoff_metrics["recall"],
        "oneoff_f1": oneoff_metrics["f1"],
        "evidence_presence_ratio": evidence_present / total,
        "per_document": per_doc,
    }
    return summary


def main() -> None:
    arg_parser = argparse.ArgumentParser(description="Evaluate parser against gold set")
    arg_parser.add_argument(
        "--gold-dir",
        default=str(_GOLD_SET_DIR),
        help="Directory containing gold set JSON files",
    )
    arg_parser.add_argument("--output", default=None, help="Write JSON report to this file")
    arg_parser.add_argument("--verbose", action="store_true", help="Print per-doc results")
    args = arg_parser.parse_args()

    settings = get_settings()
    configure_logging(settings.log_level)

    gold_dir = Path(args.gold_dir)
    if not gold_dir.exists():
        print(f"Gold set directory not found: {gold_dir}", file=sys.stderr)
        sys.exit(1)

    gold_entries = _load_gold_set(gold_dir)
    if not gold_entries:
        print(f"No gold set entries found in {gold_dir}", file=sys.stderr)
        sys.exit(1)

    parser = RuleBasedParser()
    report = evaluate(gold_entries, parser)

    print(f"\n=== Gold Set Evaluation Report ({report['total']} documents) ===")
    print(f"  Event type accuracy    : {report['event_type_accuracy']:.1%}")
    print(f"  Direction accuracy     : {report['direction_accuracy']:.1%}")
    if report["guidance_accuracy"] is not None:
        print(f"  Guidance accuracy      : {report['guidance_accuracy']:.1%}")
    print(f"  One-off F1             : {report['oneoff_f1']:.3f}")
    print(f"    Precision            : {report['oneoff_precision']:.3f}")
    print(f"    Recall               : {report['oneoff_recall']:.3f}")
    print(f"  Evidence presence      : {report['evidence_presence_ratio']:.1%}")

    if args.verbose and report.get("per_document"):
        print("\n--- Per-document results ---")
        for doc in report["per_document"]:
            status = "OK" if doc["et_match"] and doc["dir_match"] else "MISMATCH"
            print(
                f"  [{status}] {doc['document_id']}: "
                f"type={doc['pred_event_type']}({doc['gold_event_type']}) "
                f"dir={doc['pred_direction']}({doc['gold_direction']}) "
                f"conf={doc['confidence']:.2f}"
            )

    if args.output:
        out_path = Path(args.output)
        out_path.write_text(json.dumps(report, indent=2))
        print(f"\nReport written to {out_path}")


if __name__ == "__main__":
    main()