"""Teacher trajectories: a strong local model (Qwen3.8-27B via vLLM's OpenAI API) acts as the agent in our sandbox, with our tools and tool-call syntax. Only episodes that are correct AND grounded are kept, saved as canonical messages (tiny_agent.chat) with the student's system prompt. Qwen tends to overthink (thinking can run to tens of thousands of tokens), which would blow both the time budget and the student's 4k context. By default thinking is switched off in the chat template (enable_thinking=false) and the model is told to write one or two sentences of reasoning before each call instead; that text becomes the turn's block in the student data. Calls that hit max_tokens are counted ("cut"), and the log reports how many episodes survive the same filter the data build uses (scripts.gen_synth.teacher_reject). Plain chat in/out, but in Qwen's native tool format: the system prompt carries the same "# Tools" preamble its chat template renders, Qwen writes XML (asked for our JSON instead, it blends the two and ~80% of episodes broke), and the calls are converted to our canonical JSON calls here. The student's parser (tiny_agent.chat) is unchanged. Tool results go back as blocks, which is how the template renders tool turns. $TA_PY scripts/teacher.py --url http://127.0.0.1:8000/v1 --model qwen38 --hours 6 --out $TA_DATA/teacher/run1.jsonl """ import argparse import json import os import random import re import threading import time import urllib.request from collections import Counter from concurrent.futures import ThreadPoolExecutor from scripts.gen_synth import teacher_reject from tiny_agent.chat import SYSTEM_TA_V1, parse_assistant from tiny_agent.tasks import TRAIN_KINDS, check, grounded, make_task from tiny_agent.tools import TOOL_SPECS, Workspace TEACHER_SEED0 = 5_000_000 # eval 0-99,999 | RL 100,000-999,999 | scripted 1,000,000+ | teacher 5,000,000+ TEACHER_RULES = """ Results come back in blocks, in the same order. Paths are relative to /work. Never guess: every fact in your answer must come from a tool result, and compute results by running code (e.g. python3 -c), not in your head. Finish by calling submit with only the answer (a number, name, path, DONE or NOT_FOUND), no extra words.""" BRIEF_RULE = """ Before each tool call, write one or two short sentences of plain reasoning: what you will check and why. Keep it brief.""" _PARAMS = {"bash": {"command": "string"}, "read": {"path": "string", "offset": "integer", "limit": "integer"}, "edit": {"path": "string", "old_string": "string", "new_string": "string", "replace_all": "boolean"}, "write": {"path": "string", "content": "string"}, "submit": {"answer": "string"}} _REQUIRED = {"bash": ["command"], "read": ["path"], "edit": ["path", "old_string", "new_string"], "write": ["path", "content"], "submit": ["answer"]} QWEN_TOOLS = [{"type": "function", "function": { "name": n, "description": d, "parameters": {"type": "object", "properties": {k: {"type": t} for k, t in _PARAMS[n].items()}, "required": _REQUIRED[n]}}} for n, _, d in TOOL_SPECS] # what Qwen3.8's chat template renders when tools are passed (copied from chat_template.jinja) QWEN_TOOL_PREAMBLE = ( "# Tools\n\nYou have access to the following functions:\n\n" + "".join("\n" + json.dumps(t, ensure_ascii=False) for t in QWEN_TOOLS) + "\n" "\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n\n" "\n\nvalue_1\n\n" "\nThis is the value for the second parameter\nthat can span\nmultiple lines\n" "\n\n\n\n\nReminder:\n- Function calls MUST follow the " "specified format: an inner block must be nested within " "XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function " "call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, " "answer the question like normal with your current knowledge and do not tell the user about function calls\n" "") _NCALL = re.compile(r"(.*?)(?:|$)", re.S) _FUNC = re.compile(r"\s]+)>") _PARAM = re.compile(r"\s]+)>\n?(.*?)\n?", re.S) def _coerce(name, key, v): t = _PARAMS.get(name, {}).get(key, "string") if t == "integer" and re.fullmatch(r"\s*-?\d+\s*", v): return int(v) if t == "boolean" and v.strip().lower() in ("true", "false"): return v.strip().lower() == "true" return v def native_to_json(text): """Rewrite Qwen's native calls as our {json}, with parameter types from the tool schema. Anything else is left for parse_assistant to judge.""" def one(m): body = m.group(1) f = _FUNC.search(body) if not f: return m.group(0) name = f.group(1) args = {k: _coerce(name, k, v) for k, v in _PARAM.findall(body)} return "\n" + json.dumps({"name": name, "arguments": args}, ensure_ascii=False) + "\n" return _NCALL.sub(one, text.split("<|im_end|>")[0]) def chat(url, model, messages, max_tokens, temperature, top_p=0.8, top_k=20, thinking=False, timeout=600): body = json.dumps({"model": model, "messages": messages, "max_tokens": max_tokens, "temperature": temperature, "top_p": top_p, "top_k": top_k, "chat_template_kwargs": {"enable_thinking": bool(thinking)}}).encode() req = urllib.request.Request(f"{url}/chat/completions", data=body, headers={"Content-Type": "application/json"}) with urllib.request.urlopen(req, timeout=timeout) as r: out = json.loads(r.read()) choice = out["choices"][0] msg = choice["message"] text = msg.get("content") or "" reasoning = msg.get("reasoning_content") or msg.get("reasoning") if reasoning and "" not in text: text = f"\n{reasoning.strip()}\n\n{text}" if "" in text and "" not in text: # some templates open in the prompt text = "\n" + text return text, out.get("usage", {}), choice.get("finish_reason") == "length" _OVR = {"mtime": None, "args": None} def live_args(a): """Settings can be changed while running (no restart, which would end the teacher window): write e.g. {"max_tokens": 512, "thinking": 0, "temperature": 0.6} to .overrides.json.""" path = a.out + ".overrides.json" try: mt = os.path.getmtime(path) except OSError: return a if mt != _OVR["mtime"]: try: ovr = json.load(open(path)) b = argparse.Namespace(**vars(a)) for k, v in ovr.items(): if k in ("max_tokens", "thinking", "temperature", "top_p", "top_k", "max_turns"): setattr(b, k, type(getattr(a, k))(v)) _OVR.update(mtime=mt, args=b) print("overrides", ovr, flush=True) except (ValueError, OSError, TypeError) as e: print("bad overrides file:", e, flush=True) return a return _OVR["args"] def episode(seed, kind, a): task = make_task(random.Random(seed), kind) canon = task.messages() rules = TEACHER_RULES + ("" if a.thinking else BRIEF_RULE) system = QWEN_TOOL_PREAMBLE + "\n\n" + SYSTEM_TA_V1.split("\nTools:")[0] + rules conv = [{"role": "system", "content": system}, {"role": "user", "content": task.question}] toks = cut = n_calls = 0 with Workspace(task.files) as ws: for _ in range(a.max_turns): text, usage, was_cut = chat(a.url, a.model, conv, a.max_tokens, a.temperature, a.top_p, a.top_k, a.thinking) toks += usage.get("completion_tokens", 0) cut += was_cut n_calls += 1 msg = parse_assistant(native_to_json(text)) if not msg["think"] and msg["content"]: # the brief reasoning written before the calls is the student's block msg["think"], msg["content"] = msg["content"], "" conv.append({"role": "assistant", "content": text}) canon.append(msg) calls = msg["tool_calls"] if not calls: break results = [] for c in calls: results.append(c["error"] if "error" in c else ws.call(c["name"], c["arguments"])) if ws.submitted is not None: break canon.append({"role": "tool", "results": results}) if ws.submitted is not None: break conv.append({"role": "user", "content": "".join(f"\n{r}\n\n" for r in results)}) ok = check(task, ws.submitted, ws) return dict(seed=seed, kind=kind, ok=ok, grounded=grounded(task, canon), submitted=ws.submitted, answer=task.answer, completion_tokens=toks, calls=n_calls, cut=cut, messages=canon) def main(): ap = argparse.ArgumentParser() ap.add_argument("--url", default="http://127.0.0.1:8000/v1") ap.add_argument("--model", default="qwen38") ap.add_argument("--out", required=True) ap.add_argument("--n", type=int, default=10**9) ap.add_argument("--hours", type=float, default=6) ap.add_argument("--concurrency", type=int, default=48) ap.add_argument("--max_turns", type=int, default=8) ap.add_argument("--max_tokens", type=int, default=768, help="per call") ap.add_argument("--thinking", type=int, default=0, help="1 = let Qwen think (it overthinks)") ap.add_argument("--temperature", type=float, default=0.7) ap.add_argument("--top_p", type=float, default=0.8) ap.add_argument("--top_k", type=int, default=20) ap.add_argument("--kinds", default=",".join(TRAIN_KINDS)) ap.add_argument("--seed0", type=int, default=TEACHER_SEED0) a = ap.parse_args() os.makedirs(os.path.dirname(a.out), exist_ok=True) kinds = a.kinds.split(",") # resume: skip seeds already in the output file done = set() if os.path.exists(a.out): for l in open(a.out): done.add(json.loads(l)["seed"]) lock = threading.Lock() stats, t0 = Counter(), time.time() deadline = t0 + a.hours * 3600 f = open(a.out, "a") def job(i): if time.time() > deadline: return seed = a.seed0 + i if seed in done: return try: r = episode(seed, kinds[i % len(kinds)], live_args(a)) except Exception as e: # server hiccup: skip this seed with lock: stats["error"] += 1 return with lock: stats["episodes"] += 1 stats[f"ok_{r['kind']}"] += r["ok"] stats[f"n_{r['kind']}"] += 1 why, _ = teacher_reject(r) stats["usable"] += why is None stats[f"reject_{why}"] += why is not None stats["calls"] += r["calls"] stats["cut_calls"] += r["cut"] stats["completion_tokens"] += r["completion_tokens"] f.write(json.dumps(r) + "\n") f.flush() if stats["episodes"] % 50 == 0: el = time.time() - t0 print(json.dumps({"episodes": stats["episodes"], "usable": stats["usable"], "errors": stats["error"], "rejects": {k[7:]: v for k, v in stats.items() if k.startswith("reject_")}, "cut_call_rate": round(stats["cut_calls"] / max(1, stats["calls"]), 3), "tok_per_call": round(stats["completion_tokens"] / max(1, stats["calls"])), "tok_s": round(stats["completion_tokens"] / el), "ok_rate": {k: round(stats[f"ok_{k}"] / max(1, stats[f"n_{k}"]), 2) for k in kinds}, "elapsed_min": round(el / 60, 1)}), flush=True) with ThreadPoolExecutor(a.concurrency) as ex: for i in range(a.n): if time.time() > deadline: break ex.submit(job, i) # keep the queue short so the deadline is respected while ex._work_queue.qsize() > a.concurrency * 2: time.sleep(0.2) print("final", dict(stats), flush=True) if __name__ == "__main__": main()