File size: 4,816 Bytes
b2ebc95 | 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 | #!/usr/bin/env python3
"""Prepend a real orientation turn to a fraction of trajectories.
The eval harness starts the agent in an empty scratch dir (`/workspace`) with the project
somewhere else; the raw corpora all start already inside the repo, so a model trained on them
just writes files wherever it landed. This inserts, before the first assistant turn:
assistant: <short think> + bash "pwd; ls -a; echo ---; ls -d /testbed /app /workspace ..."
tool: <the real output, captured from actual task containers>
The observation is not invented — `assets/orient/*.txt` are verbatim captures from three real
SWE task images through the same broker. Run this BEFORE diversify_paths.py so the path
rewrite carries through the observation too, and the model sees roots other than /testbed.
"""
import argparse
import glob
import json
import os
import random
THINKS = [
"First, work out where I am and where the project actually is — the shell does not start inside the repo.",
"Before touching anything I should check the working directory and locate the project tree.",
"Let me orient myself: print the cwd and look for the project root.",
"I don't know which directory this shell starts in. Check that first, then find the codebase.",
"Start by locating the repository; the current directory is probably not it.",
]
PROBE = "pwd; ls -a; echo ---; ls -d /testbed /app /workspace /repo /code /srv 2>/dev/null"
FOLLOWUPS = [
"The project is at /testbed. I'll use absolute paths under it from here on.",
"Right — the repo is /testbed, and the cwd is an empty scratch dir. Work in /testbed.",
"/testbed holds the checkout. Every path from here on is absolute under /testbed.",
]
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--src", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--obs-dir", default="assets/orient")
ap.add_argument("--p", type=float, default=0.40)
ap.add_argument("--seed", type=int, default=23)
args = ap.parse_args()
rng = random.Random(args.seed)
observations = []
for f in sorted(glob.glob(args.obs_dir + "/*.txt")):
text = open(f).read().strip()
if text:
observations.append(text)
if not observations:
raise SystemExit(f"no captured observations in {args.obs_dir}")
print(f"{len(observations)} captured observations")
import pyarrow.parquet as pq
from datasets import Dataset
rows, n_aug = [], 0
for f in sorted(glob.glob(args.src + "/*.parquet")):
for batch in pq.ParquetFile(f).iter_batches(batch_size=256):
for row in batch.to_pylist():
msgs = list(row["messages"])
first_assist = next((i for i, m in enumerate(msgs) if m["role"] == "assistant"), None)
uses_testbed = "/testbed" in json.dumps(msgs)
if first_assist is not None and uses_testbed and rng.random() < args.p:
call_id = "chatcmpl-tool-%016x" % rng.getrandbits(56)
probe = [
{
"role": "assistant",
"reasoning_content": rng.choice(THINKS),
"content": "",
"tool_calls": [
{
"id": call_id,
"type": "function",
"function": {"name": "bash", "arguments": json.dumps({"command": PROBE})},
}
],
},
{"role": "tool", "tool_call_id": call_id, "name": "bash",
"content": rng.choice(observations)},
]
# fold the conclusion into the reasoning of the original first turn
head = dict(msgs[first_assist])
head["reasoning_content"] = (
rng.choice(FOLLOWUPS) + " " + (head.get("reasoning_content") or "")
).strip()
msgs = msgs[:first_assist] + probe + [head] + msgs[first_assist + 1:]
n_aug += 1
row = {**row, "messages": msgs, "n_assistant": row["n_assistant"] + 1}
rows.append(row)
rng.shuffle(rows)
os.makedirs(args.out, exist_ok=True)
ds = Dataset.from_list(rows)
shards = max(1, len(rows) // 20000 + 1)
for i in range(shards):
ds.shard(num_shards=shards, index=i).to_parquet(
os.path.join(args.out, f"train-{i:05d}-of-{shards:05d}.parquet")
)
print(f"wrote {len(rows)} rows ({n_aug} with an orientation turn) -> {args.out}")
if __name__ == "__main__":
main()
|