wjq/PlantAI-rag-new
1
0
Fork 0
forked from admin/PlantAI-rag-new
Code Pull Requests Activity

fix(aqua): batch_runner path sanitize, loader batch embed, alias expansion

Browse Source 
- 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>
This commit is contained in:
admin
2026-09-13 01:06:38 +08:00
co-authored by Claude Fable 5
parent 688edba169
commit f0a1711ebc
3 changed files with 174 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
kb-libraries/aqua/configs/entity_alias_aqua.yaml
+56
View File
@@ -5,11 +5,15 @@ Litopenaeus vannamei:
- L. vannamei
- Penaeus vannamei
- whiteleg shrimp
- white leg shrimp
- Pacific white shrimp
- Pacific whiteleg shrimp
- white shrimp
- vannamei shrimp
- 凡纳滨对虾
- 凡纳对虾
- 南美白对虾
- 白对虾
- 对虾
forbidden_expansions:
- Penaeus monodon
@@ -45,6 +49,11 @@ Oreochromis niloticus:
- tilapia
- 罗非鱼
- 尼罗罗非鱼
- 红罗非鱼
forbidden_expansions:
- Oreochromis mossambicus
- O. mossambicus
- 莫桑比克罗非鱼
Oncorhynchus mykiss:
type: species
@@ -55,6 +64,7 @@ Oncorhynchus mykiss:
- steelhead trout
- steelhead
- 虹鳟
- 虹鳟鱼
Salmo salar:
type: species
@@ -64,6 +74,7 @@ Salmo salar:
- Atlantic salmon
- salmon
- 大西洋鲑
- 大西洋鲑鱼
- 三文鱼
forbidden_expansions:
- Oncorhynchus
@@ -176,6 +187,51 @@ Apostichopus japonicus:
- 仿刺参
- 海参
Streptococcus agalactiae:
type: species
aliases:
- Streptococcus agalactiae
- S. agalactiae
- group B Streptococcus
- group B streptococcus
- GBS
- 无乳链球菌
- B族链球菌
forbidden_expansions:
- Guillain-Barré syndrome
- Guillain-Barre syndrome
- 格林巴利
Danio rerio:
type: species
aliases:
- Danio rerio
- D. rerio
- Brachydanio rerio
- zebrafish
- zebra fish
- zebra danio
- 斑马鱼
RAS:
type: culture_system
retrieval_label: recirculating aquaculture system
aliases:
- RAS
- recirculating aquaculture system
- recirculating aquaculture systems
- recirculation aquaculture system
- 循环水养殖系统
- 循环水养殖
- 工厂化循环水养殖
forbidden_expansions:
- Renin-angiotensin system
- renin angiotensin system
- RAS-related protein
- Ras GTPase
- RAS signaling
- reticular activating system
growth performance:
type: phenotype
aliases:
tooling/aqua_kb/batch_runner.py
+15 -1
View File
@@ -28,6 +28,7 @@ from __future__ import annotations
import argparse
import json
import re
import subprocess
import sys
import time
@@ -37,6 +38,15 @@ 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()
@@ -99,12 +109,16 @@ class Runner:
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
key = self.paper_key(entry)
# 默认：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 ----------
tooling/aqua_kb/qdrant_loader.py
+103 -21
View File
@@ -7,6 +7,8 @@
嵌入 text 字段，分批 upsert 到 Qdrant collection。
特性：
- 批量嵌入：input 传数组（--embed-batch 条/批），一次请求嵌入一批文本，
再按批 upsert（实测单条约 2.7 texts/s → 批量约 37 texts/s，约 14x 吞吐）
- checkpoint JSONL：每成功 upsert 一个 assertion_id 记一行 → 断点续跑（kill 后重跑跳过已传）
- 同 id 重放覆盖（Qdrant upsert 语义 + UUID 确定性 id → 幂等）
- 确定性处理顺序（文件路径 + 文件内行序排序）
@@ -15,7 +17,8 @@
python qdrant_loader.py --points <dir|file> --qdrant http://127.0.0.1:27435 \
--collection aqua_srt_v1_20260930_normalized --batch 200 \
--checkpoint <file.jsonl> \
[--embed-url http://127.0.0.1:11434/api/embed] [--embed-model bge-m3:567m]
[--embed-url http://127.0.0.1:11434/api/embed] [--embed-model bge-m3:567m] \
[--embed-batch 32]
"""
from __future__ import annotations
@@ -58,25 +61,54 @@ class EmbedClient:
self.model = model
self.dim: int | None = None
@staticmethod
def _extract_vector(item: Any) -> List[float] | None:
"""从 embedding 数组元素中提取向量（兼容 embedding 对象与裸数组两种形态）。"""
if isinstance(item, dict):
item = item.get("embedding")
if not isinstance(item, (list, tuple)):
return None
try:
return [float(value) for value in item]
except (TypeError, ValueError):
return None
def embed(self, text: str) -> List[float]:
"""单条嵌入（预热 probe / 失败回退路径）。"""
vector = self.embed_batch([text])
if not vector:
raise RuntimeError("embedding endpoint returned no vector")
return vector[0]
def embed_batch(self, texts: List[str]) -> List[List[float]]:
"""批量嵌入：input 传数组，一次请求嵌入多条（约 60 texts/s）。"""
if not texts:
return []
payload = {
"model": self.model,
"input": [t[:EMBED_INPUT_LIMIT] for t in texts],
}
for attempt in range(4):
try:
response = requests.post(
self.url,
json={"model": self.model, "input": text[:EMBED_INPUT_LIMIT]},
timeout=60,
)
response = requests.post(self.url, json=payload, timeout=120)
response.raise_for_status()
data = response.json()
embeddings = data.get("embeddings") or []
if embeddings:
vector = embeddings[0]
elif "embedding" in data:
vector = data["embedding"]
else:
raise RuntimeError("embedding endpoint returned no vector")
self.dim = len(vector)
return [float(value) for value in vector]
raw = data.get("embeddings")
if not isinstance(raw, list) or not raw:
raise RuntimeError("embedding endpoint returned no vectors")
vectors = [self._extract_vector(item) for item in raw]
if any(v is None for v in vectors):
raise RuntimeError(
f"embedding endpoint returned {len(vectors)} vectors "
f"for {len(texts)} inputs (bad shapes)"
)
if len(vectors) != len(texts):
raise RuntimeError(
f"embedding endpoint returned {len(vectors)} vectors "
f"for {len(texts)} inputs"
)
self.dim = len(vectors[0])
return vectors
except Exception as exc:
if attempt == 3:
raise
@@ -122,6 +154,29 @@ def load_checkpoint(checkpoint: Path) -> set:
return done
def embed_group(
embed: "EmbedClient", records: List[Dict[str, Any]], embed_batch: int
) -> List[List[float] | None]:
"""按 embed_batch 条/批批量嵌入 records；单条失败时按单条回退重试，仍失败则该条置 None。"""
vectors: List[List[float] | None] = [None] * len(records)
for start in range(0, len(records), embed_batch):
group = records[start:start + embed_batch]
try:
vectors[start:start + embed_batch] = embed.embed_batch(
[r.get("text") or "" for r in group]
)
except Exception as exc:
print(f"[loader] WARN embed batch failed ({exc}); retrying singly",
file=sys.stderr)
for i, record in enumerate(group):
idx = start + i
try:
vectors[idx] = embed.embed(record.get("text") or "")
except Exception:
vectors[idx] = None
return vectors
def upsert_file(
points_path: Path,
qdrant_url: str,
@@ -130,8 +185,13 @@ def upsert_file(
*,
checkpoint: Path,
embed: "EmbedClient",
embed_batch: int = 32,
) -> Dict[str, int]:
"""上传单个 qdrant_points.jsonl（断点续跑，checkpoint 追加）。"""
"""上传单个 qdrant_points.jsonl（断点续跑，checkpoint 追加）。
嵌入按 embed_batch 条/批批量调用 /api/embed，再攒 batch_size 条 upsert；
批内单条失败时逐条回退重试。checkpoint 记录口径不变（assertion_id）。
"""
checkpoint.parent.mkdir(parents=True, exist_ok=True)
done = load_checkpoint(checkpoint)
qdrant = qdrant_url.rstrip("/")
@@ -139,6 +199,7 @@ def upsert_file(
total = skipped = uploaded = 0
batch: List[Dict[str, Any]] = []
pending: List[Dict[str, Any]] = [] # 待嵌入的未 checkpoint 记录
started = time.time()
with open(checkpoint, "a", encoding="utf-8") as checkpoint_handle:
def flush_batch() -> None:
@@ -154,17 +215,35 @@ def upsert_file(
checkpoint_handle.flush()
batch = []
def flush_pending() -> None:
nonlocal pending, uploaded
if not pending:
return
vectors = embed_group(embed, pending, embed_batch)
for record, vector in zip(pending, vectors):
if vector is None:
raise RuntimeError(
f"embed failed for assertion "
f"{record['payload']['assertion_id']} "
f"(batch and single retries exhausted)"
)
record["vector"] = vector
batch.append(record)
uploaded += 1
if len(batch) >= batch_size:
flush_batch()
pending = []
for record in iter_points([points_path]):
total += 1
assertion_id = record["payload"]["assertion_id"]
if assertion_id in done:
skipped += 1
continue
record["vector"] = embed.embed(record.get("text") or "")
batch.append(record)
uploaded += 1
if len(batch) >= batch_size:
flush_batch()
pending.append(record)
if len(pending) >= embed_batch:
flush_pending()
flush_pending()
flush_batch()
return {
"uploaded": uploaded,
@@ -183,6 +262,8 @@ def main() -> None:
parser.add_argument("--checkpoint", required=True, help="JSONL of uploaded assertion_ids")
parser.add_argument("--embed-url", default=DEFAULT_EMBED_URL)
parser.add_argument("--embed-model", default=DEFAULT_EMBED_MODEL)
parser.add_argument("--embed-batch", type=int, default=32,
help="texts per /api/embed request (batch embedding)")
args = parser.parse_args()
points_arg = Path(args.points)
@@ -210,6 +291,7 @@ def main() -> None:
args.batch,
checkpoint=checkpoint,
embed=embed,
embed_batch=args.embed_batch,
)
for key in grand:
grand[key] += result[key]