Fild/diffusiongemma-26B-A4B-it-tool-selector-lora-mlx

1#!/usr/bin/env python32"""3Generate a tiny SYNTHETIC tool-selection dataset in DiffusionGemma format, so the4trainer/eval in this repo can be smoke-tested end-to-end WITHOUT any private data.5 6The real adapter was trained on private agent traces (not included). This produces7fully synthetic prompt/response pairs with the same structure: a system prompt, a8candidate tool list + a task in the user turn, the thinking-channel generation9prefill, and a dash-prefixed tool-name response ending in <turn|>.10 11  python3 make_example_data.py --out ./data12"""13import argparse, hashlib, json, random14from pathlib import Path15 16TOOLS = ["Bash", "Read", "Edit", "Write", "Grep", "Glob", "WebFetch", "WebSearch",17         "Agent", "TodoWrite", "NotebookEdit", "Task"]18SYSTEM = ("You are a tool selector. Given a task and a list of available tools, "19          "select ONLY the tools needed. Output one tool per line with a dash prefix.")20GEN_PREFILL = "<|turn>model\n<|channel>thought\n<channel|>"21 22# (task template, the tools it should select) — deterministic synthetic mapping23TASKS = [24    ("Read the config file at {path} and print its contents", ["Read"]),25    ("Find every TODO comment under {path} and list them", ["Grep", "Read"]),26    ("Fix the failing test in {path} — locate the bug and patch it", ["Read", "Edit", "Bash"]),27    ("Create a new module {path} with a hello function", ["Write"]),28    ("Search the web for the latest {topic} release notes", ["WebSearch", "WebFetch"]),29    ("Run the test suite and report failures", ["Bash"]),30    ("Rename the symbol {topic} across all files under {path}", ["Grep", "Edit"]),31    ("Summarize the open issues, then draft a plan", ["WebFetch", "TodoWrite"]),32    ("List all python files and count lines of code", ["Glob", "Bash"]),33    ("Delegate a deep research task about {topic}", ["Agent"]),34]35PATHS = ["src/parser.py", "lib/config.ts", "tests/test_api.py", "core/", "app/main.rs"]36TOPICS = ["MLX", "DiffusionGemma", "Rust async", "Postgres indexing", "WebGPU"]37 38 39def render(rng):40    template, tools = rng.choice(TASKS)41    task = template.format(path=rng.choice(PATHS), topic=rng.choice(TOPICS))42    # shuffle a candidate list that always includes the correct tools + distractors43    cands = list(set(tools) | set(rng.sample(TOOLS, k=rng.randint(6, 10))))44    rng.shuffle(cands)45    prompt = (f"<|turn>system\n{SYSTEM} <turn|>\n"46              f"<|turn>user\nAvailable tools: {', '.join(cands)}\n\n"47              f"Task: {task}\n\nSelect the tools needed:<turn|>\n{GEN_PREFILL}")48    response = "".join(f"- {t}\n" for t in tools).rstrip("\n") + "<turn|>"49    return {"prompt": prompt, "response": response}50 51 52def main():53    ap = argparse.ArgumentParser()54    ap.add_argument("--out", default="./data")55    ap.add_argument("--seed", type=int, default=7)56    args = ap.parse_args()57    out = Path(args.out); out.mkdir(parents=True, exist_ok=True)58    rng = random.Random(args.seed)59    for split, n in (("train", 120), ("valid", 24), ("test", 24)):60        rows = [render(rng) for _ in range(n)]61        f = out / f"{split}.jsonl"62        with open(f, "w") as fh:63            for r in rows:64                fh.write(json.dumps(r, ensure_ascii=False) + "\n")65        print(f"wrote {f} ({n} synthetic examples)")66    print("\nNOTE: synthetic toy data for smoke-testing the pipeline only — not the "67          "real training corpus. Expect the model to overfit this tiny set quickly.")68 69 70if __name__ == "__main__":71    main()72