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 | def main(): | 27 | def 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") | 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 | import anthropic, duckdb | 32 | import anthropic, duckdb | ||
| 31 | key = next(l.split("=", 1)[1].strip() for l in open(os.path.expanduser(a.env_file)) if l.startswith("ANTHROPIC_API_KEY=")) | 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 | client = anthropic.Anthropic(api_key=key, max_retries=5) | 34 | client = anthropic.Anthropic(api_key=key, max_retries=5) | ||
@@ -35,8 +37,12 @@ def main(): | |||||
| 35 | 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() | 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 | hp = hp.merge(ab, on="work_id"); hp = hp.sort_values(["score"], ascending=False) | 38 | hp = hp.merge(ab, on="work_id"); hp = hp.sort_values(["score"], ascending=False) | ||
| 37 | if a.limit: hp = hp.head(a.limit) | 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" | 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" | ||
| 39 | − done = {json.loads(l)["work_id"] for l in open(out_p)} if os.path.exists(out_p) else set() | 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 | 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 | 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 | with open(out_p, "a") as fh: | 47 | with open(out_p, "a") as fh: | ||
| 42 | for i in range(0, len(todo), a.batch): | 48 | for i in range(0, len(todo), a.batch): | ||
@@ -46,7 +52,7 @@ def main(): | |||||
| 46 | 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()) | 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 | for attempt in range(4): | 53 | for attempt in range(4): | ||
| 48 | try: | 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 | messages=[{"role": "user", "content": items}], output_config={"format": {"type": "json_schema", "schema": SCHEMA}, "effort": "low"}); break | 56 | messages=[{"role": "user", "content": items}], output_config={"format": {"type": "json_schema", "schema": SCHEMA}, "effort": "low"}); break | ||
| 51 | except anthropic.RateLimitError: | 57 | except anthropic.RateLimitError: | ||
| 52 | time.sleep(15 * (attempt + 1)) | 58 | time.sleep(15 * (attempt + 1)) | ||
| 53 | 59 | ||||