program-probes/run_extract_swebench.py at main · ASSERT-KTH/program-probes · GitHub

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
"""Extract hidden states from SWE-bench agent trajectories for probe training.

Each trajectory is a single linearized token sequence.  An extraction_mask marks
assistant-turn tokens; hidden states are extracted at masked positions strided by
gen_config.stride.  One .pt file is written per trajectory (activations only — no
labels).  Use run_attach_labels_swebench.py (CPU-only) to attach probe labels.

Example usage:
    uv run python run_extract_swebench.py \
        --model-config configs/models/qwen36_27b.yaml \
        --generation-config configs/generation.yaml \
        --traj-dir generations/swebench/qwen36_27b_test \
        --output-dir outputs/swebench \
        --shard-rank 0 --num-shards 1
"""
from __future__ import annotations

import argparse
import random
from pathlib import Path

import numpy as np
import torch

from src.configs import GenerationConfig, ModelConfig, load_config
from src.tasks.swe_bench_extract import load_trajectories


def _set_seeds(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def _load_model_adapter(adapter_name: str):
    if adapter_name == "qwen":
        from src.models.qwen import QwenAdapter
        return QwenAdapter()
    if adapter_name == "qwen35":
        from src.models.qwen35 import Qwen35Adapter
        return Qwen35Adapter()
    if adapter_name == "cwm":
        from src.models.cwm import CwmAdapter
        return CwmAdapter()
    if adapter_name == "laguna":
        from src.models.laguna import LagunaAdapter
        return LagunaAdapter()
    raise ValueError(f"Unknown model adapter: {adapter_name!r}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-config", required=True)
    parser.add_argument("--generation-config", required=True)
    parser.add_argument("--traj-dir", required=True, help="Directory of trajectory JSON files")
    parser.add_argument("--output-dir", default="outputs/swebench")
    parser.add_argument("--shard-rank", type=int, default=0)
    parser.add_argument("--num-shards", type=int, default=1)
    parser.add_argument("--extraction-batch-size", type=int, default=None,
                        help="Override gen_config.extraction_batch_size")
    parser.add_argument("--chunk-size", type=int, default=None,
                        help="Process sequences in chunks of this many tokens using KV cache "
                             "to bound peak GPU memory. If unset, uses a single full-sequence forward pass.")
    args = parser.parse_args()

    model_config: ModelConfig = load_config(args.model_config, ModelConfig)
    gen_config: GenerationConfig = load_config(args.generation_config, GenerationConfig)
    _set_seeds(gen_config.seed)

    print(f"Loading trajectories from {args.traj_dir}...")
    all_trajs = load_trajectories(args.traj_dir)
    trajs = all_trajs[args.shard_rank::args.num_shards]
    print(f"  {len(trajs)} trajectories (shard {args.shard_rank}/{args.num_shards})")

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    model_adapter = _load_model_adapter(model_config.adapter)
    model_adapter.load_for_extraction(model_config, gen_config)

    layer_indices = model_config.probe_layers
    stride = gen_config.stride
    batch_size = args.extraction_batch_size if args.extraction_batch_size is not None else gen_config.extraction_batch_size

    for batch_start in range(0, len(trajs), batch_size):
        batch = trajs[batch_start: batch_start + batch_size]

        pending = []
        for traj in batch:
            fname = out_dir / f"{traj.sample_id}.pt"
            if not fname.exists():
                pending.append((traj, fname))
        if not pending:
            continue

        print(f"  Extracting batch {batch_start}–{batch_start + len(pending) - 1}...")
        for traj, fname in pending:
            try:
                per_seq_hs = model_adapter.extract_hidden_states(
                    [traj.token_ids], [0], layer_indices, stride,
                    extraction_masks=[traj.extraction_mask],
                    chunk_size=args.chunk_size,
                )
            except torch.cuda.OutOfMemoryError:
                torch.cuda.empty_cache()
                print(f"    OOM — skipping {fname.name} (n_tokens={len(traj.token_ids)})")
                continue

            hs = per_seq_hs[0]
            n_steps = min(len(v) for v in hs.values()) if hs else 0

            out = {
                "activations": hs,
                "instance_id": traj.instance_id,
                "sample_id": traj.sample_id,
                "group_id": traj.instance_id,
                "outcome": traj.outcome,
                "n_captured_steps": n_steps,
                "n_tokens": len(traj.token_ids),
                "n_turns": len(traj.step_segment_indices),
            }
            torch.save(out, fname)
            print(f"    Saved {fname.name}  (n_steps={n_steps}, outcome={traj.outcome})")


if __name__ == "__main__":
    main()