Reasoning corpus builder script

A developer published a standalone Python script that regenerates the 250-row MiMo-2.6-Pro high-reasoning corpus using a different teacher model via OpenRouter. The tool inlines the MiMo taxonomy, supports category and profile filtering (coder, indiedev, genagent), preserves original row ids, and writes JSONL in the same schema with per-row usage and latency, plus a --dry-run mode to inspect selections without spending tokens.

| | /usr/bin/env python3 | | | """Regenerate the MiMo reasoning corpus with a different teacher via OpenRouter. | | | Single file: the MiMo taxonomy is inlined, no local module is imported. Needs typer , | | | requests , rich ; datasets only when the source is an HF id; transformers only for | | | --estimate . Run --check to verify the environment. | | | Reads the prompts system + user from the 250-row MiMo set and writes a new JSONL in the | | | same schema the MiMo file uses: | | | {"source": ..., "model": <openrouter model , "messages": | | | {"role": "system", "content": ...}, | | | {"role": "user", "content": ...}, | | | {"role": "assistant", "content": ..., "reasoning": ...} } | | | Extra keys per row: id original 0-based row index , category , usage , latency s . | | | id matters: categories are defined by row index, so a filtered run keeps the original ids. | | | Usage: | | | export OPENROUTER API KEY=sk-or-... | | | python scripts/openrouter reasoning corpus.py \ | | | --model deepseek/deepseek-v4-pro \ | | | --profile coder \ | | | --out data/deepseek-v4-pro-coder.jsonl | | | one teacher, everything except stats and finance, including unlabelled rows | | | python scripts/openrouter reasoning corpus.py -m moonshotai/kimi-k3.1 \ | | | --no-stats --no-finance --uncategorized -o data/kimi-k3.1-all.jsonl | | | inspect the selection without spending tokens | | | python scripts/openrouter reasoning corpus.py -m deepseek/deepseek-v4-pro --dry-run | | | """ | | | from future import annotations | | | import importlib.util | | | import json | | | import os | | | import platform | | | import random | | | import sys | | | import threading | | | import time | | | from collections import defaultdict | | | from concurrent.futures import ThreadPoolExecutor, as completed | | | from pathlib import Path | | | from typing import Any, Optional | | | import requests | | | import typer | | | from rich.console import Console | | | from rich.table import Table | | | ROOT = Path file .resolve .parent.parent | | | MiMo-2.6-Pro taxonomy, judge-labelled 2026-11-07. Inlined so this file is standalone. | | | CATEGORIES = { | | | "1 code": "0 2 30 68 96 99 117 134 137 139 143 144 145 147 160 164 169 181 184 " | | | "188 191 194 199 201 210 218 219 231 234 243 244 246", | | | "2 cs": "3 4 5 6 7 8 9 10 11 12 14 15 16 18 19 20 21 39 48 56 65 74 83 92 112 " | | | "115 120 121 158", | | | "4 stats": "23 41 59 67 76 85 94 107 141", | | | "5 finance": "25 33 42 51 60 69 77 86 95 105 116 124 202", | | | "6 business": "26 28 34 37 43 45 52 54 61 63 70 72 78 81 87 89 101 102 106 110 " | | | "114 131", | | | "8 science": "36 44 62 71 100 108 123 129", | | | "9 humanities": "1 17 24 29 38 47 50 55 64 73 82 91 98 109 122 130 136 138 140 " | | | "142 151 153 155 165 183 187 195 198 207 223 233 245", | | | } | | | PROFILES = { | | | "coder": "1 code", "2 cs", "4 stats" , | | | "indiedev": "1 code", "2 cs", "5 finance", "6 business", "9 humanities" , | | | "genagent": "5 finance", "6 business", "8 science", "9 humanities" , | | | } | | | API URL = "https://openrouter.ai/api/v1/chat/completions" | | | MODELS URL = "https://openrouter.ai/api/v1/models" | | | HF SOURCE = "xtools-at/Xiaomi-MiMo-2.6-Pro-High-Reasoning-250x" | | | LOCAL SOURCE = ROOT / "data" / "mimo-2.6-pro-high-250x.jsonl" | | | CLI flag - category key in CATEGORIES | | | FLAG TO CAT = { | | | "code": "1 code", | | | "cs": "2 cs", | | | "stats": "4 stats", | | | "finance": "5 finance", | | | "business": "6 business", | | | "science": "8 science", | | | "humanities": "9 humanities", | | | } | | | console = Console | | | app = typer.Typer add completion=False, no args is help=True | | | class ApiError RuntimeError : | | | def init self, status: int, body: str : | | | self.status = status | | | self.body = body | | | super . init f"HTTP {status}: {body :300 }" | | | ---------------------------------------------------------------- input | | | def load rows source: str, split: str - list dict : | | | path = Path source | | | if path.exists : | | | rows = json.loads line for line in path.read text .splitlines if line.strip | | | console.print f" dim source: {path} {len rows } rows /dim " | | | return rows | | | from datasets import load dataset imported lazily; local JSONL needs no datasets | | | ds = load dataset source, split=split | | | rows = dict r for r in ds | | | console.print f" dim source: {source} {split} {len rows } rows /dim " | | | return rows | | | def category index - dict int, str : | | | """row id - category. Categories are disjoint in the MiMo taxonomy; first wins if not.""" | | | out: dict int, str = {} | | | for cat, ids in CATEGORIES.items : | | | for raw in ids.split : | | | out.setdefault int raw , cat | | | return out | | | def select | | | rows: list dict , | | | profile: str, | | | flags: dict str, bool , | | | uncategorized: bool, | | | limit: int, | | | offset: int, | | | - list tuple int, Optional str , dict : | | | if profile = "all" and profile not in PROFILES: | | | raise typer.BadParameter f"unknown profile {profile r}; choose all, {', '.join PROFILES }" | | | enabled = set PROFILES profile if profile = "all" else set CATEGORIES | | | enabled = {c for c in enabled if flags.get c, True } | | | index = category index | | | picked: list tuple int, Optional str , dict = | | | for i, row in enumerate rows : | | | cat = index.get i | | | if cat is None: | | | if not uncategorized: | | | continue | | | elif cat not in enabled: | | | continue | | | picked.append i, cat, row | | | picked = picked offset: | | | if limit: | | | picked = picked :limit | | | return picked | | | ---------------------------------------------------------------- request | | | def build body | | | model: str, | | | messages: list dict , | | | temperature: float, | | | top p: Optional float , | | | max tokens: int, | | | stream: bool, | | | reasoning: bool, | | | reasoning effort: Optional str , | | | seed: Optional int , | | | extra body: Optional str , | | | - dict str, Any : | | | body: dict str, Any = { | | | "model": model, | | | "messages": messages, | | | "temperature": temperature, | | | "max tokens": max tokens, | | | "stream": stream, | | | "usage": {"include": True}, | | | } | | | if top p is not None: | | | body "top p" = top p | | | if seed is not None: | | | body "seed" = seed | | | if reasoning: | | | body "reasoning" = {"effort": reasoning effort} if reasoning effort else {"enabled": True} | | | if extra body: | | | body.update json.loads extra body | | | return body | | | def headers for api key: str, referer: Optional str , title: Optional str - dict str, str : | | | headers = { | | | "Authorization": f"Bearer {api key}", | | | "Content-Type": "application/json", | | | } | | | if referer: | | | headers "HTTP-Referer" = referer | | | if title: | | | headers "X-Title" = title | | | return headers | | | def call stream | | | body: dict, headers: dict str, str , timeout: int | | | - tuple str, str, Optional dict : | | | reasoning: list str = | | | content: list str = | | | usage: Optional dict = None | | | finish: Optional str = None | | | served: Optional str = None | | | with requests.post | | | API URL, headers=headers, json=body, stream=True, timeout= 15, timeout | | | as resp: | | | if resp.status code = 400: | | | raise ApiError resp.status code, resp.text | | | for raw in resp.iter lines decode unicode=False : | | | if not raw: | | | continue | | | line = raw.decode "utf-8", errors="replace" | | | if line.startswith ":" : | | | continue | | | if not line.startswith "data:" : | | | continue | | | payload = line 5: .strip | | | if payload == " DONE ": | | | break | | | chunk = json.loads payload | | | if chunk.get "error" : | | | raise ApiError 200, json.dumps chunk "error" :500 | | | if chunk.get "model" : | | | served = chunk "model" | | | if chunk.get "usage" : | | | usage = chunk "usage" | | | for choice in chunk.get "choices" or : | | | if choice.get "finish reason" : | | | finish = choice "finish reason" | | | delta = choice.get "delta" or {} | | | if delta.get "reasoning" : | | | reasoning.append delta "reasoning" | | | if delta.get "reasoning content" : | | | reasoning.append delta "reasoning content" | | | if delta.get "content" : | | | content.append delta "content" | | | return "".join reasoning , "".join content , usage, finish, served | | | def call plain | | | body: dict, headers: dict str, str , timeout: int | | | - tuple str, str, Optional dict : | | | resp = requests.post API URL, headers=headers, json=body, timeout= 15, timeout | | | if resp.status code = 400: | | | raise ApiError resp.status code, resp.text | | | data = resp.json | | | if data.get "error" : | | | raise ApiError 200, json.dumps data "error" :500 | | | message = data "choices" 0 "message" | | | reasoning = message.get "reasoning" or message.get "reasoning content" or "" | | | return reasoning, message.get "content" or "", data.get "usage" , data "choices" 0 .get "finish reason" , data.get "model" | | | def with retries fn, body: dict, headers: dict, timeout: int, retries: int - tuple str, str, Optional dict , Optional str , Optional str : | | | delay = 2.0 | | | for attempt in range retries + 1 : | | | try: | | | return fn body, headers, timeout | | | except ApiError as exc: | | | retryable = exc.status in 408, 409, 429, 500, 502, 503, 504 | | | if not retryable or attempt == retries: | | | raise | | | sleep for = delay 2 attempt + random.uniform 0, 1 | | | console.print f" yellow retry {attempt + 1}/{retries} in {sleep for:.1f}s /yellow dim {exc.body :120 } /dim " | | | time.sleep sleep for | | | except requests.ConnectionError, requests.Timeout as exc: | | | if attempt == retries: | | | raise ApiError 0, str exc | | | sleep for = delay 2 attempt + random.uniform 0, 1 | | | console.print f" yellow network retry {attempt + 1}/{retries} in {sleep for:.1f}s /yellow " | | | time.sleep sleep for | | | raise AssertionError "unreachable" | | | ---------------------------------------------------------------- cost | | | def fetch model entry model: str, api key: Optional str - Optional dict : | | | headers = {"Authorization": f"Bearer {api key}"} if api key else {} | | | resp = requests.get MODELS URL, headers=headers, timeout=30 | | | resp.raise for status | | | for entry in resp.json .get "data", : | | | if model in entry.get "id" , entry.get "canonical slug" : | | | return entry | | | return None | | | def resolve variant model: str, api key: Optional str , pin: bool - dict: | | | """Resolve an OpenRouter alias to the dated variant it currently points at.""" | | | entry = fetch model entry model, api key | | | if entry is None: | | | console.print f" yellow {model r} not in /models; keeping the id as-is /yellow " | | | return {"model": request model, "canonical": None, "name": None, "created": None} | | | canonical = entry.get "canonical slug" or entry.get "id" | | | return { | | | "model": canonical if pin else model, | | | "canonical": canonical, | | | "name": entry.get "name" , | | | "created": entry.get "created" , | | | } | | | def fetch pricing | | | model: str, | | | api key: Optional str , | | | price prompt: Optional float , | | | price completion: Optional float , | | | - tuple float, float : | | | """Return USD per prompt token, USD per completion token .""" | | | if price prompt is not None and price completion is not None: | | | return price prompt / 1e6, price completion / 1e6 | | | entry = fetch model entry model, api key | | | if entry is None: | | | raise typer.BadParameter | | | f"{model r} not found in OpenRouter /models; pass --price-prompt/--price-completion USD per 1M tokens " | | | | | | pricing = entry.get "pricing" or {} | | | per prompt = float pricing.get "prompt" or 0.0 | | | per completion = float pricing.get "completion" or 0.0 | | | return | | | price prompt / 1e6 if price prompt is not None else per prompt, | | | price completion / 1e6 if price completion is not None else per completion, | | | | | | def get tokenizer name: str : | | | try: | | | from transformers import AutoTokenizer | | | return AutoTokenizer.from pretrained name, trust remote code=True | | | except Exception as exc: noqa: BLE001 - estimation must not fail the run | | | console.print f" yellow tokenizer {name r} unavailable {str exc :120 } ; falling back to chars/3.6 /yellow " | | | return None | | | def count tokens tok, text: str - int: | | | if tok is None: | | | return int len text / 3.6 | | | return len tok text, add special tokens=False .input ids | | | def estimate costs | | | picked: list tuple int, Optional str , dict , | | | tok, | | | price in: float, | | | price out: float, | | | max tokens: int, | | | assume out: int, | | | - tuple float, float : | | | per: dict str, dict = defaultdict lambda: {"rows": 0, "in": 0, "out": } | | | for , cat, row in picked: | | | prompt text = "".join | | | m.get "content", "" for m in row "messages" if m.get "role" in "system", "user" | | | | | | in tok = count tokens tok, prompt text + 8 role markers | | | assistant = row "messages" -1 | | | if assume out: | | | out tok = assume out | | | else: | | | out tok = count tokens | | | tok, assistant.get "reasoning" or "" + "\n" + assistant.get "content" or "" | | | | | | bucket = per cat or "uncategorized" | | | bucket "rows" += 1 | | | bucket "in" += in tok | | | bucket "out" .append out tok | | | table = Table title="cost estimate" | | | for col, right in "category", False , "rows", True , "input tok", True , | | | "out p50", True , "out p90", True , "cost p50", True , "cost p90", True : | | | table.add column col, justify="right" if right else "left" | | | total rows = total in = 0 | | | total p50 = total p90 = 0.0 | | | for cat, bucket in sorted per.items : | | | outs = sorted bucket "out" | | | p50 = outs len outs // 2 | | | p90 = outs min len outs - 1, int len outs 0.9 | | | c50 = bucket "in" price in + bucket "rows" p50 price out | | | c90 = bucket "in" price in + bucket "rows" p90 price out | | | table.add row cat, str bucket "rows" , f"{bucket 'in' :,}", f"{p50:,}", f"{p90:,}", | | | f"${c50:.2f}", f"${c90:.2f}" | | | total rows += bucket "rows" | | | total in += bucket "in" | | | total p50 += c50 | | | total p90 += c90 | | | table.add row "TOTAL", str total rows , f"{total in:,}", "", "", | | | f"${total p50:.2f}", f"${total p90:.2f}", style="bold" | | | console.print table | | | ceiling = total in price in + total rows max tokens price out | | | console.print | | | f" dim output prior = MiMo traces tokenized with the same tokenizer; the target teacher may differ. " | | | f"Hard ceiling at max tokens={max tokens}: ${ceiling:.2f} /dim " | | | | | | return total p50, total p90 | | | ---------------------------------------------------------------- output | | | HARD DEPS = "requests", "typer", "rich" | | | OPTIONAL DEPS = {"datasets": "pip install datasets", "transformers": "pip install transformers"} | | | def installed module: str - bool: | | | return importlib.util.find spec module is not None | | | def check environment | | | , | | | source: str, | | | out: Path, | | | api key: Optional str , | | | need transformers: bool, | | | need key: bool, | | | report only: bool = False, | | | - None: | | | """Fail early with one actionable message instead of a traceback mid-run.""" | | | problems: list str = | | | if sys.version info < 3, 9 : | | | problems.append f"Python = 3.9 required, found {platform.python version }" | | | missing = m for m in HARD DEPS if not installed m | | | if missing: | | | problems.append f"missing packages: {', '.join missing } - pip install {' '.join missing }" | | | src path = Path source | | | data suffixes = ".jsonl", ".json", ".jsonl.gz", ".json.gz", ".parquet" | | | is hf = | | | not src path.exists | | | and source.count "/" == 1 | | | and not source.endswith data suffixes | | | | | | if src path.exists : | | | src state = f"file {src path} " | | | elif is hf: | | | src state = f"HF dataset {source} " | | | if not installed "datasets" : | | | problems.append "source is an HF dataset id but datasets is missing - pip install datasets" | | | else: | | | src state = "MISSING" | | | problems.append f"source not found: {src path} pass --source <jsonl\|hf-id " | | | if need transformers and not installed "transformers" : | | | console.print " yellow transformers not installed; --estimate will fall back to chars/3.6 /yellow " | | | key state = "set" if api key else "not set" | | | if need key and not api key: | | | problems.append "missing API key: pass --api-key or set OPENROUTER API KEY" | | | parent = out.parent | | | probe = parent | | | while not probe.exists and probe = probe.parent: | | | probe = probe.parent | | | if not os.access probe, os.W OK : | | | problems.append f"output directory not writable: {parent}" | | | if report only: | | | console.print f"python {platform.python version } {sys.executable} " | | | for m in HARD DEPS: | | | console.print f" {m:12s} {'ok' if installed m else 'MISSING'}" | | | for m, hint in OPTIONAL DEPS.items : | | | console.print f" {m:12s} {'ok' if installed m else 'not installed - ' + hint}" | | | console.print f" api key {key state}" | | | console.print f" source {src state}" | | | console.print f" output {out}" | | | if problems: | | | for problem in problems: | | | console.print f" red x {problem} /red " | | | raise typer.Exit 1 | | | if report only: | | | console.print " green all required checks passed /green " | | | raise typer.Exit 0 | | | def slugify model: str - str: | | | return "".join c if c.isalnum or c in "-. " else "-" for c in model .strip "-" | | | @app.command | | | def main | | | model: Optional str = typer.Option None, "--model", "-m", help="OpenRouter model id, e.g. deepseek/deepseek-v4-pro" , | | | api key: Optional str = typer.Option None, "--api-key", envvar="OPENROUTER API KEY", help="or set OPENROUTER API KEY" , | | | source: Optional str = typer.Option None, "--source", "-s", help="JSONL path or HF dataset id default: local file, else the MiMo HF dataset " , | | | split: str = typer.Option "train", "--split", help="split when --source is a HF dataset id" , | | | out: Optional Path = typer.Option None, "--out", "-o", help="default data/openrouter-<model .jsonl" , | | | profile: str = typer.Option "all", "--profile", help="all \| coder \| indiedev \| genagent" , | | | code: bool = typer.Option True, "--code/--no-code" , | | | cs: bool = typer.Option True, "--cs/--no-cs" , | | | stats: bool = typer.Option True, "--stats/--no-stats" , | | | finance: bool = typer.Option True, "--finance/--no-finance" , | | | business: bool = typer.Option True, "--business/--no-business" , | | | science: bool = typer.Option True, "--science/--no-science" , | | | humanities: bool = typer.Option True, "--humanities/--no-humanities" , | | | uncategorized: bool = typer.Option False, "--uncategorized/--no-uncategorized", help="include the 105 rows outside the taxonomy" , | | | limit: int = typer.Option 0, "--limit", help="0 = all selected rows" , | | | offset: int = typer.Option 0, "--offset" , | | | concurrency: int = typer.Option 4, "--concurrency", "-j" , | | | temperature: float = typer.Option 0.7, "--temperature" , | | | top p: Optional float = typer.Option None, "--top-p" , | | | max tokens: int = typer.Option 32768, "--max-tokens" , | | | timeout: int = typer.Option 900, "--timeout", help="read timeout per request, seconds" , | | | retries: int = typer.Option 5, "--retries" , | | | reasoning: bool = typer.Option True, "--reasoning/--no-reasoning", help="send OpenRouter's reasoning parameter" , | | | reasoning effort: Optional str = typer.Option None, "--reasoning-effort", help="low \| medium \| high" , | | | stream: bool = typer.Option True, "--stream/--no-stream" , | | | seed: Optional int = typer.Option None, "--seed" , | | | extra body: Optional str = typer.Option None, "--extra-body", help="JSON merged into the request body" , | | | resume: bool = typer.Option True, "--resume/--no-resume" , | | | retry truncated: bool = typer.Option False, "--retry-truncated", help="on resume, drop finish reason=length rows and regenerate them" , | | | dry run: bool = typer.Option False, "--dry-run" , | | | list categories: bool = typer.Option False, "--list-categories" , | | | check: bool = typer.Option False, "--check", help="print the environment/dependency report and exit" , | | | referer: Optional str = typer.Option None, "--referer", envvar="OPENROUTER REFERER" , | | | title: Optional str = typer.Option None, "--title", envvar="OPENROUTER TITLE" , | | | estimate: bool = typer.Option False, "--estimate", help="print token/cost estimate and exit" , | | | tokenizer: str = typer.Option "LiquidAI/LFM3-1.7B-Base", "--tokenizer", help="HF id or local path for counting" , | | | assume output tokens: int = typer.Option 0, "--assume-output-tokens", help="override output tokens per row" , | | | price prompt: Optional float = typer.Option None, "--price-prompt", help="USD per 1M prompt tokens" , | | | price completion: Optional float = typer.Option None, "--price-completion", help="USD per 1M completion tokens" , | | | budget: float = typer.Option 0.0, "--budget", help="stop after this many USD of measured spend" , | | | resolve model: bool = typer.Option True, "--resolve-model/--no-resolve-model", help="look up the dated variant the alias points at" , | | | pin variant: bool = typer.Option False, "--pin-variant/--no-pin-variant", help="send the canonical dated slug instead of the alias" , | | | model info: bool = typer.Option False, "--model-info", help="print the resolved model entry and exit" , | | | - None: | | | flags = { | | | FLAG TO CAT k : v | | | for k, v in { | | | "code": code, "cs": cs, "stats": stats, "finance": finance, | | | "business": business, "science": science, "humanities": humanities, | | | }.items | | | } | | | if list categories: | | | index = category index | | | console.print f" bold {len CATEGORIES } categories, {len index } labelled rows /bold " | | | for cat, ids in CATEGORIES.items : | | | rows = int x for x in ids.split | | | console.print f" {cat:14s} {len rows :3d} {'enabled' if flags cat else 'disabled'}" | | | return | | | source = source or str LOCAL SOURCE if LOCAL SOURCE.exists else HF SOURCE | | | if check: | | | check environment source=source, out=out or ROOT / "data" / "openrouter-model.jsonl", | | | api key=api key, need transformers=False, need key=False, report only=True | | | return | | | if not model: | | | raise typer.BadParameter "--model is required unless --list-categories/--check is used" | | | if profile = "all" and profile not in PROFILES: | | | raise typer.BadParameter f"unknown profile {profile r}; choose all, {', '.join PROFILES }" | | | if reasoning effort and reasoning effort not in "low", "medium", "high" : | | | raise typer.BadParameter "--reasoning-effort must be low, medium or high" | | | variant = {"model": model, "canonical": None, "name": None, "created": None} | | | if resolve model or pin variant or model info: | | | variant = resolve variant model, api key, pin variant | | | request model = variant "model" | | | out = out or ROOT / "data" / f"openrouter-{slugify request model }.jsonl" | | | check environment source=source, out=out, api key=api key, | | | need transformers=estimate, | | | need key=not dry run or estimate or model info | | | if model info: | | | entry = fetch model entry model, api key | | | if entry is None: | | | console.print f" red {model r} not found in OpenRouter /models /red " | | | raise typer.Exit 1 | | | subset = {k: entry.get k for k in "id", "canonical slug", "name", "created", | | | "hugging face id", "context length", "pricing" } | | | subset "max completion tokens" = entry.get "top provider" or {} .get "max completion tokens" | | | console.print json json.dumps subset | | | return | | | rows = load rows source, split | | | picked = select rows, profile, flags, uncategorized, limit, offset | | | if not picked: | | | console.print " red no rows selected /red " | | | raise typer.Exit 1 | | | enabled = sorted {c for , c, in picked if c} | | | counts = {c: sum 1 for , cc, in picked if cc == c for c in enabled} | | | if None in {c for , c, in picked}: | | | counts "uncategorized" = sum 1 for , c, in picked if c is None | | | console.print | | | f" bold {len picked } rows /bold selected profile={profile} " | | | f"categories=" + ", ".join f"{k}:{v}" for k, v in counts.items | | | | | | if not uncategorized: | | | unlabelled = len rows - len category index | | | if unlabelled: | | | console.print f" dim {unlabelled} unlabelled rows available; pass --uncategorized to include /dim " | | | if estimate: | | | price in, price out = fetch pricing model, api key, price prompt, price completion | | | console.print f"pricing: ${price in 1e6:.3f}/M in, ${price out 1e6:.3f}/M out" | | | if variant.get "canonical" : | | | console.print f"variant: {variant 'canonical' } {variant.get 'name' or ''}" | | | estimate costs picked, get tokenizer tokenizer , price in, price out, | | | max tokens, assume output tokens | | | return | | | err path = out.with suffix out.suffix + ".errors.jsonl" | | | done: set int = set | | | if resume and out.exists : | | | kept: list str = | | | dropped = 0 | | | for line in out.read text .splitlines : | | | if not line.strip : | | | continue | | | try: | | | record = json.loads line | | | rid = int record "id" | | | except json.JSONDecodeError, KeyError, ValueError : | | | continue | | | if retry truncated and record.get "finish reason" == "length": | | | dropped += 1 | | | continue | | | kept.append line | | | done.add rid | | | if dropped: | | | tmp = out.with suffix out.suffix + ".tmp" | | | tmp.write text "\n".join kept + "\n" if kept else "" | | | tmp.replace out | | | console.print f" dim retry: dropped {dropped} truncated rows, kept {len kept } /dim " | | | if done: | | | console.print f" dim resume: {len done } rows already in {out} /dim " | | | todo = i, c, r for i, c, r in picked if i not in done | | | if dry run: | | | sample = todo 0 if todo else picked 0 | | | msgs = m for m in sample 2 "messages" if m.get "role" in "system", "user" | | | body = build body request model, msgs, temperature, top p, max tokens, stream, reasoning, | | | reasoning effort, seed, extra body | | | console.print f" bold would write: /bold {out}" | | | canonical = f" {variant 'canonical' } " if variant.get "canonical" and variant "canonical" = request model else "" | | | console.print f" bold would call: /bold {request model}{canonical} x{len todo } concurrency={concurrency}" | | | console.print " dim first request body: /dim " | | | console.print json json.dumps body, ensure ascii=False :2000 | | | return | | | if not api key: | | | console.print " red missing API key: pass --api-key or set OPENROUTER API KEY /red " | | | raise typer.Exit 1 | | | headers = headers for api key, referer, title | | | call = call stream if stream else call plain | | | lock = threading.Lock | | | out.parent.mkdir parents=True, exist ok=True | | | stats = {"ok": 0, "fail": 0, "skipped": 0, "truncated": 0, "reasoning chars": 0, "content chars": 0, | | | "cost": 0.0, "prompt tokens": 0, "completion tokens": 0} | | | started = time.perf counter | | | def work item: tuple int, Optional str , dict - None: | | | i, cat, row = item | | | if budget 0: | | | with lock: | | | if stats "cost" = budget: | | | stats "skipped" += 1 | | | return | | | msgs = m for m in row "messages" if m.get "role" in "system", "user" | | | body = build body request model, msgs, temperature, top p, max tokens, stream, reasoning, | | | reasoning effort, seed, extra body | | | t0 = time.perf counter | | | try: | | | reasoning text, content, usage, finish, served = with retries call, body, headers, timeout, retries | | | if not content and not reasoning text: | | | raise ApiError 200, "empty completion" | | | record = { | | | "source": row.get "source", "" , | | | "model": request model, | | | "model requested": model, | | | "model canonical": variant.get "canonical" , | | | "model name": variant.get "name" , | | | "served model": served, | | | "messages": msgs + {"role": "assistant", "content": content, "reasoning": reasoning text} , | | | "id": i, | | | "category": cat, | | | "usage": usage, | | | "finish reason": finish, | | | "latency s": round time.perf counter - t0, 2 , | | | } | | | line = json.dumps record, ensure ascii=False | | | with lock: | | | with out.open "a", encoding="utf-8" as fh: | | | fh.write line + "\n" | | | stats "ok" += 1 | | | stats "reasoning chars" += len reasoning text | | | stats "content chars" += len content | | | if finish == "length": | | | stats "truncated" += 1 | | | if usage: | | | stats "prompt tokens" += int usage.get "prompt tokens" or 0 | | | stats "completion tokens" += int usage.get "completion tokens" or 0 | | | if isinstance usage.get "cost" , int, float : | | | stats "cost" += float usage "cost" | | | flag = " yellow TRUNCATED /yellow " if finish == "length" else "" | | | console.print | | | f" green {stats 'ok' :3d}/{len todo } /green id={i:3d} {str cat :14s} " | | | f"{record 'latency s' :7.1f}s r={len reasoning text :6d} c={len content :6d}{flag}" | | | | | | except Exception as exc: noqa: BLE001 - one bad row must not kill the run | | | with lock: | | | stats "fail" += 1 | | | with err path.open "a", encoding="utf-8" as fh: | | | fh.write json.dumps {"id": i, "category": cat, "model": request model, | | | "error": str exc :1000 }, ensure ascii=False + "\n" | | | console.print f" red FAIL /red id={i:3d} {str cat :14s} {str exc :160 }" | | | with ThreadPoolExecutor max workers=max 1, concurrency as pool: | | | futures = pool.submit work, item for item in todo | | | for future in as completed futures : | | | future.result | | | elapsed = time.perf counter - started | | | console.print | | | f"\n bold done /bold {stats 'ok' } ok, {stats 'fail' } failed" | | | + f", {stats 'skipped' } skipped budget " if stats "skipped" else "" | | | + f", {stats 'truncated' } truncated at max tokens" if stats "truncated" else "" | | | + f" in {elapsed / 60:.1f} min {stats 'ok' / max elapsed, 1e-9 :.2f} rows/s " | | | | | | console.print | | | f"tokens: {stats 'prompt tokens' :,} in, {stats 'completion tokens' :,} out" | | | + f" \| measured cost ${stats 'cost' :.4f}" if stats "cost" else "" | | | | | | if stats "ok" and stats "cost" : | | | full count = len select rows, profile, flags, uncategorized, 0, offset | | | console.print | | | f" dim projected for all {full count} selected rows: " | | | f"${stats 'cost' / stats 'ok' full count:.2f} /dim " | | | | | | console.print | | | f"reasoning {stats 'reasoning chars' / 1e6:.2f}M chars, " | | | f"content {stats 'content chars' / 1e6:.2f}M chars" | | | + f", cost ${stats 'cost' :.4f}" if stats "cost" else "" | | | | | | console.print f"wrote {out}" + f" errors {err path}" if stats "fail" else "" | | | if name == " main ": | | | app |