# 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)