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