train/data.py

"""3-tier streaming data loader over memmap shards (plan §1, §6 curriculum).

Samples records from gold/silver/bronze with configurable tier weights and a cleanliness
floor, yielding fixed-T char windows ready for the collator. Two phases:
  stable  weights = {gold, silver, bronze} (bronze included for coverage)
  anneal  weights = {gold: 1, silver: 0, bronze: 0}  (gold-only, washes out synth/repair bias)

Long records are chopped into T-char windows on word boundaries; short ones are used whole
(the collator packs several per row). Deterministic given (seed, rank, world_size) so runs
resume exactly and DP ranks see disjoint data.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq

# Metadata-conditioning UNK ids (insc/data/meta_vocab.py's UNK_REGION/UNK_CENTURY) -- same
# hardcoded-constant convention as train/collate.py's UNK_REGION/UNK_CENTURY, for the same
# reason (this is base-pretraining-shared code; the insc-only meta_vocab package isn't
# always on sys.path here). Shards without region_id/century_id columns (every GCB
# pretraining/gold/silver/bronze shard, built before this existed) fall back to these.
UNK_REGION, UNK_CENTURY = 14, 15


@dataclass
class TierSpec:
    path: str
    weight: float
    tier_filter: str = None      # keep only records whose `tier` column == this (None = all)


@dataclass
class DataConfig:
    tiers: dict = field(default_factory=dict)     # name -> TierSpec
    min_clean: float = 0.0
    drop_dup_frac: float = 0.7
    window_chars: int = 4096
    seed: int = 0
    exclude_holdout: bool = True   # False = train on all records (external dev set drives eval)


class ShardSet:
    """One shard directory's memmap planes + record index (memmaps shared across tiers)."""
    _cache = {}

    def __init__(self, path):
        d = Path(path)
        ip = d / "index_dedup.parquet"
        idx = pq.read_table(ip if ip.exists() else d / "index.parquet")
        self.offset = idx.column("offset").to_numpy()
        self.length = idx.column("length").to_numpy()
        self.clean = idx.column("clean").to_numpy()
        self.tier = idx.column("tier").to_numpy(zero_copy_only=False)
        self.dup = (idx.column("dup_frac").to_numpy() if "dup_frac" in idx.column_names
                    else np.zeros(len(self.offset), np.float32))
        n = len(self.offset)
        self.region_id = (idx.column("region_id").to_numpy() if "region_id" in idx.column_names
                          else np.full(n, UNK_REGION, np.int64))
        self.century_id = (idx.column("century_id").to_numpy() if "century_id" in idx.column_names
                           else np.full(n, UNK_CENTURY, np.int64))
        self.chars = np.memmap(d / "chars.bin", dtype=np.uint8, mode="r")
        self.boundary = np.memmap(d / "boundary.bin", dtype=np.uint8, mode="r")
        self.dia = np.memmap(d / "dia.bin", dtype=np.uint8, mode="r")
        self.cap = np.memmap(d / "cap.bin", dtype=np.uint8, mode="r")
        punct_p = d / "punct.bin"
        # tolerate shards built before the punctuation plane existed: falls back to all-"none"
        self.punct = np.memmap(punct_p, dtype=np.uint8, mode="r") if punct_p.exists() else None

    @classmethod
    def get(cls, path):
        if path not in cls._cache:
            cls._cache[path] = cls(path)
        return cls._cache[path]

    def eligible(self, min_clean, drop_dup, tier_filter=None, exclude_holdout=True):
        m = (self.clean >= min_clean) & (self.dup < drop_dup)
        if tier_filter is not None:
            m = m & (self.tier == tier_filter)
        idx = np.flatnonzero(m)
        if exclude_holdout:                       # reserve every HOLDOUT_MOD-th record for eval
            idx = idx[idx % HOLDOUT_MOD != 0]
        return idx


HOLDOUT_MOD = 200   # ~0.5% held out from training; eval/intrinsic selects idx % HOLDOUT_MOD == 0


class MultiTierLoader:
    def __init__(self, cfg: DataConfig, rank=0, world_size=1):
        self.cfg = cfg
        self.rank, self.world = rank, world_size
        self.sets, self.elig, self.names, self.wts = {}, {}, [], []
        for name, spec in cfg.tiers.items():
            if spec.weight <= 0:
                continue
            ss = ShardSet.get(spec.path)
            el = ss.eligible(cfg.min_clean, cfg.drop_dup_frac, spec.tier_filter,
                             exclude_holdout=cfg.exclude_holdout)
            # shard eligible records across DP ranks
            el = el[rank::world_size]
            if len(el) == 0:
                continue
            self.sets[name] = ss
            self.elig[name] = el
            self.names.append(name)
            self.wts.append(spec.weight)
        assert self.names, "no eligible tiers"
        self.wts = np.array(self.wts, float)
        self.wts /= self.wts.sum()
        self.rng = np.random.default_rng(cfg.seed + 1315423911 * rank)
        self._cursor = {n: 0 for n in self.names}
        self._perm = {n: self.rng.permutation(self.elig[n]) for n in self.names}

    def _next_record(self, name):
        ss = self.sets[name]
        p = self._perm[name]
        c = self._cursor[name]
        if c >= len(p):
            self._perm[name] = self.rng.permutation(self.elig[name])
            c = 0
        i = int(self._perm[name][c])
        self._cursor[name] = c + 1
        o, l = int(ss.offset[i]), int(ss.length[i])
        punct = (np.asarray(ss.punct[o:o+l]) if ss.punct is not None
                else np.zeros(l, dtype=np.uint8))
        return dict(chars=np.asarray(ss.chars[o:o+l]), boundary=np.asarray(ss.boundary[o:o+l]),
                    dia=np.asarray(ss.dia[o:o+l]), cap=np.asarray(ss.cap[o:o+l]), punct=punct,
                    region_id=int(ss.region_id[i]), century_id=int(ss.century_id[i]))

    def _window(self, rec):
        """Chop a long record to <= window_chars on a word boundary; else return whole."""
        W = self.cfg.window_chars
        n = len(rec["chars"])
        if n <= W:
            return rec
        b = rec["boundary"]
        start = int(self.rng.integers(0, n - W))
        # snap start to just after a boundary, end to a boundary
        we = np.flatnonzero(b[:start] >= 1)
        s = (we[-1] + 1) if len(we) else 0
        seg_end = np.flatnonzero(b[s:s+W] >= 1)
        e = (s + seg_end[-1] + 1) if len(seg_end) else min(s + W, n)
        # per-character planes get windowed; scalar per-record metadata (region_id/
        # century_id) passes through unchanged -- it describes the whole inscription/
        # papyrus, not any one character span within it.
        return {k: (v[s:e] if isinstance(v, np.ndarray) else v) for k, v in rec.items()}

    def records(self, n):
        """Yield n windowed records sampled by tier weight."""
        picks = self.rng.choice(len(self.names), size=n, p=self.wts)
        for pi in picks:
            yield self._window(self._next_record(self.names[pi]))


def stable_cfg(gdata, w=(1.0, 1.0, 0.3), window=4096, seed=0, exclude_holdout=True):
    """gold:silver:bronze default weights (bronze down-weighted as synthetic)."""
    return DataConfig(tiers={
        "gold":   TierSpec(f"{gdata}/shards/v1_punct", w[0], tier_filter="pristine"),
        "silver": TierSpec(f"{gdata}/shards/v1_punct", w[1], tier_filter="repaired"),
        "bronze": TierSpec(f"{gdata}/shards/bronze_punct", w[2], tier_filter="bronze"),
    }, window_chars=window, seed=seed, exclude_holdout=exclude_holdout)


def anneal_cfg(gdata, window=4096, seed=0, exclude_holdout=True):
    """gold-only anneal phase."""
    return DataConfig(tiers={
        "gold": TierSpec(f"{gdata}/shards/v1_punct", 1.0, tier_filter="pristine"),
    }, window_chars=window, seed=seed, exclude_holdout=exclude_holdout)