src/deberta_ime/evaluation_cli.py

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from collections.abc import Sequence
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, TextIO, cast

from .ajimee import (
    AJIMEE_LICENSE,
    AJIMEE_REPOSITORY,
    AJIMEE_REVISION,
    AJIMEE_SHA256,
)
from .evaluation import (
    EvaluationItem,
    Prediction,
    PredictionProvenance,
    evaluate_predictions,
)


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="deberta-ime-evaluate")
    parser.add_argument("--items", type=Path, required=True)
    parser.add_argument("--predictions", type=Path, required=True)
    parser.add_argument("--format", choices=("generic", "ajimee"), default="generic")
    parser.add_argument("--dataset-name", default="unspecified")
    parser.add_argument("--dataset-revision", default="unspecified")
    parser.add_argument("--dataset-license", default="unspecified")
    parser.add_argument("--candidate-limit", type=int, default=8)
    parser.add_argument("--output-dir", type=Path, default=Path("outputs"))
    parser.add_argument("--stem", default="finite_candidate_evaluation")
    return parser


def run(
    argv: Sequence[str] | None = None,
    *,
    stdout: TextIO | None = None,
) -> int:
    args = _parser().parse_args(argv)
    items = _load_items(args.items, item_format=args.format)
    predictions = _load_predictions(args.predictions)
    evaluation = evaluate_predictions(
        items,
        predictions,
        candidate_limit=args.candidate_limit,
    )
    items_artifact = _artifact(args.items)
    predictions_artifact = _artifact(args.predictions)
    pinned_ajimee = args.format == "ajimee" and items_artifact["sha256"] == AJIMEE_SHA256
    report: dict[str, Any] = {
        "schema_version": 2,
        "generated_at": datetime.now(UTC).isoformat(),
        "status": "LOCAL_FINITE_CANDIDATE_EVALUATION",
        "format": args.format,
        "dataset": (
            {
                "name": AJIMEE_REPOSITORY,
                "revision": AJIMEE_REVISION,
                "license": AJIMEE_LICENSE,
            }
            if pinned_ajimee
            else {
                "name": "AJIMEE-compatible input",
                "revision": "unverified-by-sha256",
                "license": "unspecified",
            }
            if args.format == "ajimee"
            else {
                "name": args.dataset_name,
                "revision": args.dataset_revision,
                "license": args.dataset_license,
            }
        ),
        "artifacts": {
            "items": items_artifact,
            "predictions": predictions_artifact,
        },
        "evaluation": evaluation,
        "claim_boundaries": [
            "An empty or missing prediction abstains and preserves the input text.",
            "Overcorrection is reported only for rows explicitly labelled clean.",
            "Candidates are scored as supplied; this command performs no model inference.",
            "AJIMEE provenance is pinned only when the item artifact SHA-256 matches.",
            "Prediction provenance, reason, and margin are declarations from the "
            "supplied artifact.",
        ],
    }
    args.output_dir.mkdir(parents=True, exist_ok=True)
    json_path = args.output_dir / f"{args.stem}.json"
    markdown_path = args.output_dir / f"{args.stem}.md"
    json_path.write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    markdown_path.write_text(_render_markdown(report), encoding="utf-8")
    (stdout or sys.stdout).write(
        json.dumps(
            {"ok": True, "json": str(json_path), "markdown": str(markdown_path)},
            ensure_ascii=False,
            indent=2,
        )
        + "\n"
    )
    return 0


def _load_items(path: Path, *, item_format: str) -> tuple[EvaluationItem, ...]:
    payload = _load_json_array(path)
    items: list[EvaluationItem] = []
    for row_number, raw_item in enumerate(payload, start=1):
        if not isinstance(raw_item, dict):
            raise ValueError(f"item row {row_number}: expected an object")
        if item_format == "ajimee":
            item_id = raw_item.get("index")
            input_text = raw_item.get("input")
            references = raw_item.get("expected_output")
            label: object = "unspecified"
        else:
            item_id = raw_item.get("id")
            input_text = raw_item.get("input")
            references = raw_item.get("references")
            label = raw_item.get("label", "unspecified")
        if not isinstance(item_id, str) or not isinstance(input_text, str):
            raise ValueError(f"item row {row_number}: id and input must be strings")
        if (
            not isinstance(references, list)
            or not references
            or not all(isinstance(reference, str) for reference in references)
        ):
            raise ValueError(
                f"item row {row_number}: references must be a non-empty string list"
            )
        if label not in ("clean", "typo", "unspecified"):
            raise ValueError(f"item row {row_number}: invalid label")
        items.append(
            EvaluationItem(
                item_id=item_id,
                input_text=input_text,
                references=tuple(cast(list[str], references)),
                label=label,
            )
        )
    return tuple(items)


def _load_predictions(path: Path) -> tuple[Prediction, ...]:
    payload = _load_json_array(path)
    predictions: list[Prediction] = []
    for row_number, raw_prediction in enumerate(payload, start=1):
        if not isinstance(raw_prediction, dict):
            raise ValueError(f"prediction row {row_number}: expected an object")
        item_id = raw_prediction.get("id", raw_prediction.get("index"))
        candidates = raw_prediction.get("candidates")
        if not isinstance(item_id, str):
            raise ValueError(f"prediction row {row_number}: id must be a string")
        if not isinstance(candidates, list) or not all(
            isinstance(candidate, str) and bool(candidate) for candidate in candidates
        ):
            raise ValueError(f"prediction row {row_number}: candidates must be a string list")
        provenance = raw_prediction.get("provenance", "unspecified")
        if provenance not in ("provider", "deberta", "lfm", "rule", "unspecified"):
            raise ValueError(f"prediction row {row_number}: invalid provenance")
        reason = raw_prediction.get("reason")
        if reason is not None and not isinstance(reason, str):
            raise ValueError(f"prediction row {row_number}: reason must be a string")
        margin = raw_prediction.get("margin")
        if margin is not None and (
            isinstance(margin, bool) or not isinstance(margin, (int, float))
        ):
            raise ValueError(f"prediction row {row_number}: margin must be a number")
        predictions.append(
            Prediction(
                item_id=item_id,
                candidates=tuple(cast(list[str], candidates)),
                provenance=cast(PredictionProvenance, provenance),
                reason=reason,
                margin=None if margin is None else float(margin),
            )
        )
    return tuple(predictions)


def _load_json_array(path: Path) -> list[object]:
    payload: object = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(payload, list):
        raise ValueError(f"{path}: expected a JSON array")
    return cast(list[object], payload)


def _artifact(path: Path) -> dict[str, str | int]:
    data = path.read_bytes()
    return {
        "path": str(path),
        "sha256": hashlib.sha256(data).hexdigest(),
        "size_bytes": len(data),
    }


def _render_markdown(report: dict[str, Any]) -> str:
    dataset = report["dataset"]
    evaluation = report["evaluation"]
    metrics = evaluation["metrics"]
    overcorrection = metrics["overcorrection_rate"]
    rendered_overcorrection = (
        "not available (no explicit clean rows)"
        if overcorrection is None
        else f"{overcorrection:.2%}"
    )
    reported_margin = metrics["mean_reported_margin"]
    rendered_margin = (
        "not reported"
        if reported_margin is None
        else f"{reported_margin:.4f} ({metrics['reported_margin_rows']} rows)"
    )
    provenance = (
        ", ".join(
            f"{name}={count}" for name, count in metrics["declared_provenance_counts"].items()
        )
        or "not reported"
    )
    return "\n".join(
        [
            "# Finite-candidate evaluation",
            "",
            f"Generated: {report['generated_at']}",
            f"Dataset: {dataset['name']} @ {dataset['revision']}",
            f"Dataset license/boundary: {dataset['license']}",
            f"Rows: {evaluation['rows']}",
            f"Candidate limit: {evaluation['candidate_limit']}",
            "",
            f"- Effective Acc@1: {metrics['effective_acc_at_1']:.2%}",
            f"- Candidate Recall@k: {metrics['candidate_recall_at_k']:.2%}",
            f"- Mean MinCER: {metrics['mean_min_cer']:.2%}",
            f"- Abstention: {metrics['abstention_rate']:.2%}",
            f"- Overcorrection: {rendered_overcorrection}",
            f"- Accepted candidate misses: {metrics['accepted_candidate_miss_rows']}",
            f"- Selection errors: {metrics['selection_error_rows']}",
            f"- Declared provenance: {provenance}",
            f"- Mean declared margin: {rendered_margin}",
            "",
            "An abstention preserves the input. No model inference was performed. "
            "Provenance/reason/margin are artifact declarations, not independently "
            "verified facts.",
            "",
        ]
    )


def main() -> None:
    raise SystemExit(run())


if __name__ == "__main__":
    main()