90 lines
4.0 KiB
Python
90 lines
4.0 KiB
Python
|
import os, json, random
|
||
|
|
HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
|
|
import sys; sys.path.insert(0, os.path.join(HERE, "mlb"))
|
||
|
|
from tools_common import TOOLS
|
||
|
|
|
||
|
|
OUT = os.path.join(HERE, "data", "sft_tools.jsonl")
|
||
|
|
N = 1000
|
||
|
|
random.seed(0)
|
||
|
|
|
||
|
|
FILES = ["config.py", "app.py", "main.py", "README.md", "utils.py", "src/utils.py",
|
||
|
|
"tests/test_api.py", "server.js", "index.html", "data/train.csv", "Makefile",
|
||
|
|
"requirements.txt", "model.py", "train.py", ".env.example", "docker-compose.yml"]
|
||
|
|
DIRS = ["src", "tests", "data", ".", "configs", "scripts", "app", "lib", "models"]
|
||
|
|
CMDS = ["pytest", "git status", "ls -la", "npm test", "python train.py",
|
||
|
|
"pip install -r requirements.txt", "git log --oneline", "make build",
|
||
|
|
"docker build .", "ruff check .", "python -m pytest tests/", "git diff"]
|
||
|
|
EDITS = [("5000", "8080"), ("foo", "bar"), ("debug=False", "debug=True"),
|
||
|
|
("localhost", "0.0.0.0"), ("v1", "v2"), ("GET", "POST"),
|
||
|
|
("255", "128"), ("lr=1e-3", "lr=3e-4"), ("True", "False")]
|
||
|
|
WRITES = [("hello.py", "print('hi')"), ("notes.txt", "hello world"),
|
||
|
|
(".gitignore", "__pycache__/"), ("VERSION", "1.0.0"),
|
||
|
|
("run.sh", "#!/bin/bash\npython main.py"), ("todo.md", "- ship it")]
|
||
|
|
|
||
|
|
READ_T = ["Show me the contents of {p}", "Read {p}", "Open {p}", "What's in {p}?",
|
||
|
|
"Display {p}", "Print out {p} for me", "Can you show me {p}?"]
|
||
|
|
LIST_T = ["List the files in {p}", "What's in the {p} directory?", "Show the entries in {p}",
|
||
|
|
"List everything under {p}", "ls {p}", "Show me what's inside {p}"]
|
||
|
|
WRITE_T = ["Create a file {p} containing {c}", "Make a new file {p} with the text {c}",
|
||
|
|
"Write {c} to {p}", "Save {c} into a file called {p}"]
|
||
|
|
EDIT_T = ["In {p} replace {a} with {b}", "Change {a} to {b} in {p}",
|
||
|
|
"Update {p}: swap {a} for {b}", "In the file {p}, replace {a} with {b}",
|
||
|
|
"Edit {p} and change {a} into {b}", "Modify {p} so {a} becomes {b}"]
|
||
|
|
BASH_T = ["Run {c}", "Execute {c}", "Run the command {c}", "Please run {c}",
|
||
|
|
"Can you run {c}?", "Kick off {c}", "Go ahead and run {c}"]
|
||
|
|
|
||
|
|
WEIGHTS = [("edit_file", 0.28), ("run_bash", 0.28), ("read_file", 0.16),
|
||
|
|
("list_dir", 0.14), ("write_file", 0.14)]
|
||
|
|
|
||
|
|
def pick_tool():
|
||
|
|
r, acc = random.random(), 0.0
|
||
|
|
for name, w in WEIGHTS:
|
||
|
|
acc += w
|
||
|
|
if r <= acc:
|
||
|
|
return name
|
||
|
|
return WEIGHTS[-1][0]
|
||
|
|
|
||
|
|
def make_example():
|
||
|
|
t = pick_tool()
|
||
|
|
if t == "read_file":
|
||
|
|
p = random.choice(FILES)
|
||
|
|
q = random.choice(READ_T).format(p=p); args = {"path": p}
|
||
|
|
elif t == "list_dir":
|
||
|
|
p = random.choice(DIRS)
|
||
|
|
q = random.choice(LIST_T).format(p=p); args = {"path": p}
|
||
|
|
elif t == "write_file":
|
||
|
|
p, c = random.choice(WRITES)
|
||
|
|
q = random.choice(WRITE_T).format(p=p, c=c); args = {"path": p, "content": c}
|
||
|
|
elif t == "edit_file":
|
||
|
|
p = random.choice(FILES); a, b = random.choice(EDITS)
|
||
|
|
q = random.choice(EDIT_T).format(p=p, a=a, b=b)
|
||
|
|
args = {"path": p, "old_string": a, "new_string": b}
|
||
|
|
else: # run_bash
|
||
|
|
c = random.choice(CMDS)
|
||
|
|
q = random.choice(BASH_T).format(c=c); args = {"command": c}
|
||
|
|
call = json.dumps({"name": t, "arguments": args})
|
||
|
|
assistant = f"<tool_call>\n{call}\n</tool_call>"
|
||
|
|
return {"messages": [{"role": "user", "content": q},
|
||
|
|
{"role": "assistant", "content": assistant}],
|
||
|
|
"tools": TOOLS}
|
||
|
|
|
||
|
|
def main():
|
||
|
|
seen, rows = set(), []
|
||
|
|
while len(rows) < N:
|
||
|
|
ex = make_example()
|
||
|
|
key = ex["messages"][0]["content"] + ex["messages"][1]["content"]
|
||
|
|
if key in seen:
|
||
|
|
continue
|
||
|
|
seen.add(key); rows.append(ex)
|
||
|
|
with open(OUT, "w") as f:
|
||
|
|
for r in rows:
|
||
|
|
f.write(json.dumps(r) + "\n")
|
||
|
|
dist = {}
|
||
|
|
for r in rows:
|
||
|
|
name = json.loads(r["messages"][1]["content"].split("<tool_call>\n")[1].split("\n</tool_call>")[0])["name"]
|
||
|
|
dist[name] = dist.get(name, 0) + 1
|
||
|
|
print(f"wrote {len(rows)} -> {OUT}")
|
||
|
|
print("distribution:", dist)
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|