checkpoint数据与划分

"""Build the per-checkpoint table for the checkpoint line.

One row per graded decision checkpoint (non-test diff changed), plus trajectory
totals. Features use only the visible prefix up to that step; labels are
`now` (official grade of the current patch) and `final` (trajectory's own grade).
"""
import argparse
import json
import re
from pathlib import Path
from .common import ROOT, read_json, write_json

TEST_CMD = re.compile(r'pytest|runtests\.py|unittest|bin/test|tox\b|sympy\.test|\btest\b')


def cum_output_tokens(spans):
    cum, out = 0, {}
    for step, a, b in spans:
        cum += b - a; out[step] = cum
    return out


def cum_input_tokens(spans):
    """Context length re-read at each call (prompt tokens before generation), cumulative."""
    cum, out = 0, {}
    for step, a, b in spans:
        cum += a; out[step] = cum
    return out


def step_messages(messages):
    """[(assistant_msg, [tool msgs])] indexed by step."""
    steps = []
    for m in messages:
        if m['role'] == 'assistant': steps.append((m, []))
        elif m['role'] == 'tool' and steps: steps[-1][1].append(m)
    return steps


def command_of(m):
    for tc in m.get('tool_calls') or []:
        args = tc.get('function', {}).get('arguments')
        if isinstance(args, str):
            try: args = json.loads(args)
            except Exception: args = {'command': args}
        if isinstance(args, dict) and args.get('command'): return args['command']
    for a in (m.get('extra') or {}).get('actions') or []:
        if a.get('command'): return a['command']
    return ''


def returncode_of(tool_msgs):
    for t in tool_msgs:
        c = t.get('content') or ''
        mm = re.search(r'"returncode":\s*(-?\d+)', c) or re.search(r'<returncode>(-?\d+)</returncode>', c)
        if mm: return int(mm.group(1))
    return None


def patch_stats(patch):
    add = sum(1 for l in patch.splitlines() if l.startswith('+') and not l.startswith('+++'))
    rem = sum(1 for l in patch.splitlines() if l.startswith('-') and not l.startswith('---'))
    return add, rem


def rows_for(label_file, split):
    """Decision rows: every step from the first graded production checkpoint on.

    At step t the submit-now patch is the latest graded checkpoint <= t (later
    steps before the next production change only touch tests/scratch files).
    """
    lab = read_json(label_file); tp = Path(lab['trial_path'])
    res = read_json(tp / 'result.json'); md = res['agent_result']['metadata']
    total_out = res['agent_result'].get('n_output_tokens') or 0
    cum = cum_output_tokens(md.get('tito_turn_spans') or [])
    cin = cum_input_tokens(md.get('tito_turn_spans') or []); in_total = max(cin.values()) if cin else 0
    steps = step_messages(read_json(tp / 'agent/mini-swe-agent.trajectory.json')['messages'])
    ck = [c for c in lab['checkpoints'] if c['reward'] is not None and c['prod_paths']]
    if len(ck) != sum(1 for c in lab['checkpoints'] if c['prod_paths']): return []  # ungraded checkpoint: drop trajectory
    rows = []; ci = -1; n_tests = n_pass = 0; tests_since = pass_since = 0; last_rc = None
    for t in range(len(steps)):
        cmd = command_of(steps[t][0]); rc = returncode_of(steps[t][1])
        is_test = bool(TEST_CMD.search(cmd))
        new_ck = ci + 1 < len(ck) and ck[ci + 1]['step'] == t
        if new_ck:
            ci += 1; tests_since = pass_since = 0
        if is_test:
            n_tests += 1; n_pass += rc == 0; last_rc = rc
            if not new_ck: tests_since += 1; pass_since += rc == 0
        if ci < 0: continue
        c = ck[ci]
        add, rem = patch_stats((tp / 'agent/checkpoints' / c['patch_file']).read_text()) if new_ck or not rows else (rows[-1]['lines_added'], rows[-1]['lines_removed'])
        rows.append({'task_id': lab['task_id'], 'repo': lab['task_id'].split('__')[0], 'split': split,
                     'rollout': lab['rollout'], 'traj': label_file.stem, 'step': t, 'n_steps': len(steps),
                     'ckpt_index': ci + 1, 'ckpt_step': c['step'], 'is_ckpt_step': new_ck, 'steps_since_ckpt': t - c['step'],
                     'prod_files': len(c['prod_paths']), 'lines_added': add, 'lines_removed': rem,
                     'out_tokens_at': cum.get(t, 0), 'out_tokens_total': total_out,
                     'in_tokens_at': cin.get(t, 0), 'in_tokens_total': in_total,
                     'is_test_cmd': is_test, 'rc': rc, 'n_tests_run': n_tests, 'n_tests_pass': n_pass, 'last_test_rc': last_rc,
                     'tests_since_ckpt': tests_since, 'passes_since_ckpt': pass_since,
                     'now': int(c['reward']), 'final': int(lab['trajectory_reward']), 'exit_status': md.get('exit_status'),
                     'last_ckpt_now': int(ck[-1]['reward']), 'trial_path': str(tp), 'patch_file': c['patch_file']})
    return rows


def main(a):
    rows = []; trajs = []
    splits = read_json(ROOT / 'data/splits.json')
    which = {t: s for s in ('train', 'dev', 'test') for t in splits[s]}
    for run in a.run:
        for f in sorted((ROOT / 'runs' / run / 'checkpoint-labels').glob('*.json')):
            split = which[read_json(f)['task_id']]
            if split not in a.split: continue
            r = rows_for(f, split); rows += r; lab = read_json(f)
            res = read_json(Path(lab['trial_path']) / 'result.json')
            ungraded = any(c['reward'] is None for c in lab['checkpoints'] if c['prod_paths'])
            trajs.append({'traj': f.stem, 'task_id': lab['task_id'], 'split': split, 'final': int(lab['trajectory_reward']),
                          'out_tokens_total': res['agent_result'].get('n_output_tokens') or 0, 'n_rows': len(r),
                          'in_tokens_total': r[0]['in_tokens_total'] if r else sum(a for _, a, _ in res['agent_result']['metadata'].get('tito_turn_spans') or []),
                          'dropped_ungraded': ungraded,
                          'last_ckpt_now': r[-1]['last_ckpt_now'] if r else None})
    out = Path(a.out); out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(''.join(json.dumps(r) + '\n' for r in rows))
    out.with_suffix('.trajs.jsonl').write_text(''.join(json.dumps(r) + '\n' for r in trajs))
    print(len(rows), 'rows', len(trajs), 'trajectories ->', out)


if __name__ == '__main__':
    ap = argparse.ArgumentParser()
    ap.add_argument('--run', nargs='+', required=True)
    ap.add_argument('--split', nargs='+', default=['train', 'dev'])
    ap.add_argument('--out', required=True)
    main(ap.parse_args())