1#!/usr/bin/env python32"""WP44 — Phase 3: schema-constrained LLM extraction of method / model / dataset / benchmark / claim nodes and typed edges3(introduces_method, uses_dataset, evaluates_on, extends, contradicts, reproduces, addresses_problem) for the high-priority set,4from title + abstract (evidence_type = abstract) — full text is NOT used unless a lawful OA text is attached later.5Every edge carries provenance (work_id + quoted span from the abstract), confidence (model-reported, 0–1) and6verification_status = 'unverified' (machine-verified when the quoted span is found verbatim in the abstract; human-verified7only after PI review). Model: claude-sonnet-5-5 (structured outputs, effort low), batches of 5 works, resumable JSONL.8Budget: see protocol/amendments.csv (A-004)."""9import argparse, json, os, time, pandas as pd1011SCHEMA = {"type": "object", "properties": {"works": {"type": "array", "items": {"type": "object", "properties": {12 "work_id": {"type": "integer"},13 "contribution_type": {"type": "string", "enum": ["theory", "method", "empirical", "benchmark", "engineering", "survey", "other"]},14 "entities": {"type": "array", "items": {"type": "object", "properties": {15 "name": {"type": "string"}, "kind": {"type": "string", "enum": ["method", "model", "dataset", "benchmark", "problem", "metric"]},16 "role": {"type": "string", "enum": ["introduced", "used", "evaluated_on", "extended", "compared_against", "addressed"]},17 "span": {"type": "string"}, "confidence": {"type": "number"}}, "required": ["name", "kind", "role", "span", "confidence"], "additionalProperties": False}},18 "claims": {"type": "array", "items": {"type": "object", "properties": {"claim": {"type": "string"}, "claim_kind": {"type": "string", "enum": ["result", "improvement", "negative_result", "reproduction", "contradiction"]},19 "span": {"type": "string"}, "confidence": {"type": "number"}}, "required": ["claim", "claim_kind", "span", "confidence"], "additionalProperties": False}}},20 "required": ["work_id", "contribution_type", "entities", "claims"], "additionalProperties": False}}}, "required": ["works"], "additionalProperties": False}21SYSTEM = ("You extract structured knowledge from computer-science paper abstracts for a knowledge graph. For each work: the contribution type; the named "22 "methods, models, datasets, benchmarks, metrics and research problems with their role; and the main claims. Every entity and claim must quote a "23 "short verbatim span (<= 25 words) from the abstract that supports it. Use canonical names (e.g. 'ImageNet', 'BERT', 'Transformer'). Never infer "24 "anything not stated in the abstract. Confidence is your calibrated probability that the extraction is correct.")252627def main():28 ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--model", default="claude-sonnet-5-5"); ap.add_argument("--env-file", default="~/wp44_uqo/.env")29 ap.add_argument("--batch", type=int, default=5); ap.add_argument("--limit", type=int, default=0); ap.add_argument("--max-usd", type=float, default=150.0)30 ap.add_argument("--shard", default="0/1", help="i/N: this process handles rows with index %% N == i (parallel runs write separate JSONL files)"); a = ap.parse_args(); R = a.root31 si, sn = map(int, a.shard.split("/"))32 import anthropic, duckdb33 key = next(l.split("=", 1)[1].strip() for l in open(os.path.expanduser(a.env_file)) if l.startswith("ANTHROPIC_API_KEY="))34 client = anthropic.Anthropic(api_key=key, max_retries=5)35 con = duckdb.connect(f"{R}/data/processed/wp44.duckdb", read_only=True)36 hp = pd.read_parquet(f"{R}/literature_synthesis/landmarks/high_priority.parquet")37 ab = con.execute(f"SELECT work_id, abstract FROM oa_abstracts WHERE work_id IN (SELECT work_id FROM read_parquet('{R}/literature_synthesis/landmarks/high_priority.parquet'))").df()38 hp = hp.merge(ab, on="work_id"); hp = hp.sort_values(["score"], ascending=False)39 if a.limit: hp = hp.head(a.limit)40 out_dir = f"{R}/data/processed/graph/extraction"; os.makedirs(out_dir, exist_ok=True); out_p = f"{out_dir}/extractions_{a.model}.jsonl" if sn == 1 else f"{out_dir}/extractions_{a.model}.shard{si}.jsonl"41 import glob as _g42 done = set()43 for f in _g.glob(f"{out_dir}/extractions_{a.model}*.jsonl"):44 done |= {json.loads(l)["work_id"] for l in open(f)}45 hp = hp.reset_index(drop=True); hp = hp[hp.index % sn == si]46 todo = hp[~hp.work_id.isin(done)]; cost = 0.0; price_in, price_out, price_cache = 2.0 / 1e6, 10.0 / 1e6, 0.2 / 1e6; n_calls = 047 with open(out_p, "a") as fh:48 for i in range(0, len(todo), a.batch):49 if cost >= a.max_usd:50 print("budget reached", cost); break51 b = todo.iloc[i:i + a.batch]52 items = "\n\n".join(f"### work_id {r.work_id}\nTitle: {r.title}\nYear: {r.year}\nAbstract: {(r.abstract or '')[:3000]}" for r in b.itertuples())53 for attempt in range(4):54 try:55 resp = client.messages.create(model=a.model, max_tokens=6000, system=[{"type": "text", "text": SYSTEM, "cache_control": {"type": "ephemeral"}}],56 messages=[{"role": "user", "content": items}], output_config={"format": {"type": "json_schema", "schema": SCHEMA}, "effort": "low"}); break57 except anthropic.RateLimitError:58 time.sleep(15 * (attempt + 1))59 if resp.stop_reason == "refusal":60 continue61 u = resp.usage; cost += u.input_tokens * price_in + u.output_tokens * price_out + (u.cache_read_input_tokens or 0) * price_cache; n_calls += 162 data = json.loads(next(bk.text for bk in resp.content if bk.type == "text"))63 abs_by_id = {int(r.work_id): (r.abstract or "") for r in b.itertuples()}64 for w in data["works"]:65 txt = abs_by_id.get(w["work_id"], "")66 for e in w["entities"] + w["claims"]:67 e["verification_status"] = "machine-verified" if e["span"] and e["span"].strip().lower() in txt.lower() else "unverified"68 w["model"] = a.model; w["evidence_type"] = "abstract"; fh.write(json.dumps(w) + "\n")69 fh.flush()70 if n_calls % 25 == 0:71 print(f"{n_calls} calls, {i + len(b)}/{len(todo)} works, ${cost:.2f}", flush=True)72 json.dump({"model": a.model, "calls": n_calls, "cost_usd": round(cost, 2), "works_done": len(done) + min(len(todo), i + len(b) if len(todo) else 0)}, open(f"{out_dir}/extraction_usage_{a.model}.json", "w"), indent=1)73 print("done", n_calls, "calls", round(cost, 2), "USD")747576if __name__ == "__main__":77 main()78