File size: 9,531 Bytes
2a3f489 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | """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)))
# 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()
|