path: root/src/synthetic_pair.py

"""
Frequency-Matched Synthetic Pair Test (Gemini 3.1 Pro's design).

Two synthetic tokens at IDENTICAL unigram frequency in an otherwise-normal
Shakespeare corpus:
  T_struct '@' : appears only after the trigger sequence "the " (high conditional
                 predictability, structured context).
  T_noise  '#' : injected at uniform random positions (zero conditional structure).

Hypotheses:
  Frequency-only:  J-lens norms of '@' and '#' are identical at every layer
                   (same unigram frequency, same rarity).
  Structure/workspace (Anthropic): '@' maintains a higher faithful J-lens norm,
                   especially in intermediate layers (model tracks the trigger
                   context in the residual stream).

Uses the FAITHFUL J-lens (jlens_v3: rows of W_U * J_l) plus the old proxy.

Steps:
  python3 src/synthetic_pair.py --step prep                      # build data/synth_pair
  python3 src/synthetic_pair.py --step train --seed 0 [--max_iters 3000]
  python3 src/synthetic_pair.py --step jlens --seed 0 [--layers 0,1,2,3,4,5]
  python3 src/synthetic_pair.py --step summary                   # compare @ vs #
"""
import sys, os, argparse, subprocess, pickle, random
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from typing import Any
import numpy as np
import torch

DATA_SRC = 'data/shakespeare_char/input.txt'
DATA_DIR = 'data/synth_pair'
OUT_ROOT = 'out-synth-pair'
JLENS_OUT = 'outputs/synth_pair'
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
T_STRUCT = '@'
T_NOISE = '#'
TRIGGER = 'the '
TARGET_FREQ = 0.001          # 0.1%
MODEL_ARGS: dict[str, Any] = dict(n_layer=6, n_head=6, n_embd=384, block_size=128,
                                  bias=False, dropout=0.2)


def prep(clean=False):
    with open(DATA_SRC) as f:
        text = f.read()
    n_target = max(1, int(TARGET_FREQ * len(text)))

    # positions for T_struct: after occurrences of TRIGGER
    trig_positions = []
    start = 0
    while True:
        i = text.find(TRIGGER, start)
        if i < 0:
            break
        trig_positions.append(i + len(TRIGGER))
        start = i + len(TRIGGER)
    assert len(trig_positions) >= n_target, f"only {len(trig_positions)} triggers"
    rng = np.random.RandomState(42)
    struct_pos = sorted(rng.choice(trig_positions, 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]
    insertions.sort()
    out = []
    prev = 0
    for pos, ch in insertions:
        out.append(text[prev:pos])
        out.append(ch)
        prev = pos
    out.append(text[prev:])
    modified = ''.join(out)
    assert modified.count(T_STRUCT) == modified.count(T_NOISE) == n_target
    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]
    stoi = {c: i for i, c in enumerate(chars)}
    itos = {i: c for i, c in enumerate(chars)}
    data = np.array([stoi[c] for c in modified], dtype=np.uint16)
    n = int(0.9 * len(data))
    os.makedirs(DATA_DIR, exist_ok=True)
    data[:n].tofile(os.path.join(DATA_DIR, 'train.bin'))
    data[n:].tofile(os.path.join(DATA_DIR, 'val.bin'))
    with open(os.path.join(DATA_DIR, 'meta.pkl'), 'wb') as f:
        pickle.dump({'stoi': stoi, 'itos': itos, 'vocab_size': len(chars)}, f)
    with open(os.path.join(DATA_DIR, 'input.txt'), 'w') as f:
        f.write(modified)
    print(f"prep: vocab={len(chars)}, train={n:,} val={len(data)-n:,} tokens")


def train(seed, max_iters, batch_size):
    sys.path.insert(0, '.')
    from model import GPT, GPTConfig
    torch.manual_seed(seed)
    np.random.seed(seed)
    device = 'cuda'
    train_data = np.memmap(f'{DATA_DIR}/train.bin', dtype=np.uint16, mode='r')
    val_data = np.memmap(f'{DATA_DIR}/val.bin', dtype=np.uint16, mode='r')
    with open(f'{DATA_DIR}/meta.pkl', 'rb') as f:
        meta = pickle.load(f)
    args = dict(
        n_layer=int(MODEL_ARGS['n_layer']), n_head=int(MODEL_ARGS['n_head']),
        n_embd=int(MODEL_ARGS['n_embd']), block_size=int(MODEL_ARGS['block_size']),
        bias=bool(MODEL_ARGS['bias']), dropout=float(MODEL_ARGS['dropout']),
        vocab_size=int(meta['vocab_size']),
    )
    model = GPT(GPTConfig(**args)).to(device)
    print(f"seed {seed}: params={sum(p.numel() for p in model.parameters())/1e6:.2f}M")

    opt = model.configure_optimizers(weight_decay=0.1, learning_rate=1e-3,
                                     betas=(0.9, 0.99), device_type='cuda')
    bs = batch_size
    blk = args['block_size']

    def get_batch(split):
        d = train_data if split == 'train' else val_data
        ix = torch.randint(len(d) - blk, (bs,))
        x = torch.stack([torch.from_numpy(d[i:i+blk].astype(np.int64)) for i in ix])
        y = torch.stack([torch.from_numpy(d[i+1:i+1+blk].astype(np.int64)) for i in ix])
        return x.to(device), y.to(device)

    best_val = 1e9
    out_dir = f'{OUT_ROOT}/seed{seed}'
    os.makedirs(out_dir, exist_ok=True)
    for it in range(max_iters):
        if it % 500 == 0:
            model.eval()
            lv = []
            for _ in range(50):
                X, Y = get_batch('val')
                with torch.no_grad():
                    _, loss = model(X, Y)
                lv.append(loss.item())
            v = np.mean(lv)
            model.train()
            if v < best_val:
                best_val = v
                torch.save({'model': model.state_dict(), 'model_args': args,
                            'best_val_loss': best_val}, f'{out_dir}/ckpt.pt')
            if it % 1000 == 0:
                print(f"  iter {it}: val={v:.4f}")
        X, Y = get_batch('train')
        _, loss = model(X, Y)
        loss.backward()
        opt.step()
        opt.zero_grad(set_to_none=True)
    print(f"seed {seed} done. best_val={best_val:.4f}")


def jlens(seed, layers, n_prompts):
    out_dir = f'{JLENS_OUT}/seed{seed}'
    cmd = ["python3", "-u", "src/jlens_v3.py",
           "--checkpoint", f'{OUT_ROOT}/seed{seed}/ckpt.pt',
           "--data_dir", DATA_DIR,
           "--n_prompts", str(n_prompts),
           "--layers", layers,
           "--chunk", "16",
           "--output_dir", out_dir]
    print("running:", " ".join(cmd))
    r = subprocess.run(cmd, cwd=REPO_ROOT)
    assert r.returncode == 0, "jlens_v3 failed"


def summary(layers):
    with open(f'{DATA_DIR}/meta.pkl', 'rb') as f:
        meta = pickle.load(f)
    sid = meta['stoi']
    iid_s = sid[T_STRUCT]
    iid_n = sid[T_NOISE]
    train_data = np.memmap(f'{DATA_DIR}/train.bin', dtype=np.uint16, mode='r')
    V = meta['vocab_size']
    counts = np.bincount(train_data, minlength=V).astype(float)
    freq = counts / counts.sum() * 100

    seeds = sorted([d for d in os.listdir(JLENS_OUT) if d.startswith('seed')])
    print(f"\n{'='*72}")
    print("SYNTHETIC PAIR: faithful J-lens norms of T_struct '@' vs T_noise '#'")
    print(f"('@' freq={freq[iid_s]:.3f}%, '#' freq={freq[iid_n]:.3f}%)")
    print(f"{'='*72}")
    print(f"{'seed':<5}{'layer':<6}{'@ norm':>9}{'# norm':>9}{'ratio':>8}   freq-corr r")
    for s in seeds:
        for l in map(int, layers.split(',')):
            d = torch.load(f'{JLENS_OUT}/{s}/layer{l}.pt', map_location='cpu')
            fn = d['faithful_norms']
            f_arr = np.array([freq[k] for k in range(V)])
            rf = np.corrcoef(np.array([fn[k] for k in range(V)]), f_arr)[0, 1]
            print(f"{s:<5}{l:<6}{fn[iid_s]:>9.4f}{fn[iid_n]:>9.4f}"
                  f"{fn[iid_s]/max(fn[iid_n],1e-9):>8.2f}   {rf:+.3f}")
        # per-seed ratio across middle layers
        mids = [l for l in map(int, layers.split(',')) if l in (2, 3, 4)]
        rs = []
        rn = []
        for l in mids:
            d = torch.load(f'{JLENS_OUT}/{s}/layer{l}.pt', map_location='cpu')
            fn = d['faithful_norms']
            rs.append(fn[iid_s])
            rn.append(fn[iid_n])
        print(f"  -> mean middle-layer ratio @/# = {np.mean(rs)/max(np.mean(rn),1e-9):.3f}")


if __name__ == '__main__':
    ap = argparse.ArgumentParser()
    ap.add_argument('--step', required=True, choices=['prep', 'train', 'jlens', 'summary'])
    ap.add_argument('--seed', type=int, default=0)
    ap.add_argument('--max_iters', type=int, default=3000)
    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(clean=a.clean)
    elif a.step == 'train':
        train(a.seed, a.max_iters, a.batch_size)
    elif a.step == 'jlens':
        jlens(a.seed, a.layers, a.n_prompts)
    else:
        summary(a.layers)