detection_tokens/data.py

from __future__ import annotations

import json
import random
from collections import defaultdict
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

from datasets import Dataset, load_dataset

from .config import DataConfig


@dataclass(slots=True)
class SourcePair:
    pair_id: str
    text_id: str
    source_id: str
    dataset_name: str
    source: str
    model: str
    text_type: str
    cosine_score: float | None
    ai_text: str
    human_text: str


@dataclass(slots=True)
class BinaryEvalRow:
    row_id: str
    text: str
    label: int
    text_type: str
    model: str
    source_id: str


def _valid_text(text: str, *, min_text_chars: int) -> bool:
    return isinstance(text, str) and len(text.strip()) >= min_text_chars


def _normalize_text(text: Any) -> str:
    return str(text or "").strip()


def _find_local_arrow_file(root: Path, split: str) -> Path:
    direct_path = root / f"editlens_iclr-{split}.arrow"
    if direct_path.exists():
        return direct_path
    matches = sorted(root.rglob(f"editlens_iclr-{split}.arrow"))
    if not matches:
        raise FileNotFoundError(f"Missing local dataset file for split={split!r} under {root}")
    return matches[0]


def _pick_best_row(rows: list[dict[str, Any]], *, text_key: str) -> dict[str, Any]:
    if not rows:
        raise ValueError("Cannot pick from an empty row list.")
    rows = sorted(
        rows,
        key=lambda row: (
            _normalize_text(row.get("prompt")) == "",
            _normalize_text(row.get("title")) == "",
            _normalize_text(row.get(text_key)) == "",
        ),
    )
    return rows[0]


def load_pangram_rows(config: DataConfig, *, split: str) -> Dataset:
    if config.pangram.local_dataset_path is not None:
        arrow_path = _find_local_arrow_file(config.pangram.local_dataset_path, split)
        return Dataset.from_file(str(arrow_path))
    return load_dataset(config.pangram.dataset_name, split=split)


def load_raid_rows(config: DataConfig, *, split: str) -> Dataset:
    return load_dataset(config.raid.dataset_name, split=split)


def build_pangram_binary_pairs(config: DataConfig) -> list[SourcePair]:
    rows = load_pangram_rows(config, split=config.pangram.dataset_split)
    ai_rows = []
    human_by_text_id: dict[str, list[dict[str, Any]]] = defaultdict(list)
    human_by_source_id: dict[str, list[dict[str, Any]]] = defaultdict(list)

    for row in rows:
        text = _normalize_text(row.get("text"))
        if not _valid_text(text, min_text_chars=config.min_text_chars):
            continue
        text_type = _normalize_text(row.get("text_type"))
        if text_type in config.pangram.human_text_types:
            text_id = _normalize_text(row.get("text_id"))
            source_id = _normalize_text(row.get("source_id"))
            if text_id:
                human_by_text_id[text_id].append(dict(row))
            if source_id:
                human_by_source_id[source_id].append(dict(row))
        elif text_type in config.pangram.ai_text_types:
            ai_rows.append(dict(row))

    pairs: list[SourcePair] = []
    for row in ai_rows:
        pair_source_id = _normalize_text(row.get("source_id"))
        if not pair_source_id:
            continue
        candidates = human_by_text_id.get(pair_source_id)
        if not candidates:
            candidates = human_by_source_id.get(pair_source_id)
        if not candidates:
            continue
        human_row = _pick_best_row(candidates, text_key="text")
        pairs.append(
            SourcePair(
                pair_id=f"pangram::{pair_source_id}::{_normalize_text(row.get('text_id'))}",
                text_id=_normalize_text(row.get("text_id")),
                source_id=pair_source_id,
                dataset_name="pangram",
                source=_normalize_text(row.get("source")),
                model=_normalize_text(row.get("model")),
                text_type=_normalize_text(row.get("text_type")),
                cosine_score=float(row["cosine_score"]) if row.get("cosine_score") is not None else None,
                ai_text=_normalize_text(row.get("text")),
                human_text=_normalize_text(human_row.get("text")),
            )
        )
    return pairs


def build_raid_binary_pairs(config: DataConfig) -> list[SourcePair]:
    rows = load_raid_rows(config, split=config.raid.dataset_split)
    human_rows: list[dict[str, Any]] = []
    ai_rows: list[dict[str, Any]] = []
    for row in rows:
        if config.raid.require_attack_none and _normalize_text(row.get("attack")) not in {"", "none"}:
            continue
        text = _normalize_text(row.get("generation"))
        if not _valid_text(text, min_text_chars=config.min_text_chars):
            continue
        model = _normalize_text(row.get("model"))
        if model == config.raid.human_model_name:
            human_rows.append(dict(row))
        else:
            ai_rows.append(dict(row))

    human_by_id: dict[str, list[dict[str, Any]]] = defaultdict(list)
    human_by_source_id: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in human_rows:
        row_id = _normalize_text(row.get("id"))
        source_id = _normalize_text(row.get("source_id"))
        if row_id:
            human_by_id[row_id].append(row)
        if source_id:
            human_by_source_id[source_id].append(row)

    pairs: list[SourcePair] = []
    for row in ai_rows:
        pair_source_id = _normalize_text(row.get("source_id"))
        if not pair_source_id:
            continue
        candidates = human_by_id.get(pair_source_id)
        if not candidates:
            candidates = human_by_source_id.get(pair_source_id)
        if not candidates:
            continue
        human_row = _pick_best_row(candidates, text_key="generation")
        pairs.append(
            SourcePair(
                pair_id=f"raid::{pair_source_id}::{_normalize_text(row.get('model'))}::{_normalize_text(row.get('id'))}",
                text_id=_normalize_text(row.get("id")),
                source_id=pair_source_id,
                dataset_name="raid",
                source=_normalize_text(row.get("domain")),
                model=_normalize_text(row.get("model")),
                text_type="ai_generated",
                cosine_score=None,
                ai_text=_normalize_text(row.get("generation")),
                human_text=_normalize_text(human_row.get("generation")),
            )
        )
    return pairs


def _take_pairs(
    pairs: list[SourcePair],
    *,
    take: int,
    seed: int,
) -> tuple[list[SourcePair], list[SourcePair]]:
    rng = random.Random(seed)
    shuffled = list(pairs)
    rng.shuffle(shuffled)
    if len(shuffled) < take:
        raise ValueError(f"Need at least {take} pairs, found {len(shuffled)}.")
    return shuffled[:take], shuffled[take:]


def build_training_and_eval_splits(
    config: DataConfig,
    *,
    seed: int,
) -> tuple[list[SourcePair], list[SourcePair], list[SourcePair], dict[str, int]]:
    source_pools: dict[str, list[SourcePair]] = {}
    if config.pangram.enabled:
        source_pools["pangram"] = build_pangram_binary_pairs(config)
    if config.raid.enabled:
        source_pools["raid"] = build_raid_binary_pairs(config)

    train_pairs: list[SourcePair] = []
    holdout_candidates: list[SourcePair] = []
    raid_eval_pairs: list[SourcePair] = []
    metadata = {f"{name}_pairs_available": len(pairs) for name, pairs in source_pools.items()}

    if config.raid.enabled:
        raid_eval_pairs, remaining_raid = _take_pairs(
            source_pools["raid"],
            take=config.raid.eval_holdout_pairs,
            seed=seed + 100,
        )
        source_pools["raid"] = remaining_raid

    if config.pangram.enabled:
        selected, remaining = _take_pairs(
            source_pools["pangram"],
            take=config.pangram.train_pairs,
            seed=seed + 1,
        )
        train_pairs.extend(selected)
        holdout_candidates.extend(remaining)

    if config.raid.enabled:
        selected, remaining = _take_pairs(
            source_pools["raid"],
            take=config.raid.train_pairs,
            seed=seed + 2,
        )
        train_pairs.extend(selected)
        holdout_candidates.extend(remaining)

    holdout_pairs, _ = _take_pairs(
        holdout_candidates,
        take=config.training_holdout_pairs,
        seed=seed + 3,
    )
    random.Random(seed + 4).shuffle(train_pairs)
    random.Random(seed + 5).shuffle(holdout_pairs)
    metadata.update(
        {
            "train_pairs_from_pangram": sum(pair.dataset_name == "pangram" for pair in train_pairs),
            "train_pairs_from_raid": sum(pair.dataset_name == "raid" for pair in train_pairs),
            "holdout_pairs_from_pangram": sum(pair.dataset_name == "pangram" for pair in holdout_pairs),
            "holdout_pairs_from_raid": sum(pair.dataset_name == "raid" for pair in holdout_pairs),
            "raid_eval_pairs_from_raid": len(raid_eval_pairs),
        }
    )
    return train_pairs, holdout_pairs, raid_eval_pairs, metadata


def save_pairs(path: Path, pairs: list[SourcePair]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps([asdict(pair) for pair in pairs], indent=2), encoding="utf-8")


def load_pairs(path: Path) -> list[SourcePair]:
    rows = json.loads(path.read_text(encoding="utf-8"))
    return [SourcePair(**row) for row in rows]


def load_binary_eval_rows(
    config: DataConfig,
    *,
    split: str,
    positive_text_types: set[str],
    negative_text_types: set[str],
) -> list[BinaryEvalRow]:
    rows = load_pangram_rows(config, split=split)
    payload: list[BinaryEvalRow] = []
    for index, row in enumerate(rows):
        text_type = str(row.get("text_type", "")).strip()
        if text_type in positive_text_types:
            label = 1
        elif text_type in negative_text_types:
            label = 0
        else:
            continue
        text = str(row.get("text", "")).strip()
        if not _valid_text(text, min_text_chars=config.min_text_chars):
            continue
        payload.append(
            BinaryEvalRow(
                row_id=str(row.get("text_id", index)),
                text=text,
                label=label,
                text_type=text_type,
                model=str(row.get("model", "")),
                source_id=str(row.get("source_id", "")),
            )
        )
    return payload