WP44: parallel shards for knowledge-graph extraction
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
1 changed file +10 −4
modified
knowledge_graph/03_extract_entities.py
+10 −4
@@ -26,7 +26,9 @@ SYSTEM = ("You extract structured knowledge from computer-science paper abstract | ||
| 26 | 26 | |
| 27 | 27 | def main(): |
| 28 | 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); a = ap.parse_args(); R = a.root | |
| 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.root | |
| 31 | + si, sn = map(int, a.shard.split("/")) | |
| 30 | 32 | import anthropic, duckdb |
| 31 | 33 | key = next(l.split("=", 1)[1].strip() for l in open(os.path.expanduser(a.env_file)) if l.startswith("ANTHROPIC_API_KEY=")) |
| 32 | 34 | client = anthropic.Anthropic(api_key=key, max_retries=5) |
@@ -35,8 +37,12 @@ def main(): | ||
| 35 | 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() |
| 36 | 38 | hp = hp.merge(ab, on="work_id"); hp = hp.sort_values(["score"], ascending=False) |
| 37 | 39 | if a.limit: hp = hp.head(a.limit) |
| 38 | − out_dir = f"{R}/data/processed/graph/extraction"; os.makedirs(out_dir, exist_ok=True); out_p = f"{out_dir}/extractions_{a.model}.jsonl" | |
| 39 | − done = {json.loads(l)["work_id"] for l in open(out_p)} if os.path.exists(out_p) else set() | |
| 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 _g | |
| 42 | + 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] | |
| 40 | 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 = 0 |
| 41 | 47 | with open(out_p, "a") as fh: |
| 42 | 48 | for i in range(0, len(todo), a.batch): |
@@ -46,7 +52,7 @@ def main(): | ||
| 46 | 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()) |
| 47 | 53 | for attempt in range(4): |
| 48 | 54 | try: |
| 49 | − resp = client.messages.create(model=a.model, max_tokens=12000, system=[{"type": "text", "text": SYSTEM, "cache_control": {"type": "ephemeral"}}], | |
| 55 | + resp = client.messages.create(model=a.model, max_tokens=6000, system=[{"type": "text", "text": SYSTEM, "cache_control": {"type": "ephemeral"}}], | |
| 50 | 56 | messages=[{"role": "user", "content": items}], output_config={"format": {"type": "json_schema", "schema": SCHEMA}, "effort": "low"}); break |
| 51 | 57 | except anthropic.RateLimitError: |
| 52 | 58 | time.sleep(15 * (attempt + 1)) |
| 53 | 59 | |