File size: 3,479 Bytes
60b165e | 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 | #!/usr/bin/env python3
"""
Generate a tiny SYNTHETIC tool-selection dataset in DiffusionGemma format, so the
trainer/eval in this repo can be smoke-tested end-to-end WITHOUT any private data.
The real adapter was trained on private agent traces (not included). This produces
fully synthetic prompt/response pairs with the same structure: a system prompt, a
candidate tool list + a task in the user turn, the thinking-channel generation
prefill, and a dash-prefixed tool-name response ending in <turn|>.
python3 make_example_data.py --out ./data
"""
import argparse, hashlib, json, random
from pathlib import Path
TOOLS = ["Bash", "Read", "Edit", "Write", "Grep", "Glob", "WebFetch", "WebSearch",
"Agent", "TodoWrite", "NotebookEdit", "Task"]
SYSTEM = ("You are a tool selector. Given a task and a list of available tools, "
"select ONLY the tools needed. Output one tool per line with a dash prefix.")
GEN_PREFILL = "<|turn>model\n<|channel>thought\n<channel|>"
# (task template, the tools it should select) — deterministic synthetic mapping
TASKS = [
("Read the config file at {path} and print its contents", ["Read"]),
("Find every TODO comment under {path} and list them", ["Grep", "Read"]),
("Fix the failing test in {path} — locate the bug and patch it", ["Read", "Edit", "Bash"]),
("Create a new module {path} with a hello function", ["Write"]),
("Search the web for the latest {topic} release notes", ["WebSearch", "WebFetch"]),
("Run the test suite and report failures", ["Bash"]),
("Rename the symbol {topic} across all files under {path}", ["Grep", "Edit"]),
("Summarize the open issues, then draft a plan", ["WebFetch", "TodoWrite"]),
("List all python files and count lines of code", ["Glob", "Bash"]),
("Delegate a deep research task about {topic}", ["Agent"]),
]
PATHS = ["src/parser.py", "lib/config.ts", "tests/test_api.py", "core/", "app/main.rs"]
TOPICS = ["MLX", "DiffusionGemma", "Rust async", "Postgres indexing", "WebGPU"]
def render(rng):
template, tools = rng.choice(TASKS)
task = template.format(path=rng.choice(PATHS), topic=rng.choice(TOPICS))
# shuffle a candidate list that always includes the correct tools + distractors
cands = list(set(tools) | set(rng.sample(TOOLS, k=rng.randint(6, 10))))
rng.shuffle(cands)
prompt = (f"<|turn>system\n{SYSTEM} <turn|>\n"
f"<|turn>user\nAvailable tools: {', '.join(cands)}\n\n"
f"Task: {task}\n\nSelect the tools needed:<turn|>\n{GEN_PREFILL}")
response = "".join(f"- {t}\n" for t in tools).rstrip("\n") + "<turn|>"
return {"prompt": prompt, "response": response}
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--out", default="./data")
ap.add_argument("--seed", type=int, default=7)
args = ap.parse_args()
out = Path(args.out); out.mkdir(parents=True, exist_ok=True)
rng = random.Random(args.seed)
for split, n in (("train", 120), ("valid", 24), ("test", 24)):
rows = [render(rng) for _ in range(n)]
f = out / f"{split}.jsonl"
with open(f, "w") as fh:
for r in rows:
fh.write(json.dumps(r, ensure_ascii=False) + "\n")
print(f"wrote {f} ({n} synthetic examples)")
print("\nNOTE: synthetic toy data for smoke-testing the pipeline only — not the "
"real training corpus. Expect the model to overfit this tiny set quickly.")
if __name__ == "__main__":
main()
|