local_eval.py

#!/usr/bin/env python3
"""Local evaluator for IOL-AI practice files.

Practice CSV columns:
  id, context, query, task_type, eval_type, gold

Submission CSV columns:
  id, pred

Both gold and pred are JSON list strings.
"""
from __future__ import annotations

import argparse
import csv
import json
import os
import re
import shutil
import subprocess
import sys
from collections import Counter
from pathlib import Path
from typing import Any


PRACTICE_COLUMNS = ("id", "context", "query", "task_type", "eval_type", "gold")
SUBMISSION_COLUMNS = ("id", "pred")
SPACE_RE = re.compile(r"\s+")


def score_text(value: Any) -> str:
    """Light scoring view: trim and collapse whitespace, preserving Unicode."""
    return SPACE_RE.sub(" ", str(value).strip())


def json_value_to_text(value: Any) -> str:
    if isinstance(value, str):
        return value
    if isinstance(value, (dict, list)):
        return json.dumps(value, ensure_ascii=False)
    if value is None:
        return ""
    return str(value)


def parse_json_list(value: str, column: str, row_id: str) -> list[str]:
    try:
        parsed = json.loads(value)
    except json.JSONDecodeError as exc:
        raise ValueError(f"id={row_id}: {column} is not valid JSON: {exc}") from exc

    if not isinstance(parsed, list):
        raise ValueError(f"id={row_id}: {column} must be a JSON list")
    return [json_value_to_text(item) for item in parsed]


def require_columns(path: Path, fieldnames: list[str] | None, required: tuple[str, ...]) -> None:
    missing = [column for column in required if column not in (fieldnames or [])]
    if missing:
        raise ValueError(f"{path} is missing required columns: {', '.join(missing)}")


def read_practice(path: Path) -> list[dict[str, Any]]:
    with path.open(newline="", encoding="utf-8") as handle:
        reader = csv.DictReader(handle)
        require_columns(path, reader.fieldnames, PRACTICE_COLUMNS)
        rows: list[dict[str, Any]] = []
        for row in reader:
            row_id = str(row["id"])
            rows.append({**row, "gold_list": parse_json_list(row["gold"], "gold", row_id)})
    return rows


def read_submission(path: Path) -> dict[str, list[str]]:
    with path.open(newline="", encoding="utf-8") as handle:
        reader = csv.DictReader(handle)
        require_columns(path, reader.fieldnames, SUBMISSION_COLUMNS)
        predictions: dict[str, list[str]] = {}
        for row in reader:
            row_id = str(row["id"])
            if row_id in predictions:
                raise ValueError(f"{path} contains duplicate id={row_id}")
            predictions[row_id] = parse_json_list(row["pred"], "pred", row_id)
    return predictions


def char_ngrams(text: str, n: int) -> Counter[str]:
    if n <= 0 or len(text) < n:
        return Counter()
    return Counter(text[i : i + n] for i in range(len(text) - n + 1))


def chrf_similarity(pred: str, gold: str, max_order: int = 6, beta: float = 2.0) -> float:
    """chrF-style character n-gram F-score in [0, 1]."""
    pred = score_text(pred)
    gold = score_text(gold)
    if not pred and not gold:
        return 1.0
    if not pred or not gold:
        return 0.0

    matched = 0
    pred_total = 0
    gold_total = 0
    for n in range(1, max_order + 1):
        pred_counts = char_ngrams(pred, n)
        gold_counts = char_ngrams(gold, n)
        matched += sum((pred_counts & gold_counts).values())
        pred_total += sum(pred_counts.values())
        gold_total += sum(gold_counts.values())

    if pred_total == 0 or gold_total == 0 or matched == 0:
        return 0.0

    precision = matched / pred_total
    recall = matched / gold_total
    beta_sq = beta * beta
    return (1 + beta_sq) * precision * recall / (beta_sq * precision + recall)


def align_to_gold(pred: list[str], gold_count: int) -> list[str]:
    if len(pred) > gold_count:
        return pred[:gold_count]
    if len(pred) < gold_count:
        return pred + [""] * (gold_count - len(pred))
    return pred


def score_predictions(
    practice_rows: list[dict[str, Any]], predictions: dict[str, list[str]]
) -> tuple[dict[str, float], list[dict[str, Any]]]:
    details: list[dict[str, Any]] = []
    row_exact_sum = 0
    length_ok_sum = 0
    answer_exact_sum = 0
    answer_count = 0
    chrf_sum = 0.0

    for row in practice_rows:
        row_id = str(row["id"])
        gold = row["gold_list"]
        raw_pred = predictions.get(row_id, [])
        pred = align_to_gold(raw_pred, len(gold))
        length_ok = len(raw_pred) == len(gold)

        exact_flags = [score_text(p) == score_text(g) for p, g in zip(pred, gold)]
        row_exact = length_ok and all(exact_flags)
        item_chrfs = [chrf_similarity(p, g) for p, g in zip(pred, gold)]
        row_chrf = sum(item_chrfs) / len(item_chrfs) if item_chrfs else (1.0 if not pred else 0.0)

        row_exact_sum += int(row_exact)
        length_ok_sum += int(length_ok)
        answer_exact_sum += sum(int(flag) for flag in exact_flags)
        answer_count += len(gold)
        chrf_sum += sum(item_chrfs)

        details.append(
            {
                "id": row_id,
                "gold_count": len(gold),
                "pred_count": len(raw_pred),
                "length_ok": int(length_ok),
                "row_exact": int(row_exact),
                "answer_exact": sum(int(flag) for flag in exact_flags) / len(gold) if gold else float(row_exact),
                "chrf": row_chrf,
                "gold": json.dumps(gold, ensure_ascii=False),
                "pred": json.dumps(pred, ensure_ascii=False),
            }
        )

    row_count = len(practice_rows)
    metrics = {
        "rows": float(row_count),
        "answers": float(answer_count),
        "missing_predictions": float(sum(1 for row in practice_rows if str(row["id"]) not in predictions)),
        "extra_predictions": float(len(set(predictions) - {str(row["id"]) for row in practice_rows})),
        "length_ok": length_ok_sum / row_count if row_count else 0.0,
        "row_exact": row_exact_sum / row_count if row_count else 0.0,
        "answer_exact": answer_exact_sum / answer_count if answer_count else 0.0,
        "chrf": chrf_sum / answer_count if answer_count else 0.0,
    }
    return metrics, details


def write_details(path: Path, details: list[dict[str, Any]]) -> None:
    fieldnames = ("id", "gold_count", "pred_count", "length_ok", "row_exact", "answer_exact", "chrf", "gold", "pred")
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(details)


def run_script(script_path: Path, practice_path: Path) -> None:
    env = os.environ.copy()
    input_path = Path("/tmp/data/test.csv")
    input_path.parent.mkdir(parents=True, exist_ok=True)
    shutil.copyfile(practice_path, input_path)
    subprocess.run([sys.executable, str(script_path)], check=True, env=env)


def print_summary(metrics: dict[str, float]) -> None:
    print(f"Rows: {int(metrics['rows'])}")
    print(f"Answers: {int(metrics['answers'])}")
    print(f"Missing predictions: {int(metrics['missing_predictions'])}")
    print(f"Extra predictions: {int(metrics['extra_predictions'])}")
    print(f"Length OK: {metrics['length_ok']:.4f}")
    print(f"Row exact: {metrics['row_exact']:.4f}")
    print(f"Answer exact: {metrics['answer_exact']:.4f}")
    print(f"chrF-style: {metrics['chrf']:.4f}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Score local IOL-AI practice predictions.")
    parser.add_argument("practice_csv", type=Path, help="Practice CSV with id, context, query, task_type, eval_type, gold.")
    parser.add_argument(
        "--submission",
        type=Path,
        default=Path("submission.csv"),
        help="Submission CSV with id,pred JSON-list strings. Default: submission.csv",
    )
    parser.add_argument(
        "--run-script",
        action="store_true",
        help="Run script.py with IOL_INPUT set to the practice CSV before scoring.",
    )
    parser.add_argument("--script", type=Path, default=Path("script.py"), help="Script to run with --run-script.")
    parser.add_argument("--details", type=Path, help="Optional path for per-row score details CSV.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.run_script:
        run_script(args.script, args.practice_csv)

    practice_rows = read_practice(args.practice_csv)
    predictions = read_submission(args.submission)
    metrics, details = score_predictions(practice_rows, predictions)
    print_summary(metrics)

    if args.details:
        write_details(args.details, details)
        print(f"Wrote details: {args.details}")


if __name__ == "__main__":
    main()