detection_tokens/scoring.py

from __future__ import annotations

import json
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np
import torch

from .data import BinaryEvalRow, SourcePair
from .modeling import ModelBundle, build_prompt, compute_average_logprob, compute_token_logprobs, encode_response


@dataclass(slots=True)
class ScoredText:
    text: str
    label: int
    pair_id: str
    score: float
    ai_avg_logp: float
    human_avg_logp: float


def compute_auc(labels: list[int], scores: list[float]) -> float:
    order = np.argsort(scores)
    ranks = np.empty(len(scores), dtype=np.float64)
    ranks[order] = np.arange(1, len(scores) + 1)
    pos_count = sum(labels)
    neg_count = len(labels) - pos_count
    if pos_count == 0 or neg_count == 0:
        return 0.5
    pos_rank_sum = float(sum(rank for rank, label in zip(ranks, labels) if label == 1))
    return (pos_rank_sum - pos_count * (pos_count + 1) / 2.0) / (pos_count * neg_count)


def dual_score(bundle: ModelBundle, text: str) -> tuple[float, float, float]:
    ai_prompt = build_prompt(bundle, bundle.config.model.ai_token)
    human_prompt = build_prompt(bundle, bundle.config.model.human_token)
    ai_ids, ai_prompt_len = encode_response(bundle, ai_prompt, text)
    human_ids, human_prompt_len = encode_response(bundle, human_prompt, text)
    ai_logp = float(compute_average_logprob(bundle, ai_ids, ai_prompt_len).item())
    human_logp = float(compute_average_logprob(bundle, human_ids, human_prompt_len).item())
    score_mode = bundle.config.scoring.score_mode
    if score_mode == "avg_margin":
        score = ai_logp - human_logp
    elif score_mode == "soft_token_sigmoid":
        ai_token_logps = compute_token_logprobs(bundle, ai_ids, ai_prompt_len)
        human_token_logps = compute_token_logprobs(bundle, human_ids, human_prompt_len)
        if ai_token_logps.shape[0] != human_token_logps.shape[0]:
            token_count = min(ai_token_logps.shape[0], human_token_logps.shape[0])
            ai_token_logps = ai_token_logps[:token_count]
            human_token_logps = human_token_logps[:token_count]
        margins = ai_token_logps - human_token_logps
        tau = max(bundle.config.scoring.token_sigmoid_tau, 1.0e-6)
        score = float(torch.sigmoid(margins / tau).mean().item())
    else:
        raise ValueError(f"Unsupported scoring mode: {score_mode}")
    return score, ai_logp, human_logp


def decision_threshold(bundle: ModelBundle) -> float:
    if bundle.config.scoring.score_mode == "soft_token_sigmoid":
        return 0.5
    return 0.0


def evaluate_holdout(bundle: ModelBundle, holdout_pairs: list[SourcePair], output_dir: Path) -> dict:
    threshold = decision_threshold(bundle)
    rows: list[ScoredText] = []
    for pair in holdout_pairs:
        for label, text in ((1, pair.ai_text), (0, pair.human_text)):
            score, ai_logp, human_logp = dual_score(bundle, text)
            rows.append(
                ScoredText(
                    text=text,
                    label=label,
                    pair_id=pair.pair_id,
                    score=score,
                    ai_avg_logp=ai_logp,
                    human_avg_logp=human_logp,
                )
            )

    labels = [row.label for row in rows]
    scores = [row.score for row in rows]
    tp = sum(int(row.label == 1 and row.score > threshold) for row in rows)
    fp = sum(int(row.label == 0 and row.score > threshold) for row in rows)
    tn = sum(int(row.label == 0 and row.score <= threshold) for row in rows)
    fn = sum(int(row.label == 1 and row.score <= threshold) for row in rows)
    precision = tp / max(1, tp + fp)
    recall = tp / max(1, tp + fn)
    f1 = 0.0 if precision + recall == 0.0 else 2.0 * precision * recall / (precision + recall)

    pairwise_wins = 0
    by_pair: dict[str, dict[int, ScoredText]] = {}
    for row in rows:
        by_pair.setdefault(row.pair_id, {})[row.label] = row
    for item in by_pair.values():
        pairwise_wins += int(item[1].score > item[0].score)

    summary = {
        "num_holdout_texts": len(rows),
        "score_mode": bundle.config.scoring.score_mode,
        "decision_threshold": threshold,
        "auroc": compute_auc(labels, scores),
        "accuracy": (tp + tn) / len(rows),
        "pairwise_rate": pairwise_wins / len(holdout_pairs),
        "mean_score_ai": float(np.mean([row.score for row in rows if row.label == 1])),
        "mean_score_human": float(np.mean([row.score for row in rows if row.label == 0])),
        "tp": tp,
        "fp": fp,
        "tn": tn,
        "fn": fn,
        "precision": precision,
        "recall": recall,
        "f1": f1,
    }
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "dual_eval_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    (output_dir / "dual_eval_details.json").write_text(
        json.dumps([asdict(row) for row in rows], indent=2),
        encoding="utf-8",
    )
    return summary


def evaluate_binary_rows(
    bundle: ModelBundle,
    rows: list[BinaryEvalRow],
    output_path: Path | None = None,
) -> dict:
    threshold = decision_threshold(bundle)
    scores: list[float] = []
    labels: list[int] = []
    details = []
    for row in rows:
        score, ai_logp, human_logp = dual_score(bundle, row.text)
        scores.append(score)
        labels.append(row.label)
        details.append(
            {
                "row_id": row.row_id,
                "label": row.label,
                "text_type": row.text_type,
                "model": row.model,
                "source_id": row.source_id,
                "score": score,
                "ai_avg_logp": ai_logp,
                "human_avg_logp": human_logp,
            }
        )

    labels_array = np.asarray(labels, dtype=np.int64)
    scores_array = np.asarray(scores, dtype=np.float64)
    predictions = (scores_array > threshold).astype(np.int64)

    tp = int(((labels_array == 1) & (predictions == 1)).sum())
    fp = int(((labels_array == 0) & (predictions == 1)).sum())
    tn = int(((labels_array == 0) & (predictions == 0)).sum())
    fn = int(((labels_array == 1) & (predictions == 0)).sum())
    precision = tp / max(1, tp + fp)
    recall = tp / max(1, tp + fn)
    f1 = 0.0 if precision + recall == 0.0 else 2.0 * precision * recall / (precision + recall)

    summary = {
        "num_rows": len(rows),
        "score_mode": bundle.config.scoring.score_mode,
        "decision_threshold": threshold,
        "positive_rows": int(labels_array.sum()),
        "negative_rows": int((labels_array == 0).sum()),
        "auroc": compute_auc(labels, scores),
        "accuracy_at_zero": float((predictions == labels_array).mean()),
        "tp": tp,
        "fp": fp,
        "tn": tn,
        "fn": fn,
        "precision": precision,
        "recall": recall,
        "f1": f1,
        "mean_score_positive": float(scores_array[labels_array == 1].mean()),
        "mean_score_negative": float(scores_array[labels_array == 0].mean()),
    }
    if output_path is not None:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(
            json.dumps({"summary": summary, "details": details}, indent=2),
            encoding="utf-8",
        )
    return summary