SPB Git

spb/wp44_uqo

Public
10.9 KB · 139 lines python
1#!/usr/bin/env python32"""WP44 — taxonomy validation: stratified sample + two independent LLM raters (Anthropic SDK, structured outputs) + metrics.3  --draw     : sample from DuckDB → taxonomy/validation/sample.parquet (+ sample.csv)4  --rate     : rate with --model (one rater per run), resumable, batches of 10 works per request → validation/ratings_<model>.jsonl5  --metrics  : compute agreement metrics → validation/metrics.json, disagreements.csv6API key: read from the file given by --env-file (never from the repository); the key is only used here."""7import argparse, json, os, random, time, re, duckdb, yaml, pandas as pd, numpy as np89SCHEMA = {"type": "object", "properties": {"ratings": {"type": "array", "items": {"type": "object", "properties": {10    "work_id": {"type": "integer"}, "categories": {"type": "array", "items": {"type": "string"}, "maxItems": 5},11    "primary": {"type": ["string", "null"]}, "confidence": {"type": "number"}, "rationale": {"type": "string"}},12    "required": ["work_id", "categories", "primary", "confidence", "rationale"], "additionalProperties": False}}},13    "required": ["ratings"], "additionalProperties": False}141516def draw(root, seed, per_cat, n_none):17    con = duckdb.connect(os.path.join(root, "data/processed/wp44.duckdb"), read_only=True)18    con.execute(f"CREATE OR REPLACE VIEW work_primary AS SELECT * FROM read_parquet('{root}/data/processed/classification/primary/*.parquet')")19    con.execute(f"CREATE OR REPLACE VIEW work_categories AS SELECT * FROM read_parquet('{root}/data/processed/classification/cats/*.parquet')")20    df = con.execute(f"""21        WITH base AS (SELECT c.id AS work_id, c.title, a.abstract, c.source_name, c.publication_year AS year, p.primary_category,22                             (SELECT list(category_id) FROM work_categories k WHERE k.work_id = c.id) AS rule_categories,23                             hash(c.id || '{seed}') AS h24                      FROM corpus_works c JOIN oa_abstracts a ON a.work_id = c.id LEFT JOIN work_primary p ON p.work_id = c.id WHERE c.eligible),25             ranked AS (SELECT *, row_number() OVER (PARTITION BY primary_category ORDER BY h) rn FROM base)26        SELECT * EXCLUDE (h, rn) FROM ranked WHERE (primary_category IS NOT NULL AND rn <= {per_cat}) OR (primary_category IS NULL AND rn <= {n_none})""").df()27    rng = random.Random(seed); idx = list(range(len(df))); rng.shuffle(idx)28    df["pi_review"] = False; df.loc[df.index[idx[:200]], "pi_review"] = True29    d = os.path.join(root, "taxonomy/validation"); os.makedirs(d, exist_ok=True)30    df.to_parquet(f"{d}/sample.parquet"); df.drop(columns=["abstract"]).to_csv(f"{d}/sample.csv", index=False)31    print("sample", len(df), "works;", df.primary_category.nunique(), "categories;", df.primary_category.isna().sum(), "unassigned")323334def rate(root, model, env_file, ontology, batch):35    import anthropic36    key = next(l.split("=", 1)[1].strip() for l in open(os.path.expanduser(env_file)) if l.startswith("ANTHROPIC_API_KEY="))37    client = anthropic.Anthropic(api_key=key, max_retries=5)38    ont = yaml.safe_load(open(ontology))39    defs = "\n".join(f"- {c['id']}: {c['name']} — {c['definition']}" for c in ont["categories"])40    system = ("You are an expert classifier of computer-science research papers for a systematic mapping study. Assign each paper to the "41              "applicable categories of the ontology below (zero to five; multi-label allowed; choose a primary category when at least one applies). "42              "Judge from the title, abstract, venue and year only. Be conservative: assign a category only when the paper's main "43              "contribution clearly belongs to it. If the paper is not computer-science research, return no categories.\n\nOntology:\n" + defs)44    d = os.path.join(root, "taxonomy/validation"); df = pd.read_parquet(f"{d}/sample.parquet")45    out_p = f"{d}/ratings_{model}.jsonl"; done = set()46    if os.path.exists(out_p):47        done = {json.loads(l)["work_id"] for l in open(out_p)}48    todo = df[~df.work_id.isin(done)]; usage = {"in": 0, "out": 0, "calls": 0}49    with open(out_p, "a") as fh:50        for i in range(0, len(todo), batch):51            b = todo.iloc[i:i + batch]52            items = "\n\n".join(f"### work_id {r.work_id}\nTitle: {r.title}\nVenue: {r.source_name or 'n/a'} ({r.year})\nAbstract: {(r.abstract or '')[:2500]}" for r in b.itertuples())53            for attempt in range(4):54                try:55                    resp = client.messages.create(model=model, max_tokens=8000, system=[{"type": "text", "text": system, "cache_control": {"type": "ephemeral"}}],56                                                  messages=[{"role": "user", "content": f"Rate the following {len(b)} papers. Return one rating per work_id.\n\n{items}"}],57                                                  output_config={"format": {"type": "json_schema", "schema": SCHEMA}, "effort": "low"})58                    break59                except anthropic.RateLimitError:60                    time.sleep(20 * (attempt + 1))61            if resp.stop_reason == "refusal":62                print("refusal on batch", i); continue63            text = next(bk.text for bk in resp.content if bk.type == "text"); data = json.loads(text)64            usage["in"] += resp.usage.input_tokens + (resp.usage.cache_read_input_tokens or 0); usage["out"] += resp.usage.output_tokens; usage["calls"] += 165            for r in data["ratings"]:66                r["model"] = model; fh.write(json.dumps(r) + "\n")67            fh.flush()68            if usage["calls"] % 20 == 0:69                print(f"{usage['calls']} calls, {i+len(b)}/{len(todo)} works, tokens in={usage['in']} out={usage['out']}", flush=True)70    json.dump(usage, open(f"{d}/usage_{model}.json", "w"), indent=1); print("done", usage)717273def metrics(root, ontology):74    from sklearn.metrics import cohen_kappa_score75    d = os.path.join(root, "taxonomy/validation"); df = pd.read_parquet(f"{d}/sample.parquet")76    cats = [c["id"] for c in yaml.safe_load(open(ontology))["categories"]]77    raters = {}78    for f in sorted(os.listdir(d)):79        if f.startswith("ratings_") and f.endswith(".jsonl"):80            rows = [json.loads(l) for l in open(os.path.join(d, f))]81            raters[f[8:-6]] = {r["work_id"]: r for r in rows}82    ids = [w for w in df.work_id if all(w in r for r in raters.values())]83    rule = {r.work_id: set(r.rule_categories or []) for r in df.itertuples()}84    rule_p = {r.work_id: r.primary_category for r in df.itertuples()}85    M = {"n_rated_by_all": len(ids), "raters": list(raters)}86    def prf(ref_sets, pred_sets):87        out = {}88        for c in cats:89            tp = sum(1 for w in ids if c in ref_sets[w] and c in pred_sets[w]); fp = sum(1 for w in ids if c not in ref_sets[w] and c in pred_sets[w])90            fn = sum(1 for w in ids if c in ref_sets[w] and c not in pred_sets[w])91            p = tp / (tp + fp) if tp + fp else None; r = tp / (tp + fn) if tp + fn else None92            out[c] = {"tp": tp, "fp": fp, "fn": fn, "precision": p, "recall": r, "f1": (2 * p * r / (p + r)) if p and r else (0.0 if (tp + fp + fn) else None)}93        micro_tp = sum(v["tp"] for v in out.values()); micro_fp = sum(v["fp"] for v in out.values()); micro_fn = sum(v["fn"] for v in out.values())94        mp = micro_tp / max(micro_tp + micro_fp, 1); mr = micro_tp / max(micro_tp + micro_fn, 1)95        f1s = [v["f1"] for v in out.values() if v["f1"] is not None]96        return {"per_category": out, "micro_precision": mp, "micro_recall": mr, "micro_f1": 2 * mp * mr / max(mp + mr, 1e-9), "macro_f1": float(np.mean(f1s))}97    for name, R in raters.items():98        sets = {w: set(x for x in R[w]["categories"] if x in cats) for w in ids}99        M[f"rules_vs_{name}"] = prf(sets, rule)100        M[f"kappa_primary_rules_vs_{name}"] = float(cohen_kappa_score([rule_p.get(w) or "none" for w in ids], [R[w]["primary"] or "none" for w in ids]))101    if len(raters) >= 2:102        (a, A), (b, B) = list(raters.items())[:2]103        M[f"kappa_primary_{a}_vs_{b}"] = float(cohen_kappa_score([A[w]["primary"] or "none" for w in ids], [B[w]["primary"] or "none" for w in ids]))104        cons = {w: set(A[w]["categories"]) & set(B[w]["categories"]) for w in ids}105        M["rules_vs_consensus"] = prf(cons, rule)106        # Krippendorff's alpha (interval form on MASI distance between the two raters' label sets; Passonneau 2006)107        def masi(x, y):108            x, y = set(x), set(y)109            if not x and not y: return 0.0110            j = len(x & y) / len(x | y); m = 1.0 if x == y else (0.67 if (x < y or y < x) else (0.33 if x & y else 0.0))111            return 1.0 - j * m112        pairs = [(A[w]["categories"], B[w]["categories"]) for w in ids]113        Do = float(np.mean([masi(x, y) ** 2 for x, y in pairs]))114        allsets = [x for x, _ in pairs] + [y for _, y in pairs]115        De = float(np.mean([masi(allsets[i], allsets[j]) ** 2 for i in range(len(allsets)) for j in range(i + 1, len(allsets))])) if len(allsets) <= 1200 else float(np.mean([masi(allsets[i], allsets[j]) ** 2 for i, j in np.random.default_rng(44).integers(0, len(allsets), size=(200000, 2)) if i != j]))116        M[f"krippendorff_alpha_masi_{a}_vs_{b}"] = 1.0 - Do / De if De > 0 else None117        # Jaccard agreement on multi-label sets between raters118        M[f"mean_jaccard_{a}_vs_{b}"] = float(np.mean([len(set(A[w]["categories"]) & set(B[w]["categories"])) / max(len(set(A[w]["categories"]) | set(B[w]["categories"])), 1) for w in ids]))119        dis = [{"work_id": w, "title": df.set_index("work_id").title[w], "rules": sorted(rule[w]), "rule_primary": rule_p[w], a: sorted(A[w]["categories"]), b: sorted(B[w]["categories"]),120                "pi_review": bool(df.set_index("work_id").pi_review[w])} for w in ids if not (set(A[w]["categories"]) & set(B[w]["categories"]) & rule[w]) and (rule[w] or A[w]["categories"] or B[w]["categories"])]121        pd.DataFrame(dis).to_csv(f"{d}/disagreements.csv", index=False); M["n_disagreements"] = len(dis)122    json.dump(M, open(f"{d}/metrics.json", "w"), indent=1)123    print({k: v for k, v in M.items() if not isinstance(v, dict)}); print({k: (v["micro_f1"], v["macro_f1"]) for k, v in M.items() if isinstance(v, dict) and "micro_f1" in v})124125126def main():127    ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--ontology", default=None)128    ap.add_argument("--draw", action="store_true"); ap.add_argument("--rate", action="store_true"); ap.add_argument("--metrics", action="store_true")129    ap.add_argument("--model", default="claude-haiku-5-5"); ap.add_argument("--env-file", default="~/wp44_uqo/.env"); ap.add_argument("--seed", type=int, default=44)130    ap.add_argument("--per-cat", type=int, default=50); ap.add_argument("--n-none", type=int, default=200); ap.add_argument("--batch", type=int, default=10)131    a = ap.parse_args(); ont = a.ontology or os.path.join(a.root, "taxonomy/ontology.yaml")132    if a.draw: draw(a.root, a.seed, a.per_cat, a.n_none)133    if a.rate: rate(a.root, a.model, a.env_file, ont, a.batch)134    if a.metrics: metrics(a.root, ont)135136137if __name__ == "__main__":138    main()139
Files · a549246b6e6d148c37f1b3b1d15a674cd2828d88
Menu

Theme