code/train/evaluater.py

"""Evaluation related functions."""

from flax.training import common_utils
import jax
from jax import numpy as jnp
import numpy as np

from train import model

import pdb



def valid_solution(output_seq):
    """
    This function checks if the puzzle is a valid solution by verifying if
    each row, column and box has all the numbers from 1 to 9.

    Args:
        output_seq: a numpy array of shape (243,) containing the sequence of
            output numbers

    Returns:
        int: 1 if correct solution, otherwise returns 0
    """
    # rows[i, j] keeps track if ith row has received (j + 1) number
    rows = np.zeros((9, 9))
    # cols[i, j] keeps track if ith column has received (j + 1) number
    cols = np.zeros((9, 9))
    # boxes[i, j] keeps track if ith box has received (j + 1) number
    boxes = np.zeros((9, 9))

    for j in range(81):
        # The row and column are in the range (0, 8) and puzzle values are in (1, 9)
        if int(output_seq[3 * j]) >= 9:
            return False
        if int(output_seq[3 * j + 1]) >= 9:
            return False
        if int(output_seq[3 * j + 2]) > 9:
            return False
        
        row_num = int(output_seq[3 * j])
        col_num = int(output_seq[3 * j + 1])
        
        # Mark the number in the row, column and box
        rows[row_num, int(output_seq[3 * j + 2] - 1)] += 1
        cols[col_num, int(output_seq[3 * j + 2] - 1)] += 1
        boxes[
            int(3 * (row_num // 3) + (col_num // 3)), int(output_seq[3 * j + 2] - 1)
        ] += 1

    if np.all(rows) and np.all(cols) and np.all(boxes):
        return True
    else:
        return False


def eval_step(state, batch, latent_vals, slot_pos, latent_active, config):
    pred_logits, hidden, cand_logits = model.TransformerLMHeadModel(config).apply(
        {"params": state.params}, batch, latent_values=latent_vals,
        latent_positions=slot_pos, latent_active=latent_active,
        )
    return pred_logits, hidden, cand_logits


def verify_sudoku_board(puzzle, row_num, col_num, num):
    """
    Args:
        puzzle (np.array): The correct Sudoku puzzle.
        row_num (int): The row number (0-8).
        col_num (int): The column number (0-8).
        num (int): The number predicted at the specified row and column.

    Raises:
        AssertionError: If the row_num * 9 + col_num >= 81 or if the number at the specified row and column is not equal to the given number.
    """
    if row_num * 9 + col_num >= 81: 
        assert False
  
    assert puzzle[row_num * 9 + col_num] == num


def get_eval_metrics(state, eval_data_iter, p_eval_step, config):
    """This function computes given evaluation metrics (e.g, accuracy) in eval metrics for each batch and appends the metric in the list of eval_metrics.

    Args: 
        state: contains model parameters, optimizer, etc.
        eval_data_iter: data iterator for evaluation dataset
        p_eval_step: pmap function for forward pass of model for evaluation
        config: general experiment config file

    Returns: 
        eval_metrics: contains list of evaluation metrics for each batch
    """

    eval_metrics = {
        "acc": [],  # Value/placement acc: correct digit at the model-chosen cell
        "loc_acc": [],  # Location acc: model picks the ground-truth next cell (r,c)
        "val_given_loc_acc": [],  # Correct digit AMONG steps where location matched
        "cand_bit_acc": [],  # Per-digit accuracy of predicted candidate masks
        "cand_set_acc": [],  # Exact candidate-SET match per empty cell (all 9 bits)
        "cand_set_acc_changed": [],  # ...restricted to cells that changed this stage
        "acc_complete_puzzle": []  # Accuracy of predicting correct complete puzzle
    }
    # Per-difficulty-level cell accuracy (levels 3..8). Diagnostic only: the
    # curriculum no longer keys on level.
    level_ok = {lvl: 0 for lvl in range(3, 9)}
    level_tot = {lvl: 0 for lvl in range(3, 9)}

    K = int(config.num_latent_slots)

    # Per-SLOT candidate-set accuracy, i.e. per reasoning depth. Slot j holds
    # wave snapshot j, so slot_ok[j]/slot_tot[j] is "how well is propagation
    # block j predicted". This is the signal the depth curriculum promotes and
    # backtracks on, replacing the old per-level accuracy.
    slot_ok = np.zeros(max(K, 1), dtype=np.int64)
    slot_tot = np.zeros(max(K, 1), dtype=np.int64)
    slot_ok_ch = np.zeros(max(K, 1), dtype=np.int64)
    slot_tot_ch = np.zeros(max(K, 1), dtype=np.int64)

    # Per-stage in-set rate: was the emitted digit a *member* of that stage's
    # candidate set? This is the promotion signal for the instance arm, where
    # the target is one sampled assignment rather than the unique solution, so
    # the model is right to emit any candidate. The candidate masks are read
    # here as a metric only; nothing supervises them.
    inset_ok = np.zeros(max(K, 1), dtype=np.int64)
    inset_tot = np.zeros(max(K, 1), dtype=np.int64)

    # Per-round-bin cell accuracy, for the round-count DATA curriculum: bin b is
    # unlocked at stage b, so bin_ok[b]/bin_tot[b] measures competence on the
    # puzzles that stage b introduced. This is the promotion signal for the arm
    # that has no latent slots and therefore no per-depth signal.
    n_bins = int(getattr(config, "curriculum_max_stage", 12))
    bin_ok = {b: 0 for b in range(1, n_bins + 1)}
    bin_tot = {b: 0 for b in range(1, n_bins + 1)}

    for eval_epoch in range(config.eval_epochs):
        with jax.profiler.StepTraceAnnotation("eval", step_num=eval_epoch):

            batch_tuple = next(eval_data_iter)

            # Input seq is (batchsize, 3*81 + K): clue triples, K latent
            # placeholder slots, then solution triples.
            input_seq = np.array(batch_tuple[0])

            # Puzzle solution is of the shape (batchsize, 81). Each pos in {0,.., 80}
            # for each puzzle contains value at cell (pos//9+1, pos%9 + 1)
            puzzle_sol = np.array(batch_tuple[1])
            start_index = np.array(batch_tuple[2])
            levels = np.array(batch_tuple[3]).reshape(-1)
            rbins = (np.array(batch_tuple[5]).reshape(-1)
                     if len(batch_tuple) > 5 else np.zeros_like(levels))
            total_pred, sucess_pred = 0, 0
            # Location = did the model emit the ground-truth next (r,c) cell.
            loc_tot, loc_ok, val_given_loc_ok = 0, 0, 0

            bs = input_seq.shape[0]
            bidx = np.arange(bs)
            si3 = 3 * start_index.reshape(-1)
            slot_pos = si3[:, None] + np.arange(K)[None, :]
            if getattr(config, "cand_slot_mode", "level") == "depth":
                # Eval always builds all K latents, so score all K slots.
                k_budget = np.full_like(levels, K)
            else:
                k_budget = np.clip(levels - 2, 1, K)
            active_full = np.arange(K)[None, :] < k_budget[:, None]

            def run_model(seq_batch, latent_vals, act, want_cand=False):
                sharded = common_utils.shard(
                    jax.tree_util.tree_map(np.asarray, seq_batch))
                # Explicit reshape so a zero-width slot dim (K=0 baseline)
                # shards without the ambiguous -1 inference of shard().
                _nd = jax.local_device_count()
                def _shard(x):
                    x = np.asarray(x)
                    return x.reshape((_nd, x.shape[0] // _nd) + x.shape[1:])
                lv = _shard(latent_vals)
                lp = _shard(slot_pos)
                la = _shard(act)
                logits, hidden, cand = p_eval_step(state, sharded, lv, lp, la)
                logits = np.array(logits).reshape(bs, *np.array(logits).shape[2:])
                hidden = np.array(hidden).reshape(bs, *np.array(hidden).shape[2:])
                if want_cand:
                    cand = np.array(cand).reshape(bs, *np.array(cand).shape[2:])
                    return logits, hidden, cand
                return logits, hidden

            # ---- Build the continuous latent thoughts (K recurrence passes,
            # difficulty-matched budget; causal masking means only the clue
            # region influences them). ----
            latent_vals = np.zeros((bs, K, config.emb_dim), dtype=np.float32)
            build_seq = np.array(input_seq)
            build_seq_masked = np.array(build_seq)
            # Hide the solution region during latent build (safety; causality
            # already prevents leakage into slot hiddens).
            for j in range(bs):
                build_seq_masked[j, si3[j] + K:] = 0
            # Recurrent feedback: build each latent thought from the previous
            # slot's hidden. Skipped when the model does not inject latents
            # (no-recurrence control): slots stay as static placeholders, so
            # latent_vals is left at zeros and never used.
            recurrent = bool(int(getattr(config, "recurrent_latent", 1)))
            if recurrent and K > 0:
                for j in range(K):
                    act_j = active_full & (np.arange(K)[None, :] < j)
                    _, hidden = run_model(build_seq_masked, latent_vals, act_j)
                    src = si3 - 1 + j
                    latent_vals[:, j] = hidden[bidx, src]

            # ---- Candidate-set prediction accuracy (the multi-value target) ----
            # One forward pass with the fully-built latents; read the per-slot
            # candidate head and compare to the staged bitmask targets, scored
            # only over active slots and empty cells (clue cells were zeroed).
            # Skipped entirely for the K=0 no-latent baseline (no candidate head).
            pred_bits = tgt_bits = cand_targets = None
            if K > 0:
                cand_targets = np.array(batch_tuple[4]).astype(np.int64)   # (bs, K, 81)
            # The candidate head is off in the instance arm (aux weight 0), so
            # skip its forward pass and set metrics; the masks above are still
            # read for the in-set rate.
            if K > 0 and float(getattr(config, "aux_cand_weight", 1.0)) > 0.0:
                _, _, cand_logits = run_model(
                    build_seq_masked, latent_vals, active_full, want_cand=True)  # (bs,K,81,9)
                pred_bits = (np.array(cand_logits) > 0.0)                  # sigmoid>0.5
                tgt_bits = ((cand_targets[..., None] >> np.arange(9)) & 1).astype(bool)
                valid = (cand_targets > 0) & active_full[:, :, None]       # (bs,K,81)
                if valid.sum() > 0:
                    bit_match = (pred_bits == tgt_bits)                    # (bs,K,81,9)
                    eval_metrics["cand_bit_acc"].append(
                        float(bit_match[valid].mean()))
                    eval_metrics["cand_set_acc"].append(
                        float(bit_match.all(axis=3)[valid].mean()))
                    # Same score restricted to cells whose candidate set
                    # actually changed from the previous stage. The unrestricted
                    # metrics above are dominated by cells that are unchanged
                    # copies of slot j-1, so they stay high for a head that has
                    # learned nothing but "repeat the previous slot".
                    changed = np.concatenate(
                        [np.ones_like(cand_targets[:, :1], dtype=bool),
                         cand_targets[:, 1:] != cand_targets[:, :-1]], axis=1)
                    valid_ch = valid & changed
                    if valid_ch.sum() > 0:
                        eval_metrics["cand_set_acc_changed"].append(
                            float(bit_match.all(axis=3)[valid_ch].mean()))
                    # Accumulate the same score split by slot (= depth).
                    set_match = bit_match.all(axis=3)                   # (bs,K,81)
                    slot_ok += (set_match & valid).sum(axis=(0, 2))
                    slot_tot += valid.sum(axis=(0, 2))
                    slot_ok_ch += (set_match & valid_ch).sum(axis=(0, 2))
                    slot_tot_ch += valid_ch.sum(axis=(0, 2))

            min_start_index = int(np.min(start_index))
            cur_input_seq = input_seq[:, :(min_start_index*3)]
            for i in range(min_start_index * 3, config.seq_len):
                ### In i^th iteration, i^th number in sequence will predict
                padding = np.zeros((input_seq.shape[0],
                                    config.seq_len - len(cur_input_seq[0])),
                                dtype=np.int32)
                concat_batch = np.hstack((cur_input_seq, padding))

                pred_logits, _ = run_model(concat_batch, latent_vals, active_full)

                # Positions < 3*start_index + K are given (clues + latent
                # slots); the model predicts from there on. K is a multiple
                # of 3, so the triple phase of i is unchanged.
                if i%3 == 2:
                    # Model predicts the value at the cell (cur_input_seq[j][i-2],
                    # cur_input_seq[j][i-1])
                    max_number = pred_logits[:, i-1, :].argmax(axis=-1).flatten()
                    mask_arr = np.array(i >= (3 * start_index + K)).squeeze()

                    next_number = max_number * mask_arr + (1 - mask_arr) * input_seq[:, i]

                    cur_input_seq = np.hstack(
                        (cur_input_seq, np.reshape(next_number, (-1, 1)))
                    )

                    # Iterate through all examples in batch and calculate successful
                    # predictions of numbers
                    for j in range(len(cur_input_seq)):
                        if not mask_arr[j]:
                            continue

                        total_pred += 1
                        level_tot[int(levels[j])] += 1
                        if int(rbins[j]) in bin_tot:
                            bin_tot[int(rbins[j])] += 1

                        # Location accuracy: did the model emit the ground-truth
                        # next cell (r,c) for this solver-order step?
                        loc_tot += 1
                        loc_match = (int(cur_input_seq[j][i-2]) == int(input_seq[j, i-2])
                                     and int(cur_input_seq[j][i-1]) == int(input_seq[j, i-1]))
                        if loc_match:
                            loc_ok += 1

                        # In-set rate per stage, scored at the ground-truth cell
                        # so a wrong location cannot make a digit vacuously
                        # legal. cand_targets[j, s, cell] is stage s's bitmask
                        # under cand_slot_mode="depth" (slot s <-> stage s).
                        if cand_targets is not None and loc_match:
                            cell = (int(input_seq[j, i-2]) * 9
                                    + int(input_seq[j, i-1]))
                            v = int(cur_input_seq[j][i])
                            for s in range(K):
                                bits = int(cand_targets[j, s, cell])
                                if bits <= 0:      # clue cell, not supervised
                                    continue
                                inset_tot[s] += 1
                                if 1 <= v <= 9 and (bits >> (v - 1)) & 1:
                                    inset_ok[s] += 1

                        try:
                            verify_sudoku_board(puzzle_sol[j], cur_input_seq[j][i-2], 
                                                cur_input_seq[j][i-1], cur_input_seq[j][i])
                        except AssertionError:
                            # Mistake
                            pass
                        else:
                            sucess_pred += 1
                            level_ok[int(levels[j])] += 1
                            if int(rbins[j]) in bin_ok:
                                bin_ok[int(rbins[j])] += 1
                            if loc_match:
                                val_given_loc_ok += 1
                else:
                    # Model predicts either a row number or column number
                    max_pos = pred_logits[:, i-1, :].argmax(axis=-1).flatten()
                    mask = (i >= (3 * start_index + K)).squeeze()
                    next_pos = max_pos * mask + (1 - mask) * input_seq[:, i]
                    
                    # pdb.set_trace()
                    cur_input_seq = np.hstack(
                        (cur_input_seq, np.reshape(next_pos, (-1, 1)))
                    )

            eval_metrics["acc"].append(sucess_pred * 1.0/ total_pred)
            eval_metrics["loc_acc"].append(loc_ok * 1.0 / max(loc_tot, 1))
            eval_metrics["val_given_loc_acc"].append(
                val_given_loc_ok * 1.0 / max(loc_ok, 1))

            def strip_latent_slots(seq, si):
                return np.concatenate([seq[:3*si], seq[3*si + K:]])

            # ---- Print one concrete example answer the model generated ----
            if eval_epoch == 0:
                j = 0
                si = int(start_index[j, 0])
                pred = strip_latent_slots(cur_input_seq[j], si)
                shown, n_ok, n_tot = [], 0, 0
                for k in range(si, 81):
                    r, c, v = int(pred[3*k]), int(pred[3*k+1]), int(pred[3*k+2])
                    tv = int(puzzle_sol[j][r*9+c]) if (0 <= r < 9 and 0 <= c < 9) else -1
                    ok = (0 <= r < 9 and 0 <= c < 9 and v == tv)
                    n_tot += 1; n_ok += int(ok)
                    if len(shown) < 12:
                        shown.append(f"({r},{c})->{v}[true {tv}]{'ok' if ok else 'X'}")
                print(f"EXAMPLE (level={int(levels[j])}, k={int(k_budget[j])}): "
                      f"model emitted {n_tot} (r,c)->v triples for the empty cells "
                      f"(format: (row,col)->value[true T]); first 12:", flush=True)
                print("   ", " ".join(shown), flush=True)
                print(f"EXAMPLE cells-correct={n_ok}/{n_tot}  "
                      f"valid_full_grid={valid_solution(pred)}", flush=True)

                # Instance arm: emitted digit next to the deepest stage's
                # candidate set, so it is visible whether the model is sitting
                # inside the superposition or outside it.
                if K > 0 and cand_targets is not None and pred_bits is None:
                    tgt = strip_latent_slots(input_seq[j], si)
                    shown = []
                    for t3 in range(si, min(si + 8, 81)):
                        r, c = int(tgt[3*t3]), int(tgt[3*t3+1])
                        bits = int(cand_targets[j, K-1, r*9+c])
                        cset = "".join(str(d+1) for d in range(9)
                                       if (bits >> d) & 1)
                        shown.append(f"(r{r},c{c})->{int(pred[3*t3+2])} "
                                     f"in{{{cset}}}")
                    print(f"EXAMPLE emitted vs stage-{K} candidate set:",
                          " ".join(shown), flush=True)

                # ---- Candidate-set (multi-value) prediction for this puzzle ----
                # Show, at the last active latent slot, predicted vs target
                # candidate SETS for the first few empty cells. (No latent
                # slots in the K=0 baseline, so nothing to show.)
                if K > 0 and pred_bits is not None:
                    kj = int(k_budget[j]) - 1
                    def _digs(bitrow):
                        return "".join(str(d + 1) for d in range(9) if bitrow[d])
                    cand_shown = []
                    for cell in range(81):
                        if cand_targets[j, kj, cell] <= 0:   # clue / not supervised
                            continue
                        r, c = cell // 9, cell % 9
                        pset = _digs(pred_bits[j, kj, cell])
                        tset = _digs(tgt_bits[j, kj, cell])
                        cand_shown.append(f"(r{r},c{c}) pred{{{pset}}} true{{{tset}}}")
                        if len(cand_shown) >= 8:
                            break
                    print(f"EXAMPLE candidate-set @slot{kj} (pred vs true):",
                          " ".join(cand_shown), flush=True)

            correct_eval_sudoku_puzzle = 0

            for i in range(len(cur_input_seq)):

                # increase correct_eval_sudoku_puzzle when the model output solution
                # for a given puzzle is correct
                stripped = strip_latent_slots(cur_input_seq[i], int(start_index[i, 0]))
                correct_eval_sudoku_puzzle += valid_solution(stripped)

            eval_metrics["acc_complete_puzzle"].append(
                correct_eval_sudoku_puzzle * 1.0 / len(cur_input_seq)
            )

    per_level = {lvl: (level_ok[lvl] / level_tot[lvl] if level_tot[lvl] else -1.0)
                 for lvl in range(3, 9)}
    eval_metrics["per_level_acc"] = per_level
    print("PER-LEVEL cell acc:",
          {lvl: (f"{v:.3f}" if v >= 0 else "n/a") for lvl, v in per_level.items()},
          flush=True)

    # Per-depth candidate-set accuracy, keyed by stage (slot j -> stage j+1) so
    # the curriculum controller can index it directly by stage number.
    per_slot = {j + 1: (float(slot_ok[j] / slot_tot[j]) if slot_tot[j] else -1.0)
                for j in range(K)}
    per_slot_ch = {j + 1: (float(slot_ok_ch[j] / slot_tot_ch[j])
                           if slot_tot_ch[j] else -1.0) for j in range(K)}
    eval_metrics["per_slot_acc"] = per_slot
    eval_metrics["per_slot_acc_changed"] = per_slot_ch

    # Keyed by stage (slot s -> stage s+1) to match per_slot_acc.
    per_stage_inset = {s + 1: (float(inset_ok[s] / inset_tot[s])
                               if inset_tot[s] else -1.0) for s in range(K)}
    eval_metrics["per_stage_inset_acc"] = per_stage_inset
    if K > 0 and any(v >= 0 for v in per_stage_inset.values()):
        print("PER-STAGE in-set rate (emitted digit is a stage-s candidate):",
              {s: (f"{v:.3f}" if v >= 0 else "n/a")
               for s, v in per_stage_inset.items()}, flush=True)

    per_bin = {b: (bin_ok[b] / bin_tot[b] if bin_tot[b] else -1.0)
               for b in range(1, n_bins + 1)}
    eval_metrics["per_bin_acc"] = per_bin
    if any(v >= 0 for v in per_bin.values()):
        print("PER-ROUND-BIN cell acc:",
              {b: (f"{v:.3f}" if v >= 0 else "n/a") for b, v in per_bin.items()},
              flush=True)
    if K > 0:
        print("PER-DEPTH cand-set acc:",
              {s: (f"{v:.3f}" if v >= 0 else "n/a") for s, v in per_slot.items()},
              flush=True)
        print("PER-DEPTH cand-set acc (changed cells only):",
              {s: (f"{v:.3f}" if v >= 0 else "n/a")
               for s, v in per_slot_ch.items()}, flush=True)

    return eval_metrics