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()