"""Build the data files for the v6 9B rollout-comparison Space. For 10 shared problems (sample 0) and three models (v6 SFT final / RL step 150 / RL step 250), all run with the v6 scaffold at depth 8 x width 6, write one JSON per (problem, model) with, per layer: the planner's thinking and public output (directions), every executed direction with E's public result, compactions at that layer (input size, absorbed slots, compacted memory, compactor thinking), and the frontier text the planner saw; plus the final answer and the judge verdict. Long private reasoning (E thinking is dropped; compactor and final-answer thinking are capped) keeps each file small enough to load in a browser. """ import glob, json, os, re SP = '/tmp/claude-18871/-home-aci18871mh-ian-meta-reason/f6ab329a-6c52-4319-bc08-b5751659e956/scratchpad' OUT = '/home/aci18871mh/ian/scratch/hf_space_v6compare/data' SRC = '/home/aci18871mh/ian/scratch/meta_reason_rl_v2/rc_inference/overnight_20260923/inputs' MODELS = { 'sft': dict(label='v6 SFT (midtrain final, iter 2981)', short='Midtrain (SFT)', roll=['/home/aci18871mh/ian/scratch/v6_tcseval_20261026_v5_recovery/final/rollouts/shard*.jsonl', '/home/aci18871mh/ian/scratch/v6_arxiv_imo_budget_20261026/d8_w6/rollouts/shard*.jsonl'], tcs='/home/aci18871mh/ian/scratch/v6_tcseval_20261026_v5_recovery/final/judgments/tcsbench_2026_eval.gpt-5.6-luna-high.jsonl', arxiv='/home/aci18871mh/ian/scratch/v6_arxiv_imo_budget_20261026/d8_w6/judgments/arxivmath.luna-medium.jsonl'), 's150': dict(label='RL step 150 (outcome-only run, iter_0000149)', short='RL step 150', roll=['/home/aci18871mh/ian/scratch/v6rl9b_eval_s150/d8_w6/merged/tcs.jsonl', '/home/aci18871mh/ian/scratch/v6rl9b_eval_s150/d8_w6/merged/arxiv.jsonl'], tcs='/home/aci18871mh/ian/scratch/v6rl9b_eval_s150/d8_w6/judgments/tcs.luna-high.jsonl', arxiv='/home/aci18871mh/ian/scratch/v6rl9b_eval_s150/d8_w6/judgments/arxiv.luna-medium.jsonl'), 's250': dict(label='RL step 250 (final, iter_0000249)', short='RL step 250 (final)', roll=['/home/aci18871mh/ian/scratch/v6rl9b_eval_s250/d8_w6/rollouts/shard*.jsonl'], tcs='/home/aci18871mh/ian/scratch/v6rl9b_eval_s250/d8_w6/judgments/tcs.luna-high.jsonl', arxiv='/home/aci18871mh/ian/scratch/v6rl9b_eval_s250/d8_w6/judgments/arxiv.luna-medium.jsonl'), } CAP_COMPACT_THINK = 30000 CAP_FA_THINK = 40000 CAP_E_PUBLIC = 20000 def split_think(raw): raw = raw or '' if '' in raw: t, p = raw.split('', 1) return t.replace('', '').strip(), p.strip() return raw.replace('', '').strip(), '' def cap(text, n): if len(text) <= n: return text, False head = int(n * 0.75) return text[:head] + f'\n\n[... {len(text) - n:,} characters omitted ...]\n\n' + text[-(n - head):], True def frontier_of(user): i = user.find('Exploration so far:') j = user.find('Exploration budget:') if i < 0: return '' return user[i + len('Exploration so far:'): j if j > i else None].strip() def budget_of(user): j = user.find('Exploration budget:') k = user.find('final-answer step.', j) return user[j: k + len('final-answer step.')].strip() if j >= 0 else '' def load_rolls(pats): out = {} for p in pats: for f in glob.glob(p): for line in open(f): try: r = json.loads(line) except Exception: continue out[(r['problem_id'], int(r['sample_idx']))] = r return out def load_judge(path, kind): out = {} for line in open(path): if not line.strip(): continue r = json.loads(line) k = (r['problem_id'], int(r.get('sample_idx', 0))) if kind == 'arxiv': out[k] = dict(score=1.0 if r.get('correct') else 0.0, verdict='correct' if r.get('correct') else 'incorrect', extracted=r.get('candidate') or r.get('extracted_answer') or r.get('final_extracted'), note=r.get('reason') or r.get('judge_reason') or '') else: if r.get('score') is None: continue mx = r.get('max_score') or 7 out[k] = dict(score=r['score'] / mx, verdict=f"{r['score']}/{mx}", note=(r.get('explanation') or r.get('reason') or r.get('judge_raw') or '')[:6000]) return out def build(r, judge): steps = r.get('steps') or [] comp_by_layer = {} for c in r.get('compactions') or []: th, pub = split_think(c.get('raw')) th, cut = cap(th, CAP_COMPACT_THINK) m = re.search(r'(.*?)', pub, re.S) comp_by_layer.setdefault(int(c['layer_idx']), []).append(dict( attempt=c.get('attempt'), absorbed=c.get('absorbed_slots'), error=c.get('error'), input_chars=len(c.get('user') or ''), memory=(m.group(1).strip() if m else pub), thinking=th, thinking_capped=cut, thinking_chars=len(split_think(c.get('raw'))[0]))) layers = {} for s in steps: L = int(s['layer_idx']) lay = layers.setdefault(L, dict(layer=L, mr=None, explorations=[], compactions=[])) if s['role'] == 'mr': th, pub = split_think(s.get('raw')) lay['mr'] = dict(thinking=th, public=pub, frontier=frontier_of(s.get('user') or ''), budget=budget_of(s.get('user') or ''), prompt_chars=len(s.get('user') or '')) elif s['role'] == 'e': th, pub = split_think(s.get('raw')) pub, cut = cap(pub, CAP_E_PUBLIC) lay['explorations'].append(dict(idx=s.get('direction_idx'), direction=s.get('direction') or '', ok=s.get('summary_ok'), result=pub, result_capped=cut, thinking_chars=len(th))) # A compaction with layer_idx k commits completed layer k+1 (0-based k) into the 12-slot window, i.e. it runs # while building the frontier for planner turn k+2 -- or for the final answer after the last layer. before_final = [] for k, cs in comp_by_layer.items(): if (k + 2) in layers and layers[k + 2]['mr'] is not None: layers[k + 2]['compactions'] = cs else: before_final += cs final = dict(compactions=before_final) if before_final else None for d, f in sorted((r.get('finals') or {}).items(), key=lambda x: int(x[0])): th = f.get('final_thinking') or split_think(f.get('final_raw'))[0] th, cut = cap(th, CAP_FA_THINK) final = dict(depth=int(d), answer=f.get('final_answer') or split_think(f.get('final_raw'))[1], thinking=th, thinking_capped=cut, format_fault=f.get('format_fault'), compactions=before_final) widths = {} for L, lay in layers.items(): if lay['explorations']: widths[L] = len(lay['explorations']) return dict(problem_id=r['problem_id'], sample_idx=r['sample_idx'], termination=r['termination_reason'], n_layers=r['n_layers_completed'], budget=r.get('exploration_budget'), wall_time_s=r.get('wall_time_s'), terminate_rationale=r.get('terminate_rationale') or '', judge=judge, n_compactions=sum(len(v) for v in comp_by_layer.values()), widths=widths, layers=[layers[k] for k in sorted(layers)], final=final) def main(): pick = [tuple(x) for x in json.load(open(SP + '/pick.json'))] gold = {} for line in open(f'{SRC}/arxivmath-0626.source.jsonl'): g = json.loads(line); gold[g['problem_id']] = g.get('answer') or g.get('gold_answer') or g.get('solution') index = dict(models=[dict(key=k, label=v['label'], short=v['short']) for k, v in MODELS.items()], problems=[]) loaded = {k: (load_rolls(v['roll']), {**load_judge(v['tcs'], 'tcs'), **load_judge(v['arxiv'], 'arxiv')}) for k, v in MODELS.items()} for n, key in enumerate(pick): pid = key[0] bench = 'TCSBench-2026' if pid.startswith('tcsbench') else 'ArxivMath-0626' entry = dict(id=f'p{n:02d}', problem_id=pid, bench=bench, problem=loaded['sft'][0][key]['problem'], reference=gold.get(pid) if bench.startswith('Arxiv') else None, models={}) for mk in MODELS: rolls, judges = loaded[mk] data = build(rolls[key], judges.get(key)) fn = f'{entry["id"]}__{mk}.json' json.dump(data, open(f'{OUT}/{fn}', 'w'), ensure_ascii=False) entry['models'][mk] = dict(file=fn, score=(data['judge'] or {}).get('score'), verdict=(data['judge'] or {}).get('verdict'), n_layers=data['n_layers'], termination=data['termination'], n_compactions=data['n_compactions'], mean_width=round(sum(data['widths'].values()) / max(1, len(data['widths'])), 2)) index['problems'].append(entry) print(entry['id'], pid, {m: (v['verdict'], v['n_layers'], v['mean_width'], v['n_compactions']) for m, v in entry['models'].items()}) json.dump(index, open(f'{OUT}/index.json', 'w'), ensure_ascii=False, indent=1) sizes = sorted(os.path.getsize(f'{OUT}/{f}') for f in os.listdir(OUT)) print('files', len(sizes), 'total MB', round(sum(sizes) / 1e6, 1), 'largest MB', round(sizes[-1] / 1e6, 2)) if __name__ == '__main__': main()