| """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 '</think>' in raw: |
| t, p = raw.split('</think>', 1) |
| return t.replace('<think>', '').strip(), p.strip() |
| return raw.replace('<think>', '').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'<compacted_history>(.*?)</compacted_history>', 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))) |
| |
| |
| 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() |
|
|
