code/src/jev_judge/data.py

"""Data pipeline (DESIGN §3, §5, §6.3): parquet splits -> template tokens -> kind-stratified batches.

Target layout: every sample's teacher distribution `q` is placed directly in *slot space*
(float32 [24]) so losses/metrics are uniform across kinds: noul -> slots 0-1, score -> 2-7,
choice -> 8 .. 8+n-1. Inactive slots are 0 and masked.
"""

from __future__ import annotations

import json
import math
import os
from dataclasses import dataclass
from typing import Iterator, Sequence

import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset, Sampler

from .template import (
    KIND_TO_ID,
    KINDS,
    NUM_SLOTS,
    SLOT_RANGES,
    TemplateError,
    n_options_for,
    render,
    render_parts,
    validate_options,
)

SPLITS = ("train", "validation", "calibration", "test", "test_set_30k", "ood")
UNIFORM_TOL = 1e-9


# --------------------------------------------------------------------------------------------
# Rows / placeholders
# --------------------------------------------------------------------------------------------
def is_uniform(target: Sequence[float], tol: float = UNIFORM_TOL) -> bool:
    t = np.asarray(target, dtype=np.float64)
    return bool(np.all(np.abs(t - 1.0 / len(t)) <= tol))


def d1_flags(df: pd.DataFrame, scope: str = "yuri_v1") -> np.ndarray:
    """Rows subject to the D1 policy: uniform placeholder labels (DESIGN V7).

    scope="yuri_v1": only the yuri_v1 stream (100% exact [0.5,0.5]);
    scope="all": every exactly-uniform row regardless of stream."""
    uni = df["is_uniform"].to_numpy(dtype=bool)
    if scope == "all":
        return uni
    return uni & (df["source"].to_numpy() == scope)


def apply_d1_policy(df: pd.DataFrame, policy: str, downweight: float, scope: str = "yuri_v1") -> tuple[pd.DataFrame, np.ndarray]:
    """Return (df, sample_weights). policy in {keep, downweight, drop}."""
    flags = d1_flags(df, scope)
    w = np.ones(len(df), dtype=np.float32)
    if policy == "keep":
        return df, w
    if policy == "downweight":
        w[flags] = downweight
        return df, w
    if policy == "drop":
        keep = ~flags
        return df.loc[keep].reset_index(drop=True), w[keep]
    raise ValueError(f"unknown d1 policy {policy!r}")


# --------------------------------------------------------------------------------------------
# Encoding
# --------------------------------------------------------------------------------------------
@dataclass
class Encoded:
    input_ids: list[int]
    truncated: bool


def encode_sample(tokenizer, kind: str, state: str, question: str, options: Sequence[str], max_seq_len: int,
                  allow_truncate: bool = True) -> Encoded:
    """Tokenize the rendered template as one string (natural BPE seams).

    Only when the sequence exceeds `max_seq_len` is `state` truncated -- by tokens, head 60% / tail
    40% -- decoded back to text and the template re-rendered (DESIGN §5: question/options/suffix
    are never truncated). The corpus fits in 1024 with zero truncation (V6), so this path is for
    serving edge cases."""
    ids = tokenizer.encode(render(kind, state, question, options), add_special_tokens=False)
    if len(ids) <= max_seq_len:
        return Encoded(ids, False)
    if not allow_truncate:
        raise ValueError(f"sequence of {len(ids)} tokens exceeds max_seq_len={max_seq_len}")
    prefix, _, suffix = render_parts(kind, state, question, options)
    fixed = len(tokenizer.encode(prefix + suffix, add_special_tokens=False))
    s_ids = tokenizer.encode(state, add_special_tokens=False)
    budget = max_seq_len - fixed - 4  # slack for seam tokens
    if budget < 8:
        raise ValueError("question/options alone exceed max_seq_len; refuse to truncate them")
    for _ in range(4):
        head = int(budget * 0.6)
        tail = budget - head
        st = tokenizer.decode(s_ids[:head]) + " ... " + tokenizer.decode(s_ids[-tail:])
        ids = tokenizer.encode(render(kind, st, question, options), add_special_tokens=False)
        if len(ids) <= max_seq_len:
            return Encoded(ids, True)
        budget -= (len(ids) - max_seq_len) + 4
        if budget < 8:
            break
    return Encoded(ids[:max_seq_len], True)


def target_to_slots(kind: str, target: Sequence[float]) -> np.ndarray:
    q = np.zeros(NUM_SLOTS, dtype=np.float32)
    s, _ = SLOT_RANGES[kind]
    t = np.asarray(target, dtype=np.float32)
    q[s : s + len(t)] = t / max(float(t.sum()), 1e-12)
    return q


# --------------------------------------------------------------------------------------------
# Parquet IO
# --------------------------------------------------------------------------------------------
def load_split(data_dir: str, split: str) -> pd.DataFrame:
    path = os.path.join(data_dir, f"{split}.parquet")
    df = pd.read_parquet(path)
    return df


def subset(df: pd.DataFrame, frac: float, seed: int) -> pd.DataFrame:
    """Kind×source-stratified subset (for 10% scans)."""
    if frac >= 1.0:
        return df
    rng = np.random.default_rng(seed)
    parts = []
    for _, g in df.groupby(["kind", "source"], sort=False):
        n = max(1, int(round(len(g) * frac)))
        parts.append(g.iloc[rng.permutation(len(g))[:n]])
    return pd.concat(parts).sample(frac=1.0, random_state=seed).reset_index(drop=True)


# --------------------------------------------------------------------------------------------
# Dataset
# --------------------------------------------------------------------------------------------
class JevDataset(Dataset):
    def __init__(self, df: pd.DataFrame, tokenizer, max_seq_len: int, weights: np.ndarray | None = None,
                 choice_permute_prob: float = 0.0, seed: int = 0):
        self.df = df.reset_index(drop=True)
        self.tok = tokenizer
        self.max_seq_len = max_seq_len
        self.weights = np.ones(len(self.df), dtype=np.float32) if weights is None else weights.astype(np.float32)
        self.choice_permute_prob = choice_permute_prob
        self.seed = seed
        self.epoch = 0
        self.kind_ids = self.df["kind"].map(KIND_TO_ID).to_numpy(dtype=np.int64)
        self.lengths = self.df["n_tokens"].to_numpy(dtype=np.int64) if "n_tokens" in self.df else None

    def set_epoch(self, epoch: int) -> None:
        """Augmentation randomness is a pure function of (seed, epoch, index) -> safe with DataLoader workers."""
        self.epoch = epoch

    def __len__(self) -> int:
        return len(self.df)

    def __getitem__(self, i: int) -> dict:
        row = self.df.iloc[i]
        kind = row["kind"]
        options = list(row["options"])
        target = np.asarray(row["target"], dtype=np.float32)
        if kind == "choice" and self.choice_permute_prob > 0:
            rng = np.random.default_rng([self.seed, self.epoch, int(i)])
            if rng.random() < self.choice_permute_prob:
                perm = rng.permutation(len(options))
                options = [options[j] for j in perm]
                target = target[perm]
        enc = encode_sample(self.tok, kind, row["state"], row["question"], options, self.max_seq_len)
        return {
            "input_ids": enc.input_ids,
            "kind_id": int(KIND_TO_ID[kind]),
            "n_options": int(n_options_for(kind, options)),
            "target": target_to_slots(kind, target),
            "weight": float(self.weights[i]),
            "index": i,
        }


def collate(batch: list[dict], pad_token_id: int) -> dict[str, torch.Tensor]:
    B = len(batch)
    lengths = torch.tensor([len(b["input_ids"]) for b in batch], dtype=torch.long)
    T = int(lengths.max())
    ids = torch.full((B, T), pad_token_id, dtype=torch.long)
    attn = torch.zeros((B, T), dtype=torch.long)
    for i, b in enumerate(batch):
        n = len(b["input_ids"])
        ids[i, :n] = torch.as_tensor(b["input_ids"], dtype=torch.long)
        attn[i, :n] = 1
    return {
        "input_ids": ids,
        "attention_mask": attn,
        "lengths": lengths,
        "kind_ids": torch.tensor([b["kind_id"] for b in batch], dtype=torch.long),
        "n_options": torch.tensor([b["n_options"] for b in batch], dtype=torch.long),
        "target": torch.from_numpy(np.stack([b["target"] for b in batch])),
        "weight": torch.tensor([b["weight"] for b in batch], dtype=torch.float32),
        "index": torch.tensor([b["index"] for b in batch], dtype=torch.long),
    }


# --------------------------------------------------------------------------------------------
# Kind-stratified, length-bucketed batch sampler (DESIGN §6.3)
# --------------------------------------------------------------------------------------------
class KindBatchSampler(Sampler[list[int]]):
    """Each batch draws from per-kind pools so that every primitive holds >= `kind_floor` of the batch
    (when its pool still has items); within a pool, indices are sorted by length inside windows of
    `bucket_batches` batches to reduce padding. Deterministic given `seed` + `set_epoch`."""

    def __init__(self, kind_ids: np.ndarray, lengths: np.ndarray | None, batch_size: int, kind_floor: float = 1 / 6,
                 bucket_batches: int = 64, seed: int = 0, drop_last: bool = False, mode: str = "kind_first"):
        """mode="kind_first": per-kind pools, exact kind floor, lengths only loosely aligned (pad eff. ~0.7).
        mode="length_first": global length windows -> length-homogeneous batches (pad eff. ~0.95); the kind
        floor is checked and violating batches are repaired by swapping with neighbours when possible."""
        self.mode = mode
        self.kind_ids = np.asarray(kind_ids)
        self.lengths = None if lengths is None else np.asarray(lengths)
        self.batch_size = batch_size
        self.kind_floor = kind_floor
        self.bucket_batches = bucket_batches
        self.seed = seed
        self.drop_last = drop_last
        self.epoch = 0
        self.n = len(self.kind_ids)
        self._num_batches = self.n // batch_size if drop_last else math.ceil(self.n / batch_size)

    def set_epoch(self, epoch: int) -> None:
        self.epoch = epoch

    def __len__(self) -> int:
        return self._num_batches

    def _pool(self, rng: np.random.Generator, kid: int) -> list[int]:
        idx = np.flatnonzero(self.kind_ids == kid)
        idx = idx[rng.permutation(len(idx))]
        if self.lengths is None:
            return idx.tolist()
        # sort by length inside windows
        win = max(1, int(self.batch_size * self.bucket_batches * (len(idx) / max(self.n, 1))))
        out = []
        for s in range(0, len(idx), win):
            chunk = idx[s : s + win]
            chunk = chunk[np.argsort(self.lengths[chunk], kind="stable")]
            out.extend(chunk.tolist())
        return out

    def _iter_length_first(self, rng: np.random.Generator) -> Iterator[list[int]]:
        perm = rng.permutation(self.n)
        win = self.batch_size * self.bucket_batches
        floor_n = int(math.floor(self.kind_floor * self.batch_size))
        batches: list[list[int]] = []
        for s in range(0, self.n, win):
            chunk = perm[s : s + win]
            if self.lengths is not None:
                chunk = chunk[np.argsort(self.lengths[chunk], kind="stable")]
            cur = [chunk[i : i + self.batch_size].tolist() for i in range(0, len(chunk), self.batch_size)]
            if self.drop_last:
                cur = [b for b in cur if len(b) == self.batch_size]
            # repair kind floor by swapping items between neighbouring (similar-length) batches
            for bi, b in enumerate(cur):
                counts = np.bincount(self.kind_ids[b], minlength=len(KINDS))
                for kid in range(len(KINDS)):
                    need = floor_n - counts[kid]
                    if need <= 0:
                        continue
                    for nb in (bi - 1, bi + 1):
                        if need <= 0 or not (0 <= nb < len(cur)):
                            continue
                        other = cur[nb]
                        surplus_kid = int(np.argmax(counts))
                        donors = [j for j, x in enumerate(other) if self.kind_ids[x] == kid]
                        takers = [j for j, x in enumerate(b) if self.kind_ids[x] == surplus_kid]
                        for j1, j2 in zip(donors, takers):
                            if need <= 0 or counts[surplus_kid] <= floor_n:
                                break
                            b[j2], other[j1] = other[j1], b[j2]
                            counts[kid] += 1; counts[surplus_kid] -= 1; need -= 1
            batches.extend(cur)
        for j in rng.permutation(len(batches)):
            yield batches[j]

    def __iter__(self) -> Iterator[list[int]]:
        rng = np.random.default_rng(self.seed * 1_000_003 + self.epoch)
        if self.mode == "length_first":
            yield from self._iter_length_first(rng)
            return
        pools = {kid: self._pool(rng, kid) for kid in range(len(KINDS))}
        pos = {kid: 0 for kid in pools}
        remaining = {kid: len(p) for kid, p in pools.items()}
        floor_n = int(math.floor(self.kind_floor * self.batch_size))
        batches: list[list[int]] = []
        while sum(remaining.values()) > 0:
            bs = min(self.batch_size, sum(remaining.values()))
            # floors first, then proportional fill from what remains
            take = {kid: min(floor_n, remaining[kid]) for kid in pools}
            left = bs - sum(take.values())
            if left > 0:
                tot = sum(remaining[kid] - take[kid] for kid in pools)
                for kid in pools:
                    if tot <= 0:
                        break
                    share = int(round(left * (remaining[kid] - take[kid]) / tot))
                    take[kid] += min(share, remaining[kid] - take[kid])
                # fix rounding
                gap = bs - sum(take.values())
                for kid in sorted(pools, key=lambda k: -(remaining[k] - take[k])):
                    if gap == 0:
                        break
                    room = remaining[kid] - take[kid]
                    if gap > 0 and room > 0:
                        d = min(gap, room); take[kid] += d; gap -= d
                    elif gap < 0 and take[kid] > 0:
                        d = min(-gap, take[kid]); take[kid] -= d; gap += d
            batch: list[int] = []
            for kid in pools:
                k = take[kid]
                if k:
                    batch.extend(pools[kid][pos[kid] : pos[kid] + k])
                    pos[kid] += k
                    remaining[kid] -= k
            if len(batch) < self.batch_size and self.drop_last:
                break
            batches.append(batch)
        order = rng.permutation(len(batches))
        for j in order:
            yield batches[j]


# --------------------------------------------------------------------------------------------
# Row-level validation used by prepare_data / server
# --------------------------------------------------------------------------------------------
def validate_row(row: dict) -> list[str]:
    errs = []
    kind = row.get("kind")
    if kind not in KINDS:
        errs.append(f"bad kind {kind!r}")
        return errs
    try:
        opts = validate_options(kind, row.get("options", []))
    except TemplateError as e:
        errs.append(str(e))
        return errs
    t = row.get("target")
    if not isinstance(t, (list, tuple, np.ndarray)) or len(t) != len(opts):
        errs.append(f"target length {None if t is None else len(t)} != options {len(opts)}")
        return errs
    t = np.asarray(t, dtype=np.float64)
    if np.any(t < -1e-9) or not np.isfinite(t).all():
        errs.append("target has negative/non-finite entries")
    if abs(t.sum() - 1.0) > 1e-3:
        errs.append(f"target sums to {t.sum():.5f}")
    for f in ("state", "question"):
        if not isinstance(row.get(f), str) or not row[f].strip():
            errs.append(f"empty {f}")
    return errs