ml/data.py

"""Dataset and sequence encoding for block infilling.

The masked target is a contiguous, line-aligned BLOCK of the program body (~half
of it), not a single short line. This stresses global structure (def-use across
the gap, table/segment coherence) — the regime where the thesis expects the
diffusion substrate to matter. Both models see the SAME block:

  diffusion (in-place, exact): [bos] prefix <block> suffix [eos] [pad...]
    the block positions are the denoise region; context is always visible. No gap
    padding — the block occupies exactly its real length.

  AR-FIM (causal): [bos][pre] prefix [suf] suffix [mid] block [eos] [pad...]
    loss only on the block+eos span, matching diffusion's region-only loss.

`prefix + block + suffix == source` exactly (char-offset slicing), so a correct
fill reconstructs the program and verification by execution is well-defined.
"""

from __future__ import annotations

import json

import numpy as np
import torch
from torch.utils.data import Dataset

from .config import TaskConfig
from .tokenizer import Tokenizer


def load_records(path: str) -> list[dict]:
    with open(path) as f:
        return [json.loads(line) for line in f if line.strip()]


def make_block(source: str, frac: float, rng: np.random.RandomState):
    """Slice the source into (prefix, block, suffix) where block is a contiguous,
    line-aligned run of body lines totalling ~frac of the body. Returns None if
    there is no usable body. prefix+block+suffix == source exactly."""
    nl = source.find("\n")
    if nl < 0:
        return None
    body_start = nl + 1
    stripped = source.rstrip("\n")
    end_pos = stripped.rfind("\nend")
    if end_pos < 0:
        return None
    body_end = end_pos + 1  # position of the final 'end'
    body = source[body_start:body_end]
    if not body.strip():
        return None

    # Line-start offsets within the body (each body line keeps its trailing \n).
    starts = [0] + [i + 1 for i, ch in enumerate(body) if ch == "\n" and i + 1 < len(body)]
    n = len(starts)
    target = max(1, int(len(body) * frac))

    s = rng.randint(0, n)
    acc = 0
    e = s
    while e < n and acc < target:
        seg_end = starts[e + 1] if e + 1 < n else len(body)
        acc += seg_end - starts[e]
        e += 1
    block_start = body_start + starts[s]
    block_end = body_start + (starts[e] if e < n else len(body))

    prefix = source[:block_start]
    block = source[block_start:block_end]
    suffix = source[block_end:]
    if not block.strip():
        return None
    return prefix, block, suffix


def encode_diffusion(tok: Tokenizer, prefix, block, suffix, cfg: TaskConfig):
    """(ids, region_mask, block_id, attn_mask) of length T, or None if it doesn't
    fit. block_id is the generation-block index for region positions, -1 elsewhere
    (used by the block-causal mask)."""
    pre = tok.encode(prefix)
    mid = tok.encode(block)
    suf = tok.encode(suffix)
    body = [tok.bos_id] + pre + mid + suf + [tok.eos_id]
    if len(body) > cfg.seq_len:
        return None
    ids = np.full(cfg.seq_len, tok.pad_id, dtype=np.int64)
    ids[: len(body)] = body
    region = np.zeros(cfg.seq_len, dtype=bool)
    start = 1 + len(pre)
    region[start : start + len(mid)] = True
    block_id = np.full(cfg.seq_len, -1, dtype=np.int64)
    for r in range(len(mid)):
        block_id[start + r] = r // cfg.block_len
    attn = ids != tok.pad_id
    return ids, region, block_id, attn


def encode_ar(tok: Tokenizer, prefix, block, suffix, cfg: TaskConfig):
    """(ids, loss_mask) of length T, or None if it doesn't fit."""
    pre = tok.encode(prefix)
    mid = tok.encode(block)
    suf = tok.encode(suffix)
    head = [tok.bos_id, tok.pre_id] + pre + [tok.suf_id] + suf + [tok.mid_id]
    tail = mid + [tok.eos_id]
    body = head + tail
    if len(body) > cfg.seq_len:
        return None
    ids = np.full(cfg.seq_len, tok.pad_id, dtype=np.int64)
    ids[: len(body)] = body
    loss_mask = np.zeros(cfg.seq_len, dtype=bool)
    loss_mask[len(head) : len(head) + len(tail)] = True
    return ids, loss_mask


def block_for_eval(source: str, frac: float, idx: int):
    """Deterministic block for a record (seeded by index), so diffusion and AR are
    evaluated on the IDENTICAL masked block."""
    return make_block(source, frac, np.random.RandomState(idx + 1))


# ---- token-level (lua mode): everything in token-id space ----
# Whitespace is dropped, so blocks are contiguous lexeme runs and reconstruction
# is tok.decode(prefix + block + suffix) (space-joined, executes identically).

def make_block_lua(source: str, frac: float, rng: np.random.RandomState, tok: Tokenizer):
    """Split the lexeme-id list into (pre, block, suf) id lists, block ~frac of the
    body (tokens between the signature's ')' and the final 'end'). Returns None if
    no usable body."""
    ids = tok.encode(source)
    rp = tok.stoi.get(")")
    endt = tok.stoi.get("end")
    if rp is None or endt is None or rp not in ids or endt not in ids:
        return None
    h = ids.index(rp)  # end of the parameter list
    be = max(i for i, t in enumerate(ids) if t == endt)  # final 'end'
    b0, b1 = h + 1, be
    if b1 <= b0:
        return None
    n = b1 - b0
    target = max(1, int(n * frac))
    s = rng.randint(0, n)
    e = min(s + target, n)
    a, c = b0 + s, b0 + e
    return ids[:a], ids[a:c], ids[c:]


def _ids_canvas(tok, pre, blk, suf, cfg, ar):
    if ar:
        head = [tok.bos_id, tok.pre_id] + pre + [tok.suf_id] + suf + [tok.mid_id]
        tail = blk + [tok.eos_id]
        body = head + tail
        if len(body) > cfg.seq_len:
            return None
        arr = np.full(cfg.seq_len, tok.pad_id, dtype=np.int64)
        arr[: len(body)] = body
        loss = np.zeros(cfg.seq_len, dtype=bool)
        loss[len(head): len(head) + len(tail)] = True
        return arr, loss
    body = [tok.bos_id] + pre + blk + suf + [tok.eos_id]
    if len(body) > cfg.seq_len:
        return None
    arr = np.full(cfg.seq_len, tok.pad_id, dtype=np.int64)
    arr[: len(body)] = body
    region = np.zeros(cfg.seq_len, dtype=bool)
    start = 1 + len(pre)
    region[start: start + len(blk)] = True
    bid = np.full(cfg.seq_len, -1, dtype=np.int64)
    for r in range(len(blk)):
        bid[start + r] = r // cfg.block_len
    attn = arr != tok.pad_id
    return arr, region, bid, attn


class InfillDataset(Dataset):
    """Block-infilling dataset. Each record gets a deterministic block (seeded by
    index) so runs are reproducible and the two modes train on matching blocks."""

    def __init__(self, records: list[dict], tok: Tokenizer, cfg: TaskConfig, mode: str):
        assert mode in ("diffusion", "ar")
        self.tok = tok
        self.cfg = cfg
        self.mode = mode
        self.items = []
        self.skipped = 0
        ar = mode == "ar"
        lua = tok.mode == "lua"
        for i, rec in enumerate(records):
            # Per-record block fraction sampled in [frac_lo, frac_hi] so one model
            # handles any masking ratio; eval then sweeps the ratio for free.
            rng = np.random.RandomState(i + 1)
            frac = cfg.frac_lo + (cfg.frac_hi - cfg.frac_lo) * rng.rand()
            if lua:
                blk = make_block_lua(rec["source"], frac, rng, tok)
                e = None if blk is None else _ids_canvas(tok, blk[0], blk[1], blk[2], cfg, ar)
            else:
                blk = make_block(rec["source"], frac, rng)
                e = None if blk is None else (
                    encode_ar(tok, *blk, cfg) if ar else encode_diffusion(tok, *blk, cfg))
            if blk is None or e is None:
                self.skipped += 1
            else:
                self.items.append(e)

    def __len__(self):
        return len(self.items)

    def __getitem__(self, idx):
        if self.mode == "diffusion":
            ids, region, block_id, attn = self.items[idx]
            return (torch.from_numpy(ids), torch.from_numpy(region),
                    torch.from_numpy(block_id), torch.from_numpy(attn))
        ids, loss_mask = self.items[idx]
        return torch.from_numpy(ids), torch.from_numpy(loss_mask)