synthetic_pair: add --clean boundary mode (noise token at word starts) + configurable dirs; prep reports in-word fraction

1 files changed, 30 insertions, 9 deletions

diff --git a/src/synthetic_pair.py b/src/synthetic_pair.py
index 6c6c722..4c64379 100644
--- a/src/synthetic_pair.py
+++ b/src/synthetic_pair.py
@@ -41,7 +41,7 @@ MODEL_ARGS: dict[str, Any] = dict(n_layer=6, n_head=6, n_embd=384, block_size=12
                                   bias=False, dropout=0.2)
 
 
-def prep():
+def prep(clean=False):
     with open(DATA_SRC) as f:
         text = f.read()
     n_target = max(1, int(TARGET_FREQ * len(text)))
@@ -59,10 +59,21 @@ def prep():
     rng = np.random.RandomState(42)
     struct_pos = sorted(rng.choice(trig_positions, size=n_target, replace=False).tolist())
 
-    # positions for T_noise: uniform random, disjoint from struct_pos
-    noise_pos = sorted(rng.choice(
-        [p for p in range(len(text)) if p not in set(struct_pos)],
-        size=n_target, replace=False).tolist())
+    # positions for T_noise. Default: uniform random (slices inside words,
+    # ~95% of the time — the confound the clean-boundary control fixes).
+    # clean=True: always at a word start (immediately after a space), so it
+    # sits at a clean boundary exactly like T_struct after "the ", but is
+    # still unpredictable (uniform over words).
+    if clean:
+        starts = [i + 1 for i, c in enumerate(text) if c == ' ' and i + 1 < len(text)]
+        avail = [p for p in starts if p not in set(struct_pos)]
+    else:
+        avail = [p for p in range(len(text)) if p not in set(struct_pos)]
+    noise_pos = sorted(rng.choice(avail, size=n_target, replace=False).tolist())
+
+    # sanity: fraction of noise insertions that slice inside a word
+    in_word = sum(1 for p in noise_pos
+                  if 0 < p < len(text) and (text[p - 1].isalnum() and text[p].isalnum()))
 
     # insert with offset (both sets sorted -> single merge pass)
     insertions = [(p, T_STRUCT) for p in struct_pos] + [(p, T_NOISE) for p in noise_pos]
@@ -76,9 +87,11 @@ def prep():
     out.append(text[prev:])
     modified = ''.join(out)
     assert modified.count(T_STRUCT) == modified.count(T_NOISE) == n_target
-    print(f"prep: '{T_STRUCT}' x{n_target} after '{TRIGGER.strip()}', "
-          f"'{T_NOISE}' x{n_target} random, "
-          f"freq each = {n_target/len(modified):.4%}")
+    mode = "clean-boundary" if clean else "random"
+    print(f"prep({mode}): '{T_STRUCT}' x{n_target} after '{TRIGGER.strip()}', "
+          f"'{T_NOISE}' x{n_target} {mode}, "
+          f"freq each = {n_target/len(modified):.4%}, "
+          f"in-word '#' = {in_word}/{n_target} ({in_word/n_target:.1%})")
 
     # build vocab (existing chars + the two synthetic)
     chars = sorted(set(text)) + [T_STRUCT, T_NOISE]
@@ -214,9 +227,17 @@ if __name__ == '__main__':
     ap.add_argument('--batch_size', type=int, default=16)
     ap.add_argument('--layers', default='0,1,2,3,4,5')
     ap.add_argument('--n_prompts', type=int, default=10)
+    ap.add_argument('--clean', action='store_true',
+                    help='prep: place the noise token at word starts (after '
+                         'random spaces) instead of uniform random positions')
+    ap.add_argument('--data_dir', default=DATA_DIR)
+    ap.add_argument('--out_root', default=OUT_ROOT)
+    ap.add_argument('--jlens_out', default=JLENS_OUT)
     a = ap.parse_args()
+    if a.data_dir != DATA_DIR or a.out_root != OUT_ROOT or a.jlens_out != JLENS_OUT:
+        DATA_DIR, OUT_ROOT, JLENS_OUT = a.data_dir, a.out_root, a.jlens_out
     if a.step == 'prep':
-        prep()
+        prep(clean=a.clean)
     elif a.step == 'train':
         train(a.seed, a.max_iters, a.batch_size)
     elif a.step == 'jlens':