spb/wpai0002_uqo

1#!/usr/bin/env python32"""Plan and run rewrites. Default provider is the deterministic mock (zero cost). Live Anthropic calls require3``--live`` and respect the budget ceiling in the config; ``--batch`` uses Message Batches (submit now,4``--collect`` later)."""5from __future__ import annotations67import argparse8import json9import random10import sys11from pathlib import Path1213ROOT = Path(__file__).resolve().parents[1]14sys.path.insert(0, str(ROOT / "src"))15from htr.rewrite.engine import BudgetExceeded, RewriteEngine, merge_plan, plan_requests  # noqa: E40216from htr.rewrite.mock_provider import MockProvider  # noqa: E40217from htr.utils.io import read_jsonl  # noqa: E402181920def select_passages(corpus_dir: Path, n: int, seed: int, splits_path: Path | None, allowed_splits: set[str], stratify: bool = True) -> list[dict]:21    docs = {d["doc_id"]: d for d in read_jsonl(corpus_dir / "documents.jsonl")}22    split_of = {a["doc_id"]: a["split"] for a in read_jsonl(splits_path)} if splits_path else {}23    ps = [p for p in read_jsonl(corpus_dir / "passages.jsonl") if p["eligible"] and (not split_of or split_of.get(p["doc_id"]) in allowed_splits)]24    rng = random.Random(seed)25    if not stratify or n >= len(ps):26        rng.shuffle(ps)27        return ps[:n]28    # stratified by (discipline, role): round-robin over strata29    strata: dict[tuple, list[dict]] = {}30    for p in ps:31        strata.setdefault((docs[p["doc_id"]].get("subject_area"), p["section_role"]), []).append(p)32    for v in strata.values():33        rng.shuffle(v)34    out: list[dict] = []35    keys = sorted(strata, key=lambda k: -len(strata[k]))36    while len(out) < n and any(strata.values()):37        for k in keys:38            if strata[k]:39                out.append(strata[k].pop())40                if len(out) >= n:41                    break42    return out434445def main() -> None:46    ap = argparse.ArgumentParser()47    ap.add_argument("--corpus", default=str(ROOT / "data/parsed/pilot"))48    ap.add_argument("--job", required=True, help="job dir under data/pairs/")49    ap.add_argument("--config", default=str(ROOT / "configs/rewriting/pilot.json"))50    ap.add_argument("--n-passages", type=int, default=50)51    ap.add_argument("--seed", type=int, default=1)52    ap.add_argument("--splits", default=None, help="splits.jsonl; only passages of --allowed-splits are rewritten")53    ap.add_argument("--allowed-splits", default="train,val,test,heldout_domain,heldout_period")54    ap.add_argument("--variants", default=None)55    ap.add_argument("--live", action="store_true", help="use the real Anthropic API (spends money)")56    ap.add_argument("--model", default=None, help="override model (e.g. a held-out model)")57    ap.add_argument("--batch", action="store_true")58    ap.add_argument("--collect", action="store_true", help="collect finished batches and exit")59    ap.add_argument("--max-requests", type=int, default=None)60    ap.add_argument("--estimate-only", action="store_true")61    a = ap.parse_args()62    cfg = json.loads(Path(a.config).read_text())63    model = a.model or cfg["model"]64    variants = (a.variants.split(",") if a.variants else cfg["variants"])65    corpus = Path(a.corpus)66    job_dir = ROOT / "data/pairs" / a.job67    if a.live:68        from htr.rewrite.anthropic_provider import AnthropicProvider, load_pricing6970        pricing = load_pricing()71        provider = AnthropicProvider(model, effort=cfg.get("effort", "low"), pricing=pricing)72    else:73        from htr.rewrite.anthropic_provider import load_pricing7475        pricing = load_pricing() if a.estimate_only else None76        provider = MockProvider(seed=a.seed)77        if a.estimate_only:78            provider.model = model  # estimate for the configured real model79            provider.name = cfg["provider"]  # so that planned request ids match the live job's ids80    engine = RewriteEngine(job_dir, provider, cfg["budget_usd_ceiling"], batch=a.batch, concurrency=cfg.get("concurrency", 4))81    sel = select_passages(corpus, a.n_passages, a.seed, Path(a.splits) if a.splits else None, set(a.allowed_splits.split(",")))82    pmap = {p["passage_id"]: p for p in read_jsonl(corpus / "passages.jsonl")}83    reqs = plan_requests(sel, variants, cfg["samples_per_variant"], provider.name, model, cfg["max_tokens"], cfg.get("effort", "low"))84    est = engine.estimate(reqs, pmap, pricing)85    print("estimate:", est)86    merge_plan(job_dir, reqs)  # the plan accumulates every selection (needed to interpret batch results later)87    engine.write_manifest({"corpus": str(corpus), "n_passages": len(sel), "variants": variants, "selection_seed": a.seed, "estimate": est, "live": a.live})88    if a.estimate_only:89        return90    try:91        if a.collect:92            print("collect:", engine.collect_batches(reqs, pmap))93        elif a.batch:94            if not a.live:95                print("batch mode requires --live (mock runs synchronously); running mock live instead")96                print("run:", engine.run_live(reqs, pmap, a.max_requests))97            else:98                print("submitted batches:", engine.submit_batches(reqs, pmap, max_requests=a.max_requests, pricing=pricing))99        else:100            print("run:", engine.run_live(reqs, pmap, a.max_requests))101    except BudgetExceeded as e:102        print("BUDGET STOP:", e)103        sys.exit(3)104105106if __name__ == "__main__":107    main()108