"""STEP1 (MLX), scaled (STATUS 2026-06-10, next-step 1): 30 scenario triples instead of 5.
For each scenario, compress turn-1 into SP and tokenize a REFERENTIAL turn-2 (needs a turn-1
value, which therefore lives ONLY in the SP) and a matched CONTROL turn-2 (self-contained —
the value is in the question, so the SP isn't needed). Same SP for both. Categories cover
arithmetic chains, codes/IDs, names, schedule, preferences and measurements so the probe's
heads aren't an artifact of one task family. Save for attn_probe3.py.

Run next to sp_mlx.py in the HF repo:  python3 attn_export3.py
"""
import numpy as np, mlx.core as mx
import sp_mlx

M = sp_mlx.get()
tok, pooler, embT = M["tok"], M["pooler"], M["embT"]
emb = lambda ids: embT(mx.array([ids]))

from attn_scenarios import SCEN
out = {"n": np.array([len(SCEN)]),
       "bos": np.array([tok.bos_token_id if tok.bos_token_id is not None else tok.encode("")[0]]),
       "cat": np.array([c for c, *_ in SCEN])}
for i, (cat, t1, ref, ctrl) in enumerate(SCEN):
    sp = pooler.forward(emb(tok.encode(t1, add_special_tokens=False)).astype(mx.float32))
    out[f"sp_{i}"] = np.array(sp.astype(mx.float32))[0]
    out[f"ref_{i}"] = np.array(tok.encode(ref, add_special_tokens=False))
    out[f"ctrl_{i}"] = np.array(tok.encode(ctrl, add_special_tokens=False))
np.savez("attn_probe3.npz", **out)
print(f"saved {len(SCEN)} scenarios, SP shape {out['sp_0'].shape}")
print("ATTN_EXPORT3_DONE")