forked from admin/PlantAI-rag-new
adminandClaude Fable 5
f0a1711ebc
- batch_runner: sanitize paper_id to disk dir name (was skipping all papers whose paper_id contains :/()); keeps explicit result path compat; verified 5/5 no-skip on paren-DOI papers - qdrant_loader: batch embed via Ollama array input (32/batch, 14x: 2.7 -> 36.9 texts/s, 24.5k points ~10min vs 2.5h); per-batch fail falls back to single; checkpoint format unchanged - entity_alias_aqua: +Streptococcus agalactiae (GBS), RAS (forbidden_expansions), Danio rerio; merged vernacular/Chinese aliases for O. niloticus / O. mykiss / S. salar / L. vannamei; g02 gap entities now normalize correctly - note for P3: ingest manifest MUST carry paper sidecar (pilot manifest_pilot.json does not - would emit paper_id="result" dirt) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
254 lines
11 KiB
Python
254 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""tooling/aqua_kb/batch_runner.py —— 分批建库调度器
|
||
|
||
按 manifest 编排建库全链，全程断点续跑：
|
||
for each paper in manifest:
|
||
result.json 已存在 → 跳过抽取（外部/D2 负责抽取，本工具只调度后续步骤）
|
||
→ phase2b_converter.convert()（neo4j.cypher + qdrant_points.jsonl）
|
||
→ 攒批（--convert-batch 篇）后：
|
||
Neo4j：逐文件 cypher-shell 执行（--neo4j-shell 路径），成功记 checkpoint
|
||
Qdrant：qdrant_loader.load_files()（embed + upsert），成功记 checkpoint
|
||
→ 失败写 build/failed.jsonl（含错误），全部步骤可重跑
|
||
|
||
状态文件（<out>/state/，JSONL 每行一个 paper_id）：
|
||
converted.jsonl / neo4j_loaded.jsonl / qdrant_loaded.jsonl / failed.jsonl
|
||
|
||
用法：
|
||
python batch_runner.py --manifest pilot_manifest.json --result-root <dir> \
|
||
--out <build dir> --qdrant http://127.0.0.1:27435 \
|
||
--collection aqua_srt_v1_20260930_normalized \
|
||
--neo4j-shell /opt/neo4j-community-5.15.0/bin/cypher-shell \
|
||
--neo4j-uri bolt://127.0.0.1:19889 --neo4j-user neo4j \
|
||
--neo4j-password-file kb-libraries/aqua/data/neo4j/.auth \
|
||
--embed-url http://127.0.0.1:11434/api/embed --embed-model bge-m3:567m \
|
||
--convert-batch 50 [--limit 20]
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import re
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
import phase2b_converter
|
||
import qdrant_loader
|
||
|
||
# 与抽取链落盘一致的目录名净化：非 [0-9a-zA-Z._-] 一律替换为 `_`（保留 doi: 前缀）。
|
||
# 例：doi:10.1016/0044-8486(81)90060-0 → doi_10.1016_0044-8486_81_90060-0
|
||
_SANITIZE_RE = re.compile(r"[^0-9a-zA-Z._-]")
|
||
|
||
|
||
def sanitize_paper_key(value: str) -> str:
|
||
"""把 paper_id 规范成与磁盘一致的目录名形态（md_batch/<sanitized>/result.json）。"""
|
||
return _SANITIZE_RE.sub("_", str(value))
|
||
|
||
|
||
def load_ids(path: Path) -> set:
|
||
ids: set = set()
|
||
if path.exists():
|
||
with open(path, "r", encoding="utf-8") as handle:
|
||
for line in handle:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
ids.add(json.loads(line)) # 兼容 json.dumps 写入的带引号行
|
||
except json.JSONDecodeError:
|
||
ids.add(line)
|
||
return ids
|
||
|
||
|
||
def append_id(path: Path, value: str) -> None:
|
||
with open(path, "a", encoding="utf-8") as handle:
|
||
handle.write(json.dumps(value) + "\n")
|
||
|
||
|
||
class Runner:
|
||
def __init__(self, args: argparse.Namespace) -> None:
|
||
self.args = args
|
||
self.out = Path(args.out)
|
||
self.state_dir = self.out / "state"
|
||
self.state_dir.mkdir(parents=True, exist_ok=True)
|
||
self.converted_path = self.state_dir / "converted.jsonl"
|
||
self.neo4j_loaded_path = self.state_dir / "neo4j_loaded.jsonl"
|
||
self.qdrant_loaded_path = self.state_dir / "qdrant_loaded.jsonl"
|
||
self.failed_path = self.out / "failed.jsonl"
|
||
self.converted = load_ids(self.converted_path)
|
||
self.neo4j_loaded = load_ids(self.neo4j_loaded_path)
|
||
self.qdrant_loaded = load_ids(self.qdrant_loaded_path)
|
||
self._embed: Optional[qdrant_loader.EmbedClient] = None
|
||
|
||
@property
|
||
def embed(self) -> qdrant_loader.EmbedClient:
|
||
if self._embed is None:
|
||
self._embed = qdrant_loader.EmbedClient(
|
||
self.args.embed_url, self.args.embed_model
|
||
)
|
||
return self._embed
|
||
|
||
# ---------- manifest ----------
|
||
|
||
def read_manifest(self) -> List[Dict[str, Any]]:
|
||
with open(self.args.manifest, "r", encoding="utf-8") as handle:
|
||
manifest = json.load(handle)
|
||
entries = manifest.get("papers") if isinstance(manifest, dict) else manifest
|
||
if not isinstance(entries, list):
|
||
raise RuntimeError("manifest must be a dict with 'papers' list or a list")
|
||
if self.args.limit:
|
||
entries = entries[: self.args.limit]
|
||
return entries
|
||
|
||
@staticmethod
|
||
def paper_key(entry: Dict[str, Any]) -> str:
|
||
return str(entry.get("paper_id") or entry.get("id") or entry.get("doi")
|
||
or entry.get("result") or json.dumps(entry, sort_keys=True))
|
||
|
||
def resolve_result(self, entry: Dict[str, Any]) -> Path:
|
||
# 兼容 1：manifest 行显式给出 result 路径（相对 result-root 或绝对路径）
|
||
if entry.get("result"):
|
||
path = Path(entry["result"])
|
||
if not path.is_absolute():
|
||
path = Path(self.args.result_root) / path
|
||
return path
|
||
# 默认：result-root/<sanitized(paper_id)>/result.json
|
||
# sanitize 与抽取链落盘目录名一致（doi:/、/、() 等 → _），
|
||
# 直接用原始 paper_id 拼路径会因含 :/() 而找不到文件。
|
||
key = sanitize_paper_key(self.paper_key(entry))
|
||
return Path(self.args.result_root) / key / "result.json"
|
||
|
||
# ---------- steps ----------
|
||
|
||
def step_convert(self, key: str, result_path: Path, paper_json: Optional[Dict[str, Any]]) -> Dict[str, Any]:
|
||
"""转换幂等（产物确定性），checkpoint 只写一次。返回产物路径 stats。"""
|
||
stats = phase2b_converter.convert(
|
||
result_path, self.out, self.args.release_id, paper_json
|
||
)
|
||
if key not in self.converted:
|
||
append_id(self.converted_path, key)
|
||
print(f"[runner] converted {key}: {stats['assertion_count']} assertions")
|
||
return stats
|
||
|
||
def step_neo4j(self, key: str, cypher_path: Path) -> None:
|
||
if key in self.neo4j_loaded:
|
||
print(f"[runner] neo4j skip {key} (checkpoint)")
|
||
return
|
||
shell = self.args.neo4j_shell
|
||
auth = [f"--format=plain", f"--address={self.args.neo4j_uri}",
|
||
f"--username={self.args.neo4j_user}"]
|
||
with open(self.args.neo4j_password_file, "r", encoding="utf-8") as handle:
|
||
password = handle.read().strip().splitlines()[0]
|
||
command = [shell, *auth, f"--password={password}", "-f", str(cypher_path)]
|
||
proc = subprocess.run(command, capture_output=True, text=True, timeout=300)
|
||
if proc.returncode != 0:
|
||
raise RuntimeError(f"cypher-shell failed for {key}: {proc.stderr[-2000:]}")
|
||
append_id(self.neo4j_loaded_path, key)
|
||
print(f"[runner] neo4j loaded {key}")
|
||
|
||
def step_qdrant(self, key: str, points_path: Path) -> None:
|
||
if key in self.qdrant_loaded:
|
||
print(f"[runner] qdrant skip {key} (checkpoint)")
|
||
return
|
||
embed = self.embed
|
||
probe = embed.embed("aquaculture warm-up probe")
|
||
qdrant_loader.ensure_collection(
|
||
self.args.qdrant, self.args.collection, len(probe)
|
||
)
|
||
qdrant_loader.upsert_file(
|
||
points_path,
|
||
self.args.qdrant,
|
||
self.args.collection,
|
||
self.args.batch,
|
||
checkpoint=self.state_dir / "qdrant_checkpoint.jsonl",
|
||
embed=embed,
|
||
)
|
||
append_id(self.qdrant_loaded_path, key)
|
||
print(f"[runner] qdrant loaded {key}")
|
||
|
||
# ---------- main ----------
|
||
|
||
def run(self) -> None:
|
||
entries = self.read_manifest()
|
||
started = time.time()
|
||
stats = {"total": len(entries), "converted": 0, "neo4j": 0, "qdrant": 0,
|
||
"failed": 0}
|
||
for index, entry in enumerate(entries, 1):
|
||
key = self.paper_key(entry)
|
||
if key in self.failed_keys:
|
||
continue
|
||
paper_json = entry.get("paper") if isinstance(entry.get("paper"), dict) else None
|
||
try:
|
||
result_path = self.resolve_result(entry)
|
||
if not result_path.exists():
|
||
print(f"[runner] SKIP {key}: result.json missing "
|
||
f"({result_path}) — 抽取由外部完成")
|
||
continue
|
||
conv = self.step_convert(key, result_path, paper_json)
|
||
stats["converted"] += 1
|
||
cypher_path = Path(conv["cypher"])
|
||
points_path = Path(conv["points"])
|
||
self.step_neo4j(key, cypher_path)
|
||
stats["neo4j"] += 1
|
||
self.step_qdrant(key, points_path)
|
||
stats["qdrant"] += 1
|
||
if index % 10 == 0:
|
||
print(f"[runner] progress {index}/{stats['total']} "
|
||
f"({time.time() - started:.0f}s)")
|
||
except Exception as exc:
|
||
stats["failed"] += 1
|
||
self.record_failure(key, str(exc))
|
||
print(f"[runner] FAILED {key}: {exc}", file=sys.stderr)
|
||
print(f"[runner] DONE {json.dumps(stats, ensure_ascii=False)} in "
|
||
f"{time.time() - started:.1f}s")
|
||
|
||
@property
|
||
def failed_keys(self) -> set:
|
||
keys: set = set()
|
||
if self.failed_path.exists():
|
||
with open(self.failed_path, "r", encoding="utf-8") as handle:
|
||
for line in handle:
|
||
try:
|
||
keys.add(json.loads(line)["paper_key"])
|
||
except Exception:
|
||
continue
|
||
return keys
|
||
|
||
def record_failure(self, key: str, error: str) -> None:
|
||
with open(self.failed_path, "a", encoding="utf-8") as handle:
|
||
handle.write(json.dumps(
|
||
{"paper_key": key, "error": error,
|
||
"ts": time.strftime("%Y-%m-%dT%H:%M:%S%z")},
|
||
ensure_ascii=False,
|
||
) + "\n")
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(description="aqua KB batch runner")
|
||
parser.add_argument("--manifest", required=True)
|
||
parser.add_argument("--result-root", required=True)
|
||
parser.add_argument("--out", required=True, help="build output root")
|
||
parser.add_argument("--release-id", default="aqua_srt_v1_20260930_normalized")
|
||
parser.add_argument("--qdrant", required=True)
|
||
parser.add_argument("--collection", default="aqua_srt_v1_20260930_normalized")
|
||
parser.add_argument("--batch", type=int, default=200)
|
||
parser.add_argument("--neo4j-shell", default="cypher-shell")
|
||
parser.add_argument("--neo4j-uri", default="bolt://127.0.0.1:19889")
|
||
parser.add_argument("--neo4j-user", default="neo4j")
|
||
parser.add_argument("--neo4j-password-file", required=True)
|
||
parser.add_argument("--embed-url", default="http://127.0.0.1:11434/api/embed")
|
||
parser.add_argument("--embed-model", default="bge-m3:567m")
|
||
parser.add_argument("--convert-batch", type=int, default=50,
|
||
help="batch size for Neo4j/Qdrant loading (checkpoint per paper)")
|
||
parser.add_argument("--limit", type=int, default=None,
|
||
help="only process first N manifest entries")
|
||
args = parser.parse_args()
|
||
Runner(args).run()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|