Upload code/wavecurriculum_run/train/evaluater.py with huggingface_hub

code/wavecurriculum_run/train/evaluater.py

@@ -32,6 +32,78 @@ def _lcs_len(a, b):
     return prev[-1]
 
 
+def _inst_ce_stats(logp_digits, bits, qcounts):
+    """CE(q || p) against the ceiling log|S|, per cell.
+
+    q is the empirical per-cell digit distribution of this puzzle's stored
+    instances at this stage: the POST-constraint distribution the data actually
+    teaches. That is the whole point of using it instead of Uniform(S). By
+    Gibbs, CE(Uniform(S) || p) >= log|S| for every p, so comparing that to
+    log|S| is a tautology (it is exactly `excess >= 0`). CE(q || p) instead has
+    floor H(q), and H(q) < log|S| strictly whenever the constraints prune
+    anything, so the comparison
+
+        CE(q || p)  <=  log|S|          equivalently   KL(q || p) <= log|S| - H(q)
+
+    is both satisfiable and meaningful: the model's divergence from the truth
+    must be smaller than the information the constraints carry. It is also a
+    complete test on its own, unlike mass/spread, because it punishes every
+    failure mode at once -- dropping a candidate that instances do use sends
+    -log p(d) up, and leaking mass outside S lowers p on the digits that are
+    actually targets.
+
+    A cell is only counted when the bar is above the floor, i.e. when
+    H(q) < log|S| strictly. Two kinds of cell are dropped:
+
+      |S| = 1     the ceiling is log 1 = 0 and no finite CE can meet it.
+      H(q) = log|S|  q is already uniform on S, so the bar EQUALS the floor and
+                  by Gibbs only p == q exactly passes. Asking for that is asking
+                  for zero error, not for learning.
+
+    On the real pool this second case is 8-27% of cells depending on stage, so
+    leaving them in would cap ceq_ok well below 1 and make any fixed threshold
+    mean something different at every stage. The exclusion depends only on q and
+    S, never on p, so it cannot be gamed by the model.
+
+    Args:
+        logp_digits: (n, 9) log-probabilities of digits 1..9 under the model.
+        bits: (n,) int bitmask; bit d-1 set iff digit d is a candidate.
+        qcounts: (n, 9) instance-digit counts; column d-1 counts digit d.
+
+    Returns:
+        Sums over the countable cells plus the count they divide by.
+    """
+    out = {"ceq": 0.0, "hq": 0.0, "ceq_ok": 0, "ceq_gap": 0.0, "ceq_count": 0,
+           "ceq_all": 0, "ceq_groups": 0, "ceq_skip": 0}
+    in_s = ((bits[:, None] >> np.arange(9)) & 1).astype(np.float64)
+    size = in_s.sum(axis=1)
+    tot = qcounts.sum(axis=1)
+    cand = (size >= 2) & (tot > 0)
+    if not cand.any():
+        return out
+    q_all = qcounts[cand].astype(np.float64) / tot[cand][:, None]
+    h_all = -(q_all * np.log(np.maximum(q_all, 1e-12))).sum(axis=1)
+    l_all = np.log(size[cand])
+    # Keep only cells where the constraints left headroom under the ceiling.
+    room = (l_all - h_all) > 1e-9
+    out["ceq_skip"] = int((~room).sum())
+    if not room.any():
+        return out
+    q = q_all[room]
+    lp = logp_digits[cand][room]
+    ce_q = -(q * lp).sum(axis=1)
+    h_q = h_all[room]
+    gap = l_all[room] - ce_q
+    out["ceq"] = float(ce_q.sum())
+    out["hq"] = float(h_q.sum())
+    out["ceq_ok"] = int((gap >= 0).sum())
+    out["ceq_gap"] = float(gap.sum())
+    out["ceq_count"] = int(room.sum())
+    out["ceq_all"] = int(bool((gap >= 0).all()))
+    out["ceq_groups"] = 1
+    return out
+
+
 def _cand_value_stats(logp_digits, bits):
     """Value-head statistics against one stage's candidate sets.
 
@@ -301,6 +373,14 @@ def get_eval_metrics(state, eval_data_iter, p_eval_step, config):
     vhgap_sum = np.zeros(max(K, 1), dtype=np.float64)
     vhall_sum = np.zeros(max(K, 1), dtype=np.int64)
     vhgrp_sum = np.zeros(max(K, 1), dtype=np.int64)
+    # CE(q||p) vs the ceiling log|S|, q = the instances' own digit frequencies.
+    vceq_sum = np.zeros(max(K, 1), dtype=np.float64)
+    vhq_sum = np.zeros(max(K, 1), dtype=np.float64)
+    vceq_ok = np.zeros(max(K, 1), dtype=np.int64)
+    vceq_gap = np.zeros(max(K, 1), dtype=np.float64)
+    vceq_cnt = np.zeros(max(K, 1), dtype=np.int64)
+    vceq_all = np.zeros(max(K, 1), dtype=np.int64)
+    vceq_grp = 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
@@ -326,6 +406,10 @@ def get_eval_metrics(state, eval_data_iter, p_eval_step, config):
             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))
+            # (bs, K, 81, 9) instance-digit counts; zeros when the run has no
+            # instance pool, in which case CE(q||p) is simply not reported.
+            q_counts = (np.array(batch_tuple[6]).astype(np.int64)
+                        if len(batch_tuple) > 6 else None)
             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
@@ -476,6 +560,17 @@ def get_eval_metrics(state, eval_data_iter, p_eval_step, config):
                         vhgap_sum[s] += st["hgap"]
                         vhall_sum[s] += st["hbound_all"]
                         vhgrp_sum[s] += st["hbound_groups"]
+                        if q_counts is not None:
+                            qs = _inst_ce_stats(
+                                logp_d[sel], bits[sel],
+                                q_counts[j, s, np.where(ok_cell, cells, 0)][sel])
+                            vceq_sum[s] += qs["ceq"]
+                            vhq_sum[s] += qs["hq"]
+                            vceq_ok[s] += qs["ceq_ok"]
+                            vceq_gap[s] += qs["ceq_gap"]
+                            vceq_cnt[s] += qs["ceq_count"]
+                            vceq_all[s] += qs["ceq_all"]
+                            vceq_grp[s] += qs["ceq_groups"]
 
             min_start_index = int(np.min(start_index))
             cur_input_seq = input_seq[:, :(min_start_index*3)]
@@ -743,6 +838,20 @@ def get_eval_metrics(state, eval_data_iter, p_eval_step, config):
     eval_metrics["per_stage_val_hbound"] = _per_stage(vhok_sum, vspread_cnt)
     eval_metrics["per_stage_val_hgap"] = _per_stage(vhgap_sum, vspread_cnt)
     eval_metrics["per_stage_val_hbound_puz"] = _per_stage(vhall_sum, vhgrp_sum)
+    # The single promotion criterion: CE(q||p) <= log|S| per cell, counted.
+    #   ceq      mean CE(q||p), to compare against floor H(q) and ceiling log|S|
+    #   hq       mean H(q), the best CE any model could achieve on this data
+    #   ceq_ok   fraction of |S|>=2 cells clearing the ceiling  <- gate on this
+    #   ceq_gap  mean slack log|S| - CE(q||p); negative means worse than uniform
+    #   ceq_puz  fraction of puzzles where every such cell clears it
+    # This subsumes mass and spread: leak lowers p on real targets, and dropping
+    # a candidate the instances use sends -log p(d) up, so both are penalized by
+    # the one number. log|S| needs no tuning and recalibrates per stage.
+    eval_metrics["per_stage_val_ceq"] = _per_stage(vceq_sum, vceq_cnt)
+    eval_metrics["per_stage_val_hq"] = _per_stage(vhq_sum, vceq_cnt)
+    eval_metrics["per_stage_val_ceq_ok"] = _per_stage(vceq_ok, vceq_cnt)
+    eval_metrics["per_stage_val_ceq_gap"] = _per_stage(vceq_gap, vceq_cnt)
+    eval_metrics["per_stage_val_ceq_puz"] = _per_stage(vceq_all, vceq_grp)
     if _verbose_eval() and K > 0 and any(v >= 0 for v in per_stage_vce.values()):
         print("PER-STAGE val_excess (0 = uniform over candidate set):",
               {s: (f"{v:.3f}" if per_stage_vce[s] >= 0 else "n/a")