latent_backtrack / scripts /inspect_stage0.py
Avra98's picture
Add training code (same as GitHub reasoning-by-superposition-latent)
8f46582 verified
Raw
History Blame Contribute Delete
685 Bytes
import json
from dataset import expand_data
from stokenizer import StepTokenizer
tok = StepTokenizer()
for L in (10, 12, 14, 16):
d = json.load(open(f"data/star_2arm_L{L}_valid_fo_bfs.json"))
s = d[0]
print(f"===== L{L} =====")
print("keys:", list(s.keys()))
print("steps len:", len(s.get("steps", [])),
"target:", s.get("target"),
"root:", s.get("root"),
"neg:", s.get("neg_target"))
q, c = expand_data(s, 1, len(s["steps"]), neg_sampling=False)
print("stage0 Q:", q[:300])
print("stage0 cont:", repr(c))
ids = tok.encode(q, add_special_tokens=False)
print("q tokens:", len(ids), "max_id:", max(ids))
print()