File size: 3,717 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
"""Path-diversify an SFT corpus.
99% of the converted trajectories live at /testbed, because that is where SWE-bench-style
container images put the repo. That is correct for swe-bench-verified (all 500 tasks) but it
would teach an unconditional "/testbed" reflex, which is wrong for terminal-bench-2 (whose
tasks live at /app and name their paths in the prompt).
For a fraction of rows this rewrites /testbed to another absolute root, consistently across
every message, tool argument and observation in the trajectory, AND states the location in the
user turn. Rows left alone keep /testbed with a bare prompt. The pair teaches:
prompt names a location -> work there
prompt names nothing -> the repo is at /testbed
Nothing is invented: only a path string is substituted, and the added sentence is true of the
rewritten trajectory.
"""
import argparse
import glob
import json
import os
import random
import re
ROOTS = [
"/app",
"/app/{repo}",
"/workspace",
"/workspace/{repo}",
"/srv/{repo}",
"/opt/{repo}",
"/home/user/{repo}",
"/code",
"/repo",
"/project/{repo}",
]
SAY = [
"The project is at {root}.",
"The repository you need to work on is checked out at {root}.",
"Work in {root}.",
"You will find the code in {root}.",
"The codebase lives at {root} — use absolute paths under it.",
]
REPO_WORDS = [
"src", "pkg", "core", "lib", "service", "backend", "engine", "toolkit",
"app", "server", "client", "runtime", "sdk", "platform",
]
def sub_all(obj, old, new):
if isinstance(obj, str):
return obj.replace(old, new)
if isinstance(obj, list):
return [sub_all(x, old, new) for x in obj]
if isinstance(obj, dict):
return {k: sub_all(v, old, new) for k, v in obj.items()}
return obj
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--src", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--p", type=float, default=0.30, help="fraction of rows to relocate")
ap.add_argument("--seed", type=int, default=17)
args = ap.parse_args()
rng = random.Random(args.seed)
import pyarrow.parquet as pq
from datasets import Dataset
rows = []
n_moved = 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 = row["messages"]
blob = json.dumps(msgs)
if "/testbed" in blob and rng.random() < args.p:
repo = f"{rng.choice(REPO_WORDS)}{rng.choice(['', '-' + rng.choice(REPO_WORDS), '_' + str(rng.randint(2, 99))])}"
root = rng.choice(ROOTS).format(repo=repo)
msgs = sub_all(msgs, "/testbed", root)
# state the location in the first user turn
for m in msgs:
if m["role"] == "user":
m["content"] = rng.choice(SAY).format(root=root) + "\n\n" + (m.get("content") or "")
break
n_moved += 1
row = {**row, "messages": msgs, "source": row["source"] + "+moved"}
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_moved} relocated) -> {args.out}")
if __name__ == "__main__":
main()
|