| """Write moth training chats with a teacher behind an OpenAI-compatible server (vLLM). |
| |
| python -m moth.gen_dialogs --pool --n-per-folder 250 # teacher-written home files -> out/file_pool.json |
| python -m moth.gen_dialogs --dialogs 12000 --start 0 # chats -> out/dialogs.jsonl (resumable) |
| |
| One chat: sample a home and a (persona, scenario, length); the teacher plays the human (USER_SIM_PROMPT) |
| and moth (TEACHER_BIBLE + real bash/read/write tools run in the sandbox). Records use tiny_agent's |
| canonical message format; filtering and lowercasing happen later in build_sft.py. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import asyncio |
| import collections |
| import json |
| import os |
| import random |
| import re |
| import sys |
| import time |
|
|
| from openai import AsyncOpenAI |
|
|
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| from moth import _tools as tools |
| from moth.agent import (AGENT_CHAT_REMIND, AGENT_LENGTHS, AGENT_LENSES, AGENT_PROMPT, AGENT_REMIND, |
| AGENT_SCENARIOS, AGENT_TAG, AGENT_TASKS, AGENT_USER_PERSONAS, HARNESS_STYLES, |
| MOTH_AGENT_RULES, OUTCOMES, agent_turn) |
| from moth.home import load_pool, sample_home |
| from moth.persona import (FILE_POOL_PROMPT, FOLDER_HINTS, FOLDERS, HOME_SCENARIOS, SCENARIOS, |
| TALK_SCENARIOS, TEACHER_BIBLE, TOPICS, USER_PERSONAS, USER_SIM_PROMPT) |
|
|
| if os.environ.get("MOTH_NO_BWRAP") == "1": |
| def _plain(workdir, command, timeout): |
| return ["timeout", "-k", "1", str(timeout), "env", "-i", f"HOME={workdir}", "PATH=/usr/bin:/bin", |
| "LANG=C.UTF-8", "bash", "-c", f"cd {workdir} && ulimit -t {timeout}; {command}"] |
| tools.bwrap_cmd = _plain |
|
|
| OAI_TOOLS = [ |
| {"type": "function", "function": {"name": "bash", "description": "Run a shell command in your home folder ~ (ls, cat, wc, grep...).", |
| "parameters": {"type": "object", "properties": {"command": {"type": "string"}}, "required": ["command"]}}}, |
| {"type": "function", "function": {"name": "read", "description": "Read a file with line numbers.", |
| "parameters": {"type": "object", "properties": {"path": {"type": "string"}, "offset": {"type": "integer"}, |
| "limit": {"type": "integer"}}, "required": ["path"]}}}, |
| {"type": "function", "function": {"name": "write", "description": "Create or overwrite a file.", |
| "parameters": {"type": "object", "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, |
| "required": ["path", "content"]}}}, |
| ] |
| SIM_SWITCH = ("[DIRECTOR NOTE: for this message only, you are NOT moth. You write the HUMAN side of the chat. " |
| "Output only the human's next message.]\n\n") |
| SIM_SYSTEM = ("You write one side of a chat for a story: the HUMAN who is chatting with moth, a tiny cute chatbot moth. " |
| "You are never moth. Type like a real person in a chat app, in the voice of the person you play: a grandparent " |
| "or formal person writes proper sentences with capitals and punctuation, a kid or teen types lowercase with slang, " |
| "a sloppy typer makes typos. Short (often 1-10 words), no greeting after the first message, rarely emojis, never " |
| "the same message twice. Reply with only the human's next chat message: no preface, no quotes, no notes.") |
| SIM_SHORT = ("\n\nIMPORTANT: this time write a tiny lazy reply of 1 to 5 words to moth's last message, typed your way, " |
| "that sounds {mood}. if it names something, use a word moth really said.") |
| SHORT_MOODS = ["amused", "doubtful", "impressed", "teasing", "agreeing", "curious about one word moth used", "bored", |
| "surprised", "touched", "a bit confused", "playfully annoyed", "sleepy", "delighted", "unconvinced"] |
| |
| |
| WHAT_X = ["what {x}?", "what {x}", "wait, {x}?", "{x}?", "what do you mean {x}?", "huh, {x}?", "What {x}?", "a {x}?", |
| "{x}? what", "why {x}?", "{x}?? lol", "what's {x}"] |
| WHAT_GENERIC = ["what?", "huh?", "what do you mean?", "wdym", "What do you mean?", "Pardon?", "wait what", "??", |
| "i don't get it", "what does that mean", "hm?", "Sorry, what?"] |
| WORD_STOP = set("that this with have from they them then than there their what when where which would could should " |
| "about maybe think really very just like little tiny small some much many more here know your yours " |
| "mine been being sure okay also still only once ever never always good nice thing things something " |
| "sorry thank thanks wait well kind sort feel feels want need make made said says look tell told it's " |
| "i'm don't that's what's you're can't didn't isn't doesn't i'll i've let's every other same into over " |
| "back down come came goes going gone does done want these those because while again even most " |
| "moth friend idea".split()) |
| DET = set("the a an my your his her our some this that these those its their any no".split()) |
| ADJ = set("little big tiny small soft warm old new nice good long short whole same other only first last best very cosy " |
| "cozy fuzzy dusty bright dark sleepy happy sweet great real funny tall round".split()) |
| WHATX_RE = re.compile(r"\b(?:what|which|why)(?: is| a| an| the)? ([a-z']+)\W*$", re.I) |
| RECALL_FACTS = [("your name (make one up)", "your name"), ("your pet's name and what animal it is (make it up)", "your pet"), |
| ("your favourite food", "your favourite food"), ("one small thing that happened to you today", |
| "what happened to you today"), ("your favourite colour", "your favourite colour"), |
| ("where you're going this weekend", "your weekend plan"), ("what your job is", "your job"), |
| ("your best friend's name (make one up)", "your best friend")] |
| FACT_WORDS = {"name", "favourite", "favorite", "food", "colour", "color", "pet", "today", "friend", "best", "job", "weekend", |
| "work", "love", "really", "go", "call", "happen", "moth", "thing", "day"} |
| RECALL_SHARE = {"p": 0.15} |
| LONG_SHARE = {"p": 0.15} |
| SIM_NO_REPEAT = "\n\nYou already sent: {prev}. Write something new, not any of these." |
| NO_THINK = {"chat_template_kwargs": {"enable_thinking": False}} |
| MOTH_REMIND = ("\n\n(stay moth: one short lowercase line, usually under 15 words. don't say what you are (a moth, small, " |
| "on a card) unless they ask what you are. files only if they ask about your home or want something saved: " |
| "one tiny command.)") |
| QBACK = ("\n(this time hand the turn back, on the same line: end with one small question about them, or bring back " |
| "something they told you earlier.)") |
| NO_QBACK = "\n(this time just answer: no question at the end, unless a game needs one.)" |
| RECALL_REMIND = "\n(they told you this earlier in this chat: answer from the chat.)" |
| LEAKS = ("<think>", "</think>", "<|tool_call", "<|im_") |
| TURNS = [3, 4, 5, 6, 6, 7, 8, 8, 9, 10, 12, 14] |
| HOME_SHARE = {"p": 0.15} |
| THINK = {"on": False, "budget": 256} |
| TURN_CHOICES: list[int] = [] |
| SCENARIO_IDS: list[int] = [] |
|
|
|
|
| def _sft_keep(): |
| """build_sft.keep (the build-time filter), so moth_turn resamples replies it would cut. build_sft imports torch, |
| tokenizers and tiny_agent at the top for the token rows only: stub whichever this env lacks.""" |
| import importlib.util |
| import types |
| try: |
| for name, attrs in (("torch", {}), ("tokenizers", {"Tokenizer": object})): |
| if name not in sys.modules and importlib.util.find_spec(name) is None: |
| sys.modules[name] = types.ModuleType(name) |
| sys.modules[name].__dict__.update(attrs) |
| if "tiny_agent.chat" not in sys.modules: |
| sys.modules.setdefault("tiny_agent", types.ModuleType("tiny_agent")) |
| sys.modules["tiny_agent.chat"] = types.ModuleType("tiny_agent.chat") |
| sys.modules["tiny_agent.chat"].render = None |
| from moth.build_sft import keep |
| return keep |
| except Exception as e: |
| print(f"WARNING: build_sft.keep not importable ({type(e).__name__}: {e}); no build-time checks", flush=True) |
| return None |
|
|
|
|
| SFT_KEEP = _sft_keep() |
|
|
|
|
| def sft_problem(meta: dict, msgs: list[dict], text: str) -> str | None: |
| """Why build_sft would cut the chat at this new reply (None if it keeps it, or the chat is already cut earlier).""" |
| if SFT_KEEP is None: |
| return None |
|
|
| def cut(t): |
| why = collections.Counter() |
| m = msgs + [{"role": "assistant", "think": None, "content": t, "tool_calls": []}] |
| out = SFT_KEEP({**meta, "home_after": {}, "messages": m}, why) |
| return None if out is not None and len(out) == len(m) else (next(iter(why), "cut").split(":")[0]) |
| try: |
| r = cut(text) |
| return r if r and not cut("ok.") else None |
| except Exception: |
| return None |
|
|
|
|
| def coerce(name: str, args: dict) -> dict: |
| if name == "read": |
| for k in ("offset", "limit"): |
| if k in args: |
| try: |
| args[k] = int(args[k]) |
| except (TypeError, ValueError): |
| args.pop(k) |
| return args |
|
|
|
|
| def transcript(msgs: list[dict]) -> str: |
| out = [] |
| for m in msgs: |
| if m["role"] == "user" and m["content"].startswith(AGENT_TAG): |
| body = m["content"][len(AGENT_TAG):].strip() |
| out.append("[your AI agent's latest message to you, which moth can also see]\n" + body[:600]) |
| elif m["role"] == "user": |
| out.append(f"you: {m['content']}") |
| elif m["role"] == "assistant": |
| for c in m.get("tool_calls") or []: |
| a = c["arguments"] |
| out.append(f"[moth used {c['name']}: {a.get('command') or a.get('path')}]") |
| if m.get("content"): |
| out.append(f"moth: {m['content']}") |
| elif m["role"] == "tool": |
| for r in m["results"]: |
| out.append("[output]\n" + (r if len(r) < 400 else r[:400] + " ...")) |
| return "\n".join(out) or "(nothing yet: write your first message)" |
|
|
|
|
| async def chat(client, model, messages, budget=None, **kw): |
| |
| |
| extra = {"thinking_token_budget": THINK["budget"] if budget is None else budget} if THINK["on"] else NO_THINK |
| for attempt in range(4): |
| try: |
| return await client.chat.completions.create(model=model, messages=messages, extra_body=extra, **kw) |
| except Exception as e: |
| if attempt == 3: |
| raise |
| await asyncio.sleep(5 * (attempt + 1)) |
| last = e |
|
|
|
|
| META = re.compile(r"\bthe user\b|\bin character\b|\bmoth (should|would|is asked)\b|\blet me (think|check|re)|" |
| r"\bno markdown\b|\bkeep it (short|simple)\b|\b\d+ words\b|\blet me try\b|\bmisunderstand", re.I) |
| COMPLEX_CMD = re.compile(r"\bfor\b|\bwhile\b|\||&&|;|\*|-la\b|2>|\$\(|\becho\b|<<|\s-\w*R") |
| _P = r"(?:~/)?(?:[\w.-]+/)?[\w.-]+" |
| TINY_CMD = re.compile(rf"\s*(?:ls(?: -[1a])?(?: (?:~/?|(?:~/)?[\w-]+/?))*|cat {_P}(?: {_P})*|" |
| rf"(?:wc|head)(?: -[a-z]*\d*| -n \d+)? {_P})\s*") |
| |
| AGENT_BAD = re.compile(r"\b[a-z]\w*_\w+\b|\w\(\)|\b\w+\.(?:py|js|ts|tsx|json|csv|md|ya?ml|toml|sh|rs|go|html|css|sql|cfg|ini)\b|" |
| r"\b(?:ruff|mypy|pytest|npm|pip|git|github|docker|eslint|webpack|cargo|kubectl|ci|lint|linter|" |
| r"endpoints?|commits?|regex|api|json|csv|yaml|http|https|redis|schema|parser|tuple|repo)\b|" |
| r"\btests? (?:all )?(?:pass|passed|passes|passing)\b|\b(?:all )?\d+ tests\b|^(?:\S+\s+){25}\S|" |
| r"\bi(?:['’]?m| am) (?:just |only )?a (?:little |tiny |small )?moth\b|" |
| r"\bcan(?:['’]?t|not) (?:understand|read|follow)\b|\bdon['’]?t (?:understand|get it)\b") |
| BUG = re.compile(r"\bbugs?\b|\bgrand ?ma\b", re.I) |
| MOTIF = re.compile(r"\bfans?\b", re.I) |
|
|
|
|
| def grams(t: str, n=5) -> set: |
| w = re.sub(r"[^\w\s']", " ", t.lower()).split() |
| return {tuple(w[i:i + n]) for i in range(len(w) - n + 1)} |
|
|
|
|
| EXAMPLE_GRAMS = set().union(*map(grams, re.findall(r'-> "([^"]+)"', MOTH_AGENT_RULES) + |
| ["ohhh it fixed a bug! i'm a bug too. i hope it was nice to it."])) |
|
|
|
|
| STOP = set("the and for that this with you your are was were has have had not but all can its into from what when then " |
| "than they them there their our out just one two too very like little tiny she her his who how why yes now " |
| "new agent said".split()) |
| VIBES = [(r"\bsorry|\bapolog", lambda m: re.search(r"sorry|apolog", m)), (r"\bask|\bquestion", lambda m: "?" in m), |
| (r"\bwords\b|\blong\b|\bbusy\b|\btired\b", lambda m: len(m.split()) > 80), |
| (r"\bproud\b|\bexcited\b|\bhappy\b|\bpolite\b", lambda m: True)] |
|
|
|
|
| def stems(t: str) -> set: |
| return {re.sub(r"(es|s|ed|ing)$", "", w) for w in re.findall(r"[a-z]{3,}", t.lower()) if w not in STOP} |
|
|
|
|
| def agent_stale(text: str, msgs: list[dict]) -> bool: |
| """r3 agent chats: LFM copied the few-shot lines, repeated itself, said "bug" to 1 in 3 agent messages that never |
| mentioned one (the bible's grandma lore), and riffed on words that aren't in the message at all.""" |
| mine = set().union(*(grams(m["content"]) for m in msgs if m["role"] == "assistant" and m.get("content"))) |
| if grams(text) & (EXAMPLE_GRAMS | mine): |
| return True |
| if MOTIF.search(text) and any(MOTIF.search(m["content"] or "") for m in msgs if m["role"] == "assistant"): |
| return True |
| last = msgs[-1]["content"] |
| if not last.startswith(AGENT_TAG): |
| return False |
| if BUG.search(text) and not re.search(r"\bbugs?\b", last, re.I): |
| return True |
| body, t = last[len(AGENT_TAG):].lower(), text.lower() |
| return not (stems(t) & stems(body) or any(re.search(c, t) and ok(body) for c, ok in VIBES)) |
|
|
|
|
| def complex_cmd(cmd: str) -> bool: |
| return bool(cmd) and bool(COMPLEX_CMD.search(cmd) or not TINY_CMD.fullmatch(cmd)) |
|
|
|
|
| def bash_cmd(tc) -> str: |
| try: |
| return str(json.loads(tc.function.arguments or "{}").get("command", "")) if tc.function.name == "bash" else "" |
| except (json.JSONDecodeError, AttributeError): |
| return "" |
|
|
|
|
| def teacher_view(u: str, remind: str | None = None, extra: str = "") -> str: |
| if remind: |
| return u + remind |
| return u + (MOTH_REMIND if THINK["on"] else "") + extra |
|
|
|
|
| def reply_text(r, cut_ok=False) -> str | None: |
| """The reply, or None when the call broke: cut off by max_tokens, or reasoning/tool markup leaked into the text.""" |
| c = r.choices[0] |
| t = (c.message.content or "").strip() |
| if (c.finish_reason == "length" and not cut_ok) or any(x in t for x in LEAKS): |
| return None |
| return t |
|
|
|
|
| def last_moth(msgs: list[dict]) -> str: |
| return next((m["content"] for m in reversed(msgs) if m["role"] == "assistant" and m.get("content")), "") |
|
|
|
|
| def confused(last: str, mine: list[str]) -> str: |
| """A "what X?" turn about a word moth really said (or a plain "what?"), built in code.""" |
| toks, words = [w.strip("'") for w in re.findall(r"[a-z']+", last.lower())], [] |
| for j, w in enumerate(toks): |
| if w in DET: |
| x = next((v for v in toks[j + 1:j + 4] if v not in ADJ), "") |
| if len(x) >= 3 and x not in WORD_STOP and x not in DET: |
| words.append(x) |
| for _ in range(4): |
| t = random.choice(WHAT_X).format(x=random.choice(words)) if words and random.random() < 0.75 \ |
| else random.choice(WHAT_GENERIC) |
| if t not in mine: |
| return t |
| return t |
|
|
|
|
| async def user_turn(client, model, persona, scenario, turns, i, msgs, topic="anything", directive="") -> str | None: |
| p = USER_SIM_PROMPT.format(persona=persona, scenario=scenario, topic=topic, turns=turns, i=i, transcript=transcript(msgs)) |
| mine = [m["content"] for m in msgs if m["role"] == "user" and not m["content"].startswith(AGENT_TAG)] |
| last = last_moth(msgs) |
| short = not directive and msgs and random.random() < 0.4 |
| if short and last and random.random() < 0.3: |
| return confused(last, mine) |
| if mine: |
| p += SIM_NO_REPEAT.format(prev=" / ".join(f'"{x[:100]}"' for x in mine[-8:])) |
| p += directive |
| if THINK["on"]: |
| |
| if short: |
| p += SIM_SHORT.format(mood=random.choice(SHORT_MOODS)) |
| r = await chat(client, model, [{"role": "system", "content": SIM_SYSTEM}, {"role": "user", "content": p}], |
| budget=0, temperature=1.0, top_p=0.95, max_tokens=100) |
| else: |
| |
| r = await chat(client, model, [{"role": "system", "content": TEACHER_BIBLE}, {"role": "user", "content": SIM_SWITCH + p}], |
| tools=OAI_TOOLS, tool_choice="none", temperature=1.0, top_p=0.95, max_tokens=80) |
| t = reply_text(r, cut_ok=True) |
| if not t or "[END]" in t: |
| return None |
| lines = [x.strip() for x in t.split("\n") if x.strip()] |
| if len(lines) > 1 and lines[0].endswith(":"): |
| lines = lines[1:] |
| t = lines[0].strip('"\u201c\u201d').strip() |
| t = re.sub(r"^(you|me|user|friend|human)\s*:\s*", "", t, flags=re.I).strip() |
| if not t or re.match(r"moth\s*:", t, re.I): |
| return None |
| w = WHATX_RE.search(t) if len(t.split()) <= 5 and last else None |
| if w and w.group(1).lower() not in WORD_STOP | {"do", "does", "is", "are", "you", "u", "mean", "happened", "now", "else", |
| "about", "it", "that", "this", "then", "not", "tho", "so"} \ |
| and stems(w.group(1)).isdisjoint(stems(last)) and w.group(1).lower() not in last.lower(): |
| return confused(last, mine) |
| return t[:300] |
|
|
|
|
| async def moth_turn(client, model, oai: list[dict], msgs: list[dict], ws, stats, talk_only=False, must=None, |
| meta=None) -> bool | str: |
| """Extend oai (teacher view) and msgs (canonical) with moth's reply to the last user message.""" |
| avoid: list[str] = [] |
| for k in range(8 if talk_only else 5): |
| |
| |
| tk = {} if talk_only else {"tools": OAI_TOOLS, "tool_choice": "auto" if k < 3 else "none"} |
| view = oai if not avoid else oai[:-1] + [{**oai[-1], "content": oai[-1]["content"] + "\n(not these, they are " |
| "stale, dodgy or too techy: " + " / ".join(avoid[-3:]) + ")"}] |
| r = await chat(client, model, view, temperature=1.0 if avoid else 0.85, top_p=0.95, |
| max_tokens=THINK["budget"] + 300 if THINK["on"] else 400, **tk) |
| m = r.choices[0].message |
| if r.choices[0].finish_reason == "length": |
| continue |
| if m.tool_calls and THINK["on"] and any(complex_cmd(bash_cmd(tc)) for tc in m.tool_calls): |
| continue |
| if m.tool_calls: |
| calls, results, oai_calls = [], [], [] |
| for tc in m.tool_calls[:3]: |
| try: |
| args = coerce(tc.function.name, json.loads(tc.function.arguments or "{}")) |
| except json.JSONDecodeError: |
| return False |
| if tc.function.name not in ("bash", "read", "write"): |
| return False |
| res = await asyncio.to_thread(ws.call, tc.function.name, args) |
| calls.append({"name": tc.function.name, "arguments": args}) |
| results.append(res) |
| oai_calls.append({"id": tc.id, "type": "function", |
| "function": {"name": tc.function.name, "arguments": json.dumps(args)}}) |
| oai.append({"role": "assistant", "content": None, "tool_calls": oai_calls}) |
| for tc, res in zip(oai_calls, results): |
| oai.append({"role": "tool", "tool_call_id": tc["id"], "content": res}) |
| msgs.append({"role": "assistant", "think": None, "content": "", "tool_calls": calls}) |
| msgs.append({"role": "tool", "results": results}) |
| stats["tool_calls"] += len(calls) |
| continue |
| text = re.sub(r"^moth\s*:\s*", "", reply_text(r) or "", flags=re.I) |
| if not text or META.search(text): |
| continue |
| if must and k < 4 and stems(text).isdisjoint(must): |
| stats["recall_resample"] += 1 |
| continue |
| why = sft_problem(meta, msgs, text) if meta else None |
| if why and (talk_only or k < 4): |
| stats[f"sft_{why}"] += 1 |
| avoid.append(text) |
| continue |
| if talk_only and AGENT_BAD.search(text.lower()): |
| stats["agent_resample"] += 1 |
| avoid.append(text) |
| continue |
| if talk_only and agent_stale(text, msgs): |
| stats["agent_stale"] += 1 |
| avoid.append(text) |
| continue |
| oai.append({"role": "assistant", "content": text}) |
| msgs.append({"role": "assistant", "think": None, "content": text, "tool_calls": []}) |
| if why: |
| stats[f"sft_stop_{why}"] += 1 |
| return "stop" |
| return True |
| return False |
|
|
|
|
| AGENT_POOL: list[dict] = [] |
| AGENT_SHARE = {"p": 0.25} |
| MID_RUN = re.compile(r":\s*\Z|\b(?:Now )?[Ll]et me (?!know)|\bPhase \d|\bI'll now\b|\bLet's\b|^(?:Now|Next),? |" |
| r"^I (?:see|notice)\b|^(?:Excellent|Perfect|Great)! Now\b", re.M) |
|
|
|
|
| async def agent_message(client, model, rng) -> tuple[str, str]: |
| """(raw message, source): a real one from the pool most of the time, else one written by the teacher.""" |
| if AGENT_POOL and rng.random() < 0.75: |
| m = rng.choice(AGENT_POOL) |
| return m["text"], m["source"] |
| p = AGENT_PROMPT.format(style=rng.choice(HARNESS_STYLES), task=rng.choice(AGENT_TASKS), outcome=rng.choice(OUTCOMES), |
| length=rng.choice(AGENT_LENGTHS)) |
| r = await chat(client, model, [{"role": "user", "content": p}], temperature=0.9, top_p=0.95, |
| max_tokens=THINK["budget"] + 900 if THINK["on"] else 700) |
| return reply_text(r) or "", "teacher" |
|
|
|
|
| async def agent_dialog(client, model, seed, rng, files, stats) -> dict | None: |
| """moth watches a big coding agent: agent message -> moth's silly take -> the human vents or chats -> repeat.""" |
| persona, scenario = rng.choice(AGENT_USER_PERSONAS), rng.choice(AGENT_SCENARIOS) |
| msgs: list[dict] = [] |
| oai = [{"role": "system", "content": TEACHER_BIBLE}] |
| sources = [] |
|
|
| async def say(u: str, remind: str) -> bool: |
| msgs.append({"role": "user", "content": u}) |
| oai.append({"role": "user", "content": teacher_view(u, remind)}) |
| if await moth_turn(client, model, oai, msgs, ws, stats, talk_only=True, |
| meta={"type": "agent", "home": files, "topic": ""}): |
| return True |
| msgs.pop() |
| return False |
|
|
| with tools.Workspace(files) as ws: |
| ok = True |
| for k in range(rng.randint(2, 4)): |
| if k == 0 and rng.random() < 0.4: |
| u = await user_turn(client, model, persona, scenario, 6, 1, msgs, "the AI agent") |
| if u and not (ok := await say(u, AGENT_CHAT_REMIND)): |
| break |
| raw, src = await agent_message(client, model, rng) |
| if not raw: |
| break |
| if not (ok := await say(agent_turn(raw), AGENT_REMIND.format(lens=rng.choice(AGENT_LENSES)))): |
| break |
| sources.append(src) |
| for j in range(rng.choice([0, 1, 1, 2, 3])): |
| u = await user_turn(client, model, persona, scenario, 6, j + 2, msgs, "the AI agent") |
| if u is None or not (ok := await say(u, AGENT_CHAT_REMIND)): |
| break |
| if not ok: |
| break |
| if not sources: |
| stats["bad"] += 1 |
| return None |
| stats["agent_cut"] += not ok |
| return {"seed": seed, "type": "agent", "persona": persona, "scenario": scenario, "agent_sources": sources, |
| "turns": len(msgs), "home": files, "home_after": files, "messages": msgs} |
|
|
|
|
| async def one_dialog(client, model, seed, pool, stats) -> dict | None: |
| rng = random.Random(seed) |
| files = sample_home(rng, pool) |
| if not SCENARIO_IDS and rng.random() < AGENT_SHARE["p"]: |
| return await agent_dialog(client, model, seed, rng, files, stats) |
| if SCENARIO_IDS: |
| scen = [SCENARIOS[i] for i in SCENARIO_IDS] |
| else: |
| scen = HOME_SCENARIOS if rng.random() < HOME_SHARE["p"] else TALK_SCENARIOS |
| persona, scenario, turns = rng.choice(USER_PERSONAS), rng.choice(scen), rng.choice(TURN_CHOICES or TURNS) |
| topic = rng.choice(TOPICS) |
| talk = scen is TALK_SCENARIOS |
| if talk and not TURN_CHOICES and rng.random() < LONG_SHARE["p"]: |
| turns = rng.randint(16, 24) |
| recall = None |
| if talk and rng.random() < RECALL_SHARE["p"]: |
| fact, short = rng.choice(RECALL_FACTS) |
| fi = rng.choice([1, 2]) |
| recall = {"fact": short, "share_turn": fi, "ask_turn": fi + rng.randint(5, 8)} |
| turns = max(turns, recall["ask_turn"] + rng.randint(1, 3)) |
| msgs: list[dict] = [] |
| oai = [{"role": "system", "content": TEACHER_BIBLE}] |
| with tools.Workspace(files) as ws: |
| for i in range(1, turns + 1): |
| directive, extra = "", QBACK if rng.random() < 0.2 else NO_QBACK |
| if recall and i == recall["share_turn"]: |
| directive = f"\n\nIn THIS message, casually tell moth {fact}." |
| elif recall and i == recall["ask_turn"]: |
| directive = (f"\n\nIn THIS message, ask moth if it remembers {short} (you told it earlier). Ask in your " |
| "own words and don't give the answer.") |
| extra = RECALL_REMIND |
| if (recall and i < recall["ask_turn"]) or (turns >= 16 and i < turns - 3): |
| directive += "\n\nThe chat is not finished: do not write [END] yet." |
| u = await user_turn(client, model, persona, scenario, turns, i, msgs, topic, directive) |
| for _ in range(3 if "[END]" in directive else 0): |
| if u is not None: |
| break |
| u = await user_turn(client, model, persona, scenario, turns, i, msgs, topic, directive) |
| if u is None: |
| break |
| if recall and i == recall["share_turn"]: |
| recall["said"] = u |
| must = None |
| if recall and i == recall["ask_turn"] and recall.get("said"): |
| recall["asked"] = u |
| must = stems(recall["said"]) - FACT_WORDS or None |
| msgs.append({"role": "user", "content": u}) |
| oai.append({"role": "user", "content": teacher_view(u, extra=extra)}) |
| r = await moth_turn(client, model, oai, msgs, ws, stats, must=must, |
| meta={"type": "home" if scen is HOME_SCENARIOS else "talk", "home": files, "topic": topic}) |
| if not r: |
| stats["bad"] += 1 |
| return None |
| if r == "stop": |
| break |
| if not any(m["role"] == "assistant" for m in msgs): |
| return None |
| final = {} |
| for root, _, names in os.walk(ws.root): |
| for n in names: |
| p = os.path.join(root, n) |
| final[os.path.relpath(p, ws.root)] = open(p, errors="replace").read()[:2000] |
| return {"seed": seed, "type": "home" if scen is HOME_SCENARIOS else "talk", "persona": persona, |
| "scenario": scenario, "topic": topic, "turns": turns, "recall": recall, "home": files, |
| "home_after": final, "messages": msgs} |
|
|
|
|
| async def run_dialogs(a, client, model): |
| pool = load_pool(a.pool_file) |
| path = os.path.join(a.out, "dialogs.jsonl") |
| done = set() |
| if os.path.exists(path): |
| for line in open(path): |
| try: |
| done.add(json.loads(line)["seed"]) |
| except (ValueError, KeyError): |
| pass |
| todo = [s for s in range(a.start, a.start + a.dialogs) if s not in done] |
| stats = collections.Counter(ok=0, bad=0, err=0, tool_calls=0, agent_resample=0, agent_stale=0, agent_cut=0, |
| recall_resample=0) |
| sem = asyncio.Semaphore(a.concurrency) |
| f = open(path, "a") |
| t0 = time.time() |
|
|
| async def go(seed): |
| async with sem: |
| try: |
| d = await one_dialog(client, model, seed, pool, stats) |
| except Exception as e: |
| stats["err"] += 1 |
| print(f"seed {seed}: {type(e).__name__}: {e}", flush=True) |
| return |
| if d: |
| f.write(json.dumps(d, ensure_ascii=False) + "\n") |
| f.flush() |
| stats["ok"] += 1 |
| if stats["ok"] % 100 == 0: |
| el = time.time() - t0 |
| print(f"{stats} {stats['ok'] / el * 60:.0f} dialogs/min, {len(todo) - stats['ok']} left", flush=True) |
|
|
| await asyncio.gather(*(go(s) for s in todo)) |
| print(f"done {stats} in {(time.time() - t0) / 60:.1f} min", flush=True) |
|
|
|
|
| async def run_pool(a, client, model): |
| path = a.pool_file |
| pool = json.load(open(path)) if os.path.exists(path) else {} |
| sem = asyncio.Semaphore(a.concurrency) |
|
|
| async def batch(folder, k): |
| async with sem: |
| have = list(pool.get(folder, {})) |
| avoid = ", ".join(random.Random(k).sample(have, min(40, len(have)))) or "none" |
| p = FILE_POOL_PROMPT.format(n=10, folder=folder, hint=FOLDER_HINTS[folder], avoid=avoid, bible=TEACHER_BIBLE) |
| try: |
| r = await chat(client, model, [{"role": "user", "content": p}], temperature=1.0, top_p=0.95, |
| max_tokens=6000 if THINK["on"] else 1500) |
| txt = r.choices[0].message.content or "" |
| items = json.loads(txt[txt.index("["): txt.rindex("]") + 1]) |
| except Exception: |
| return |
| for it in items: |
| name = re.sub(r"[^a-z0-9_.]", "_", str(it.get("name", "")).lower().strip()) |
| if not name.endswith(".txt") or len(name) > 48 or not it.get("content"): |
| continue |
| pool.setdefault(folder, {}).setdefault(name, str(it["content"]).strip()[:400]) |
|
|
| await asyncio.gather(*(batch(fo, k) for fo in FOLDERS for k in range(a.n_per_folder // 10))) |
| json.dump(pool, open(path, "w"), indent=1, ensure_ascii=False) |
| print("pool sizes:", {k: len(v) for k, v in pool.items()}, flush=True) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--base-url", default="http://127.0.0.1:8000/v1") |
| ap.add_argument("--model", default=None) |
| ap.add_argument("--out", default="out") |
| ap.add_argument("--pool", action="store_true") |
| ap.add_argument("--n-per-folder", type=int, default=250) |
| ap.add_argument("--pool-file", default=None) |
| ap.add_argument("--dialogs", type=int, default=0) |
| ap.add_argument("--start", type=int, default=0) |
| ap.add_argument("--concurrency", type=int, default=96) |
| ap.add_argument("--agent-share", type=float, default=0.25, help="fraction of chats watching a big agent") |
| ap.add_argument("--agent-pool", default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "agent_pool.json")) |
| ap.add_argument("--eval-agent-msgs", default=os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), |
| "data", "eval_agent_msgs.json"), |
| help="held-out eval agent messages: dropped from the agent pool") |
| ap.add_argument("--home-share", type=float, default=0.15, help="fraction of chats about moth's home/files") |
| ap.add_argument("--recall-share", type=float, default=0.15, help="talk chats: share a fact early, ask 5+ turns later") |
| ap.add_argument("--long-share", type=float, default=0.15, help="talk chats with 16-24 user turns") |
| ap.add_argument("--think", action="store_true", help="teacher always reasons first (LFM2.5): cap it, split it off") |
| ap.add_argument("--think-budget", type=int, default=256, help="--think: max reasoning tokens per moth/agent call") |
| ap.add_argument("--turns", default="", help="comma-separated user-turn counts to sample from") |
| ap.add_argument("--scenario-ids", default="", help="comma-separated indexes into persona.SCENARIOS") |
| a = ap.parse_args() |
| SCENARIO_IDS.extend(int(x) for x in a.scenario_ids.split(",") if x) |
| TURN_CHOICES.extend(int(x) for x in a.turns.split(",") if x) |
| THINK["on"], THINK["budget"] = a.think, a.think_budget |
| HOME_SHARE["p"] = a.home_share |
| RECALL_SHARE["p"], LONG_SHARE["p"] = a.recall_share, a.long_share |
| AGENT_SHARE["p"] = a.agent_share |
| if os.path.exists(a.agent_pool): |
| AGENT_POOL.extend(m for m in json.load(open(a.agent_pool)) if not MID_RUN.search(m["text"])) |
| n = len(AGENT_POOL) |
| if os.path.exists(a.eval_agent_msgs): |
| norm = lambda t: " ".join(str(t).split()) |
| ev = [m["text"] if isinstance(m, dict) else m for m in json.load(open(a.eval_agent_msgs))] |
| exact, heads = {norm(t) for t in ev}, {norm(t)[:200] for t in ev} |
| AGENT_POOL[:] = [m for m in AGENT_POOL if norm(m["text"]) not in exact and norm(m["text"])[:200] not in heads] |
| print(f"agent pool: removed {n - len(AGENT_POOL)} eval messages ({len(ev)} held out)", flush=True) |
| else: |
| print(f"agent pool: WARNING no eval file at {a.eval_agent_msgs}, nothing held out", flush=True) |
| print(f"agent pool: {len(AGENT_POOL)} final messages", flush=True) |
| os.makedirs(a.out, exist_ok=True) |
| a.pool_file = a.pool_file or os.path.join(a.out, "file_pool.json") |
| client = AsyncOpenAI(base_url=a.base_url, api_key="x", timeout=300) |
|
|
| async def run(): |
| model = a.model or (await client.models.list()).data[0].id |
| if a.pool: |
| await run_pool(a, client, model) |
| if a.dialogs: |
| await run_dialogs(a, client, model) |
| asyncio.run(run()) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
