code/train/data.py

# coding=utf-8
# Copyright 2024 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Data loading procedure for Othello and Sudoku game.
"""

import itertools
import os
import pickle

import jax
import numpy as np
import tensorflow as tf
from tensorflow.compat.v1 import gfile

import pdb


class CurriculumState:
    """Mutable curriculum stage shared between the train loop and the sampler.

    stage t (1..max_stage) is a REASONING-DEPTH stage: t latent slots are
    active and the first t wave snapshots are supervised. Every puzzle is
    available from step 0 -- difficulty is not gated.

    The difficulty tag is deliberately unused as a curriculum axis. It takes
    only 6 values (3..8 = hardest solver-strategy digit needed by any cell), so
    it cannot express a 12-step ladder, it is uncorrelated with puzzle size
    (r=-0.003 vs empty-cell count), and it explains only ~19% of the variance
    in solver round count. The ladder is over propagation depth instead.
    """

    def __init__(self, stage=1, max_stage=12):
        self.stage = stage
        self.max_stage = max_stage

    def unlocked_levels(self):
        """All levels, always. Kept for logging/compat with per-level reports."""
        return list(range(3, 9))


def compute_puzzle_levels(strategy_codes, start_index):
    """Per-puzzle difficulty level = max strategy digit over solution cells.

    strategy_codes: (N, 81) int64 chain codes (digits = strategy applications).
    Clue cells have code 0. Level clipped to [3, 8] (rare all-lone-single
    puzzles fold into level 3).
    """
    n = len(strategy_codes)
    levels = np.zeros(n, dtype=np.int32)
    chunk = 200000
    for lo in range(0, n, chunk):
        x = strategy_codes[lo:lo + chunk].astype(np.int64).copy()
        m = np.zeros_like(x)
        while x.any():
            np.maximum(m, x % 10, out=m)
            x //= 10
        levels[lo:lo + chunk] = m.max(axis=1).astype(np.int32)
    return np.clip(levels, 3, 8)


def create_dataset(config, bs, train, curriculum=None):
    """Create Sudoku dataset according to the config.

    Args:
        config: a config object containing the hyparameters for the dataset
            creation.
        bs: batch size
        train: whether the dataset is for train or eval

    Returns:
        a tf.data.Dataset object
    """
    ds, output_types, output_shapes = None, None, None
    ds = SudokuDataset(config, train=train, curriculum=curriculum)
    # Each example is (sequence with latent slots, solution, start_index,
    # difficulty level, per-slot candidate-set targets, round-bin). The
    # round-bin is 0 unless the round-count data curriculum is enabled.
    K = int(getattr(config, "num_latent_slots", 0))
    output_types = (tf.int32, tf.int32, tf.int32, tf.int32, tf.int32, tf.int32)
    output_shapes = (
        tf.TensorShape([config.seq_len]),
        tf.TensorShape([config.block_size]),
        tf.TensorShape([1]),
        tf.TensorShape([1]),
        tf.TensorShape([K, config.block_size]),
        tf.TensorShape([1]),
    )

    # Create a tf.data.Dataset object from the generator.
    tf_ds = tf.data.Dataset.from_generator(
        generator=ds, output_types=output_types, output_shapes=output_shapes)

    # Repeat the dataset indefinitely.
    tf_ds = tf_ds.repeat()
    # Shuffle the dataset with a buffer size of 8 * bs and a seed of 0.
    tf_ds = tf_ds.shuffle(8 * config.minibatch_size, seed=0)
    # Batch the dataset with a batch size of bs.
    tf_ds = tf_ds.batch(bs)
    return tf_ds



def prepare_tf_data(xs):
    """Convert a input batch from tf Tensors to numpy arrays."""
    def _prepare(x):
        return x._numpy()  # pylint: disable=protected-access
    
    return jax.tree_map(_prepare, xs)


def create_iter(config, bs, train, curriculum=None):
    tf_ds = create_dataset(config, bs, train=train, curriculum=curriculum)
    it = map(prepare_tf_data, tf_ds)
    return it

class SudokuDataset:
    """Sudoku dataset."""
    def __init__(self, config, train=True, curriculum=None):
        self.config = config
        self.train = train
        self.curriculum = curriculum
        self.num_latent_slots = int(getattr(config, "num_latent_slots", 0))
        self.latent_token_id = int(getattr(config, "latent_token_id", 10))
        self.rng = np.random.RandomState(config.seed if hasattr(config, "seed") else 0)
        self.preprocess_sudoku()
        self._load_candidate_masks()
        self._load_round_bins()
        self._load_instances()

    def _load_instances(self):
        """Load superposition instances: one concrete assignment per row.

        Replaces the multi-hot candidate target with ordinary value tokens. For
        a puzzle at stage s there are several assignments, each picking one digit
        per cell from that cell's stage-s candidate set, so the candidate set is
        recoverable across instances instead of being supervised as a set.

        assignments: (M, 81) uint8, cell = r*9+c
        starts/counts: (N, S) row range for each (puzzle, stage)
        """
        self.instances = None
        d = getattr(self.config, "instance_dir", None)
        if not d:
            return
        split = "train" if self.train else "test"
        self.instances = np.load(
            os.path.join(d, f"{split}_assignments.npy"), mmap_mode="r")
        self.inst_starts = np.load(os.path.join(d, f"{split}_starts.npy"))
        self.inst_counts = np.load(os.path.join(d, f"{split}_counts.npy"))
        self.inst_stages = int(self.inst_starts.shape[1])
        print(f"[inst] loaded {split} instances {self.instances.shape} "
              f"over {self.inst_starts.shape[0]} puzzles, "
              f"{self.inst_stages} stages", flush=True)

    def instance_stage(self):
        """Stage whose instances this example should target.

        Bound to the curriculum stage so the target's ambiguity matches the
        latent depth: stage t runs t latent slots and supervises the stage-(t-1)
        assignments, ending at the unique solution when t == S.
        """
        S = self.inst_stages
        if self.curriculum is None:
            return S - 1
        t = int(np.clip(self.curriculum.stage, 1, S))
        return t - 1

    def instance_values(self, idx, stage):
        """One assignment for (puzzle idx, stage), sampled uniformly."""
        n = int(self.inst_counts[idx, stage])
        if n <= 0:
            return None
        row = int(self.inst_starts[idx, stage]) + self.rng.randint(n)
        return np.asarray(self.instances[row])

    def apply_instance(self, seq, idx, stage):
        """Rewrite the value token of every triple to this instance's digit.

        The (row, col) order is untouched, so the clue block and the solver-order
        output sequence are exactly as before; only the values change.
        """
        vals = self.instance_values(idx, stage)
        if vals is None:
            return seq
        seq = seq.copy()
        cells = seq[0::3].astype(np.int64) * 9 + seq[1::3].astype(np.int64)
        seq[2::3] = vals[cells].astype(seq.dtype)
        return seq

    def _load_round_bins(self):
        """Load per-puzzle solver round counts and bin them into max_stage
        equal-count bins, for the round-count DATA curriculum.

        The round count (waves needed to reach the unique solution, 5..38) is
        the same propagation-depth axis the latent curriculum supervises, but
        used to order the *puzzles* instead of the supervision. Bin edges are
        always computed on the train split and reused for eval so a bin index
        means the same thing in both.
        """
        self.round_bins = None
        self.num_bins = int(getattr(self.config, "curriculum_max_stage", 12))
        if str(getattr(self.config, "data_curriculum", "none")) != "rounds":
            return
        tr_path = getattr(self.config, "train_meta_path", None)
        path = tr_path if self.train else getattr(
            self.config, "test_meta_path", None)
        if not (path and tr_path):
            raise ValueError(
                "data_curriculum='rounds' needs SUDOKU_TRAIN_META and "
                "SUDOKU_TEST_META (the *_meta.npy written by "
                "staged_candidate_gen.py; column 2 is num_rounds)")
        rounds = np.load(path, mmap_mode="r")[:, 2].astype(np.int32)
        # Cut points from the TRAIN split, so a bin index means the same thing
        # in eval. Round counts are integers with a peaked distribution (mean
        # 22, sd 4), so raw quantiles collide -- the 12-bin quantiles repeat 20
        # twice on the full corpus, which would leave a bin permanently empty
        # and strand its stage with no frontier to measure. Force the cuts
        # strictly increasing so every bin is reachable.
        train_rounds = np.load(tr_path, mmap_mode="r")[:, 2].astype(np.int32)
        edges = np.quantile(train_rounds, np.linspace(0, 1, self.num_bins + 1))
        cuts = np.round(edges[1:-1]).astype(np.int64)
        for i in range(1, len(cuts)):
            if cuts[i] <= cuts[i - 1]:
                cuts[i] = cuts[i - 1] + 1
        # bin j (1-based) = stage that first unlocks the puzzle.
        self.round_bins = np.clip(
            np.searchsorted(cuts, rounds, side="right") + 1,
            1, self.num_bins).astype(np.int32)
        self.bin_index = {b: np.where(self.round_bins == b)[0]
                          for b in range(1, self.num_bins + 1)}
        counts = {b: int(len(v)) for b, v in self.bin_index.items()}
        print(f"[rounds] {'train' if self.train else 'eval'} bin counts:",
              counts, flush=True)
        print(f"[rounds] cuts: {cuts.tolist()}  (rounds "
              f"{int(rounds.min())}..{int(rounds.max())})", flush=True)
        empty = [b for b, c in counts.items() if c == 0]
        if empty and self.train:
            raise ValueError(
                f"round-bin curriculum has empty train bins {empty}; those "
                f"stages would have no puzzles and no frontier signal")

    def _load_candidate_masks(self):
        """Load staged candidate-set masks (N, S, 81) uint16, aligned by puzzle
        index with the loaded .npy. Row i here == puzzle i in the base file."""
        if self.train:
            path = getattr(self.config, "train_cand_masks_path", None)
        else:
            path = getattr(self.config, "test_cand_masks_path", None)
        self.cand_masks = None
        self.num_stages = 0
        if path:
            self.cand_masks = np.load(path, mmap_mode="r")
            self.num_stages = int(self.cand_masks.shape[1])
            print(f"[cand] loaded {path} shape {self.cand_masks.shape}", flush=True)

    def slot_budget(self, level):
        """Number of latent slots this example activates (see cand_slot_mode).

        "depth" mode must agree with the recurrence depth used by the train
        step, since build_latent_state only ever writes slots [0, num_passes):
        supervising a slot the recurrence never filled would train the head off
        an all-zero latent.
        """
        K = self.num_latent_slots
        if getattr(self.config, "cand_slot_mode", "level") == "depth":
            stage = self.curriculum.stage if self.curriculum is not None \
                else getattr(self.config, "curriculum_max_stage", 6)
            pps = int(getattr(self.config, "passes_per_stage", 1))
            return int(np.clip(pps * stage, 1, K))
        return int(np.clip(level - 2, 1, K))

    def _slot_stage_targets(self, idx, level, clue_cells=None):
        """Return (K, 81) int32 candidate bitmasks, one per latent slot.

        The S stored stages are mapped onto the example's k active slots; see
        cand_slot_mode for the two mappings ("level" re-paces the whole shrink
        sequence into k slots, "depth" assigns slot j to stage j). Inactive slots
        (j>=k) default to the final stage; they are masked out of the loss.

        Although the array is laid out over all 81 cell positions (for a fixed
        batch shape), the supervised targets are only the *empty* cells: clue
        cells are zeroed out here as a sentinel (a genuine empty cell always has
        >=1 candidate at every stage), and the loss ignores zero rows. So the
        effective target per puzzle is (#empty cells) x 9, in solver order."""
        K = self.num_latent_slots
        if self.cand_masks is None or K == 0:
            return np.zeros((K, 81), dtype=np.int32)
        S = self.num_stages
        k = self.slot_budget(level)
        depth_mode = getattr(self.config, "cand_slot_mode", "level") == "depth"
        stages = self.cand_masks[idx].astype(np.int32)   # (S, 81)
        out = np.zeros((K, 81), dtype=np.int32)
        for j in range(K):
            if j >= k:
                # Inactive slot: masked out of the loss, value is irrelevant.
                s = S - 1
            elif depth_mode:
                # Identity: slot j holds propagation block j, so growing the
                # recurrence depth extends the chain instead of re-pacing it.
                # The solution is only reached at full depth, which is what
                # makes this a curriculum over reasoning depth.
                s = min(j, S - 1)
            else:
                # Active slots span the full shrink sequence: slot 0 -> stage 0
                # (widest candidate set, genuinely multi-valued), last active
                # slot -> final stage (solution). For k==1 the single slot maps
                # to the WIDEST set (stage 0), not the solution, so even level-3
                # puzzles give the candidate head a real multi-candidate target
                # (the LM head still produces the unique answer).
                s = int(round(j * (S - 1) / max(k - 1, 1)))
            out[j] = stages[s]
        # Sentinel-zero the clue cells so only the empty cells are supervised.
        if clue_cells is not None and len(clue_cells) > 0:
            out[:, clue_cells] = 0
        return out

    def _build_level_index(self, levels):
        """Map difficulty level -> array of puzzle indices."""
        return {lvl: np.where(levels == lvl)[0] for lvl in range(3, 9)}

    def insert_latent_slots(self, seq, start_index):
        """Insert K latent placeholder tokens between clues and solution.

        seq: (243,) triple sequence. Returns (243 + K,) sequence:
        [clues (3*si)] [K placeholders] [solution triples].
        """
        k = self.num_latent_slots
        if k == 0:
            return seq
        si3 = 3 * int(start_index)
        return np.concatenate([
            seq[:si3],
            np.full(k, self.latent_token_id, dtype=seq.dtype),
            seq[si3:],
        ])


    def convert_to_fixed_or_random_order(self, inputs, start_index):
        """Convert the sequence of moves to either a fixed or random order.

        Args:
            inputs: a numpy array of shape (num_puzzles, seq_len) containing the
                sequence of moves for each puzzle
            start_index: a numpy array of shape (num_puzzles, 1) containing the starting
                index for each puzzle

        Returns:
            transformed_input: a numpy array of shape (num_puzzles, seq_len) containing the
                sequence of moves for each puzzle in either a fixed or random order
        """
        transformed_input = np.zeros_like(inputs)
        
        for i in range(len(inputs)):
            cur_seq = inputs[i]
            cur_start_index = start_index[i, 0]
            
            # Split the sequence into input and output prompts
            inp_prompt = cur_seq[ :(3 * cur_start_index) ].reshape(-1, 3)
            out_prompt = cur_seq[ (3 * cur_start_index): ].reshape(-1, 3)
            
            # Sort the input prompts in a fixed order
            if self.config.seq_order == "fixed":
                transformed_input[i, :(3 * cur_start_index) ] = inp_prompt[ np.lexsort( inp_prompt[:, ::-1].T ) ].flatten()
            # Randomly shuffle the input prompts
            elif self.config.seq_order == "random":
                transformed_input[i, :(3 * cur_start_index) ] = np.random.permutation(inp_prompt).flatten()
            
            # Sort the output prompts in a fixed order
            if self.config.seq_order == "fixed":
                transformed_input[i, (3 * cur_start_index): ] = out_prompt[ np.lexsort( out_prompt[:, ::-1].T ) ].flatten()
            # Randomly shuffle the output prompts
            elif self.config.seq_order == "random":
                transformed_input[i, (3 * cur_start_index): ] = np.random.permutation(out_prompt).flatten()
        
        return transformed_input

    def get_puzzles_start_index(self, path):
        """Get the puzzles, start index, inputs and difficulty levels.

        Returns:
            inputs: (num_puzzles, 243) move sequences (strategy column removed)
            puzzles: (num_puzzles, 81) solutions
            start_index: (num_puzzles, 1) number of clue cells
            levels: (num_puzzles,) puzzle difficulty level in [3, 8]
                (= hardest solver-strategy digit needed by any cell)
        """
        with gfile.Open(path, "rb") as f:
            inputs_with_start_index = np.load(f)
        start_index = inputs_with_start_index[:, 0]  # Get the start index

        rest = inputs_with_start_index[:, 1:]
        # Strategy chain codes (4th entry of each cell quadruple); keep them to
        # derive the curriculum difficulty level, then remove from the inputs.
        strategy_codes = rest.reshape(len(rest), 81, 4)[:, :, 3]
        levels = compute_puzzle_levels(strategy_codes, start_index)
        inputs = np.delete( rest, np.arange(81) * 4 + 3, axis=1)
        
        puzzles = np.zeros((len(inputs), 81), dtype=np.int8)  # Initialize puzzles
        for j in range(81):
            cell_id = inputs[:, 3 * j] * 9 + inputs[:, 3 * j + 1]  # Get the cell id
            puzzles[np.arange(len(inputs)), cell_id] = inputs[:, 3 * j + 2]  # Set the puzzle
        
        return inputs, puzzles, start_index.reshape(-1, 1), levels
    
    
    def preprocess_sudoku(self):
        """Preprocess the sudoku for train and test datasets.
        
        Depending on the `train` flag, this method loads and processes the
        sudoku puzzles and their start indices from the appropriate paths, and
        optionally converts them to a fixed or random order based on the 
        configuration.
        """
        if self.train is True:
            # Load train puzzles, inputs, and start indices
            (self.train_inputs, self.train_puzzles, self.train_start_index,
             self.train_levels) = (
                self.get_puzzles_start_index(self.config.train_puzzle_path)
            )
            # Convert train inputs to fixed or random order if specified
            if self.config.seq_order in {"fixed", "random"}:
                self.train_inputs = self.convert_to_fixed_or_random_order(self.train_inputs, self.train_start_index)
            self.level_index = self._build_level_index(self.train_levels)
            print("train level counts:",
                  {l: len(v) for l, v in self.level_index.items()}, flush=True)
        
        elif self.train is False:
            # Load evaluation puzzles, inputs, and start indices
            (self.eval_inputs, self.eval_puzzles, self.eval_start_index,
             self.eval_levels) = (
                self.get_puzzles_start_index(self.config.test_puzzle_path)
            )
            # Convert evaluation inputs to fixed or random order if specified
            if self.config.seq_order in {"fixed", "random"}: 
                self.eval_inputs = self.convert_to_fixed_or_random_order(self.eval_inputs, self.eval_start_index)
            self.level_index = self._build_level_index(self.eval_levels)

    def __len__(self):
        if self.train is True:
            return len(self.train_puzzles)
        elif self.train is False:
            return len(self.eval_puzzles)
        
    def __getitem__(self, idx):
        """Returns one example: (sequence with latent slots, solution,
        start_index, difficulty level).

        The base sequence is 243 tokens of (row, column, value) triples; K
        latent placeholder tokens are inserted after the clue block, giving
        243 + K tokens. start_index is the number of clue cells; level in
        [3, 8] is the hardest solver strategy needed by any cell.
        """
        if self.train is True:
            inputs, puzzles = self.train_inputs, self.train_puzzles
            start_index, levels = self.train_start_index, self.train_levels
        else:
            inputs, puzzles = self.eval_inputs, self.eval_puzzles
            start_index, levels = self.eval_start_index, self.eval_levels

        base = inputs[idx, :]
        if self.instances is not None and self.train:
            # Same input prompt, different output prompt: the clue triples are
            # untouched (their instance digit is the clue) while the empty cells
            # take one draw from the stage's candidate sets. Averaged over the
            # instances the target IS the candidate set, so the superposition is
            # learned from ordinary next-token CE instead of a set head. Eval
            # keeps the unique solution: the sequence it scores is generated, and
            # `puzzles` must stay the ground truth the accuracy is measured on.
            base = self.apply_instance(base, idx, self.instance_stage())
        seq = self.insert_latent_slots(base, start_index[idx, 0])
        # Clue cells = the first `start_index` (r,c,v) triples; their cell ids
        # are excluded from candidate supervision (only empty cells are scored).
        si = int(start_index[idx, 0])
        clue_triples = inputs[idx, :3 * si].reshape(-1, 3)
        clue_cells = (clue_triples[:, 0] * 9 + clue_triples[:, 1]).astype(np.int64)
        cand_targets = self._slot_stage_targets(idx, int(levels[idx]), clue_cells)
        rbin = (int(self.round_bins[idx]) if self.round_bins is not None else 0)
        return (
            seq,
            puzzles[idx, :],
            start_index[idx],
            np.array([levels[idx]], dtype=np.int32),
            cand_targets,
            np.array([rbin], dtype=np.int32),
        )

    def _sample_level(self):
        """Uniform over levels that have puzzles. Used for level-balanced mode
        and for eval, where per-level reporting needs every level represented."""
        avail = [l for l in range(3, 9) if len(self.level_index[l]) > 0]
        return avail[self.rng.randint(len(avail))]

    def _sample_round_gated(self):
        """Uniform over puzzles whose round-bin is already unlocked (bin<=stage).
        Reading the stage at yield time lets the pool grow on promotion."""
        stage = self.curriculum.stage if self.curriculum is not None \
            else self.num_bins
        stage = int(np.clip(stage, 1, self.num_bins))
        pool = np.concatenate([self.bin_index[b] for b in range(1, stage + 1)])
        return int(pool[self.rng.randint(len(pool))])

    def __call__(self):
        # Infinite generator. Train draws puzzles uniformly from the whole
        # corpus (natural difficulty mix, ~68% level 3): the difficulty tag
        # selects nothing, since the curriculum axis is propagation depth.
        # Eval stays level-balanced so per-level accuracy is measurable and
        # comparable across runs.
        #
        # The exception is the round-count DATA curriculum, where the train pool
        # is restricted to puzzles needing at most stage-many propagation waves.
        # Eval is never gated: it must score the whole corpus at every stage.
        round_gated = (self.train and self.round_bins is not None)
        level_balanced = (not self.train) or bool(
            int(getattr(self.config, "level_balanced_sampling", 0)))
        n = len(self.train_puzzles) if self.train else len(self.eval_puzzles)
        while True:
            if round_gated:
                idx = self._sample_round_gated()
            elif level_balanced:
                idx_arr = self.level_index[self._sample_level()]
                idx = int(idx_arr[self.rng.randint(len(idx_arr))])
            else:
                idx = int(self.rng.randint(n))
            yield self.__getitem__(idx)