implement Addendum 8: 6-arm token-space recurrence suite (arms A, B, C, D1, D2, D3), tests, and jobs/e8.csv

1 files changed, 10 insertions, 8 deletions

diff --git a/src/train.py b/src/train.py
index a0874c4..2523325 100644
--- a/src/train.py
+++ b/src/train.py
@@ -38,6 +38,7 @@ def make_gpu_dataset(examples: list, cfg: Config, device: torch.device, max_cach
             "y": torch.empty((0, out_max), dtype=torch.long, device=device),
             "y_in": torch.empty((0, out_max), dtype=torch.long, device=device),
             "y_mask": torch.empty((0, out_max), dtype=torch.bool, device=device),
+            "loss_mask": torch.empty((0, out_max), dtype=torch.bool, device=device),
             "size": 0,
             "is_cached": True,
         }
@@ -50,6 +51,7 @@ def make_gpu_dataset(examples: list, cfg: Config, device: torch.device, max_cach
         "y": batch["y"].to(device, non_blocking=True),
         "y_in": batch["y_in"].to(device, non_blocking=True),
         "y_mask": batch["y_mask"].to(device, non_blocking=True),
+        "loss_mask": batch["loss_mask"].to(device, non_blocking=True),
         "size": len(examples),
         "is_cached": True,
     }
@@ -77,13 +79,13 @@ def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int | bool], examples:
             x = gpu_ds["x"][bi: bi + eval_bs]
             y = gpu_ds["y"][bi: bi + eval_bs]
             y_in_eval = gpu_ds["y_in"][bi: bi + eval_bs]
-            mask = gpu_ds["y_mask"][bi: bi + eval_bs]
+            mask = gpu_ds["loss_mask"][bi: bi + eval_bs]
         else:
             batch = make_batch(chunk_examples, cfg)
             x = batch["x"].to(device)
             y = batch["y"].to(device)
             y_in_eval = batch["y_in"].to(device)
-            mask = batch["y_mask"].to(device)
+            mask = batch["loss_mask"].to(device)
 
         with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
             out = model(x, y_in_eval)
@@ -93,10 +95,9 @@ def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int | bool], examples:
         token_total += int(mask.sum().item())
 
         B = x.shape[0]
-        max_out_len = getattr(cfg, "max_out_len", 6)
         from src.data import _global_lengths
         _, out_max = _global_lengths(cfg)
-        max_len = max(max_out_len, out_max)
+        max_len = max(getattr(cfg, "max_out_len", 6), out_max)
 
         y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=device)
         with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
@@ -106,10 +107,11 @@ def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int | bool], examples:
                 y_in = torch.cat([y_in, nxt[:, None]], dim=1)
         gen = y_in[:, 1:].cpu()
 
-        for i, (xrow, target) in enumerate(chunk_examples):
+        for i, ex in enumerate(chunk_examples):
+            xrow, target = ex[0], ex[1]
             pred = decode_tokens(gen[i].tolist(), cfg)
             target_n = decode_tokens(target, cfg)
-            ok = pred == target_n
+            ok = (pred == target_n) and (pred != -1)
             em_correct += int(ok)
             per_example.append((tuple(xrow), target_n, pred, ok))
 
@@ -268,7 +270,7 @@ def main() -> None:
                 x = train_gpu["x"][idx]
                 y = train_gpu["y"][idx]
                 y_in = train_gpu["y_in"][idx]
-                mask = train_gpu["y_mask"][idx]
+                mask = train_gpu["loss_mask"][idx]
 
                 with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
                     out = model(x, y_in)
@@ -317,7 +319,7 @@ def main() -> None:
                 x = batch["x"].to(device)
                 y = batch["y"].to(device)
                 y_in = batch["y_in"].to(device)
-                mask = batch["y_mask"].to(device)
+                mask = batch["loss_mask"].to(device)
 
                 with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
                     out = model(x, y_in)