# Reasoning corpus builder script

> Source: <https://gist.github.com/AJ-Gonzalez/31f1f8cfe45c233911e055d5cfed6da0>
> Published: 2026-11-11 12:24:31+00:00

|  | #!/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() |