Buckets:
| """Thinking chains from a teacher served on OpenRouter, for the items the local generation will not reach. | |
| python teacher_gen.py --n 4 --budget-usd 6.5 --deadline <unix time> | |
| The local run works through the pool from the front; this takes items from the back (tasksource only by default), | |
| so the two never overlap. Same prompt, same per-item option order (gen_order), same sampling as the local chains | |
| (temperature 0.6, top_p 0.95, top_k 20), no reasoning budget (the chain ends where the model ends it; max_tokens is | |
| a cost cap only and a chain that hits it counts as not closed, like a local chain that runs out of context). | |
| Records are written in gen.jsonl's format with "source": "teacher"; run_v2.sh merges them before training. | |
| The key is read from $OPENROUTER_KEY_FILE (default /workspace/.cache/openrouter_key) and is never logged. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import asyncio | |
| import json | |
| import os | |
| import random | |
| import time | |
| import httpx | |
| from transformers import AutoTokenizer | |
| from common import messages, parse_letter, read_jsonl, render, split_chain | |
| from gen_reasoning import gen_order | |
| URL = "https://openrouter.ai/api/v1/chat/completions" | |
| def to_sample(tok, msg: dict, finish: str | None, k: int) -> dict: | |
| """The local sample format: text = chain + '\\n</think>\\n\\n' + answer, as vLLM writes it after '<think>\\n'.""" | |
| reasoning = (msg.get("reasoning") or "").strip() | |
| content = (msg.get("content") or "").strip() | |
| closed = bool(reasoning) and finish == "stop" | |
| text = (reasoning + "\n</think>\n\n" + content) if closed else reasoning | |
| chain, answer = split_chain(text) | |
| letter = parse_letter(answer, k) if chain is not None else None | |
| return {"text": text, "tokens": len(tok(text, add_special_tokens=False)["input_ids"]), "finish": finish, | |
| "closed": chain is not None, "answer": letter, "looped": chain is None} | |
| async def one(client, key, args, item, order, seed, stats): | |
| body = {"model": args.model, "messages": messages(item, order), "max_tokens": args.max_tokens, | |
| "temperature": 0.6, "top_p": 0.95, "top_k": 20, "seed": seed, "reasoning": {"exclude": False}, | |
| "provider": {"require_parameters": True, "quantizations": ["fp8", "bf16", "fp16", "unknown"], | |
| "max_price": {"prompt": 0.3, "completion": 2.0}}} | |
| for attempt in range(6): | |
| try: | |
| r = await client.post(URL, json=body, headers={"Authorization": f"Bearer {key}"}, timeout=900) | |
| if r.status_code == 402: | |
| stats["out_of_credit"] = True | |
| return None | |
| if r.status_code in (408, 429, 500, 502, 503, 504): | |
| stats["retries"] += 1 | |
| await asyncio.sleep(2 ** attempt + random.random()) | |
| continue | |
| d = r.json() | |
| if "choices" not in d: | |
| stats["errors"] += 1 | |
| stats["last_error"] = str(d.get("error", d))[:300] | |
| await asyncio.sleep(2 ** attempt) | |
| continue | |
| u = d.get("usage") or {} | |
| stats["cost"] += float(u.get("cost") or 0.0) | |
| stats["completion_tokens"] += int(u.get("completion_tokens") or 0) | |
| stats["prompt_tokens"] += int(u.get("prompt_tokens") or 0) | |
| stats["providers"][d.get("provider", "?")] = stats["providers"].get(d.get("provider", "?"), 0) + 1 | |
| ch = d["choices"][0] | |
| return ch["message"], ch.get("finish_reason"), d.get("provider"), u | |
| except (httpx.HTTPError, json.JSONDecodeError) as e: | |
| stats["errors"] += 1 | |
| stats["last_error"] = repr(e)[:300] | |
| await asyncio.sleep(2 ** attempt) | |
| return None | |
| async def run(args): | |
| key = open(os.environ.get("OPENROUTER_KEY_FILE", "/workspace/.cache/openrouter_key")).read().strip() | |
| tok = AutoTokenizer.from_pretrained(args.tokenizer) | |
| pool = read_jsonl(args.pool) | |
| taken = set() | |
| for p in (args.local_gen, args.out): | |
| if os.path.exists(p): | |
| taken |= {r["id"] for r in read_jsonl(p)} | |
| tail = [x for x in reversed(pool) if x["id"] not in taken and (args.sets == "all" or x["id"].split(":")[0] in args.sets.split(","))] | |
| if args.limit: | |
| tail = tail[:args.limit] | |
| stats = {"cost": 0.0, "completion_tokens": 0, "prompt_tokens": 0, "retries": 0, "errors": 0, "providers": {}, | |
| "items_done": 0, "items_dropped": 0, "samples": 0, "closed": 0, "parsed": 0, "correct": 0} | |
| print(json.dumps({"candidates": len(tail), "model": args.model, "n": args.n, "budget_usd": args.budget_usd}), flush=True) | |
| sem = asyncio.Semaphore(args.concurrency) | |
| out = open(args.out, "a") | |
| t0 = time.time() | |
| per_item_est = [args.est_item_usd] | |
| async def item_job(item): | |
| order = gen_order(item) | |
| async with sem: | |
| res = await asyncio.gather(*[one(client, key, args, item, order, 1000 * j + 7, stats) for j in range(args.n)]) | |
| if any(r is None for r in res): | |
| stats["items_dropped"] += 1 | |
| return | |
| samples = [] | |
| for msg, finish, prov, u in res: | |
| s = to_sample(tok, msg, finish, len(item["options"])) | |
| s["provider"], s["api_completion_tokens"], s["api_cost"] = prov, u.get("completion_tokens"), u.get("cost") | |
| samples.append(s) | |
| stats["samples"] += 1 | |
| stats["closed"] += s["closed"] | |
| stats["parsed"] += s["answer"] is not None | |
| stats["correct"] += s["answer"] is not None and order[s["answer"]] == item["gold"] | |
| rec = {"id": item["id"], "order": order, "prompt_tokens": len(tok(render(tok, item, think=True, order=order), | |
| add_special_tokens=False)["input_ids"]), | |
| "samples": samples, "source": "teacher", "teacher": args.model} | |
| out.write(json.dumps(rec, ensure_ascii=False) + "\n") | |
| out.flush() | |
| stats["items_done"] += 1 | |
| per_item_est.append(sum(float(u.get("cost") or 0) for _, _, _, u in res)) | |
| async def reporter(): | |
| while True: | |
| await asyncio.sleep(60) | |
| el = time.time() - t0 | |
| print("progress " + json.dumps({"time": time.strftime("%H:%M:%S", time.gmtime()), "elapsed_min": round(el / 60, 1), | |
| **{k: v for k, v in stats.items()}, | |
| "out_tok_per_s": round(stats["completion_tokens"] / el)}), flush=True) | |
| limits = httpx.Limits(max_connections=args.concurrency * args.n + 8, max_keepalive_connections=args.concurrency * args.n) | |
| async with httpx.AsyncClient(limits=limits, http2=False) as client: | |
| rep = asyncio.create_task(reporter()) | |
| tasks = [] | |
| for item in tail: | |
| # stop submitting when the money (with what is in flight) or the time runs out | |
| while len(tasks) - stats["items_done"] - stats["items_dropped"] >= args.concurrency: | |
| await asyncio.sleep(0.5) | |
| in_flight = len(tasks) - stats["items_done"] - stats["items_dropped"] | |
| avg = sum(per_item_est[-200:]) / len(per_item_est[-200:]) | |
| if stats["cost"] + (in_flight + 1) * avg * 1.3 > args.budget_usd or stats.get("out_of_credit") or time.time() > args.deadline: | |
| break | |
| tasks.append(asyncio.create_task(item_job(item))) | |
| await asyncio.gather(*tasks) | |
| rep.cancel() | |
| out.close() | |
| summary = {k: v for k, v in stats.items()} | {"minutes": round((time.time() - t0) / 60, 1), | |
| "items_submitted": len(tasks)} | |
| print("EXIT " + json.dumps(summary), flush=True) | |
| os.makedirs("results", exist_ok=True) | |
| json.dump(summary, open("results/teacher_generation.json", "w"), indent=1) | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--model", default="qwen/qwen3.5-35b-a3b:nitro") | |
| ap.add_argument("--tokenizer", default="Qwen/Qwen3.5-4B") | |
| ap.add_argument("--pool", default="data/train_pool.jsonl") | |
| ap.add_argument("--local-gen", default="data/gen.jsonl") | |
| ap.add_argument("--out", default="data/gen_teacher.jsonl") | |
| ap.add_argument("--sets", default="ts", help="comma list of id prefixes (proc, ts) or 'all'") | |
| ap.add_argument("--n", type=int, default=4) | |
| ap.add_argument("--max-tokens", type=int, default=12288, help="cost cap per chain; not a reasoning budget") | |
| ap.add_argument("--concurrency", type=int, default=24, help="items in flight (each sends n requests)") | |
| ap.add_argument("--budget-usd", type=float, default=6.5) | |
| ap.add_argument("--est-item-usd", type=float, default=0.012, help="starting guess, replaced by measured cost") | |
| ap.add_argument("--deadline", type=float, default=float("inf")) | |
| ap.add_argument("--limit", type=int, default=None) | |
| args = ap.parse_args() | |
| asyncio.run(run(args)) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 8.96 kB
- Xet hash:
- 96753273c32a97f3d946bba5525af57c29ff411ab42f6d7fe86c59dc9adf462a
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.