Build deduplicated section documents and harden LightRAG ingestion

- build_section_docs.py: indications, contraindications, warnings and interactions
  only, one representative per active-substance set
- ingest: --dedupe, --retry-failed
- embedding server: bounded PyTorch allocator cache and serialized encoding
  (fixes memory growth that crashed long ingests)
- rebuild_vdb.py: recompute vector stores from the graph after corruption
This commit is contained in:
Oleh Poiasnik 2026-11-02 16:14:32 +01:00
parent e1fdd93261
commit 83ec1aca24
4 changed files with 375 additions and 9 deletions

View File

@ -10,9 +10,17 @@
-d '{"model": "local-embed", "input": "test"}'
"""
import os
# Без этого кэш CPU-аллокатора PyTorch растёт с каждой новой длиной входа
# (за ночь сервер дорос до ~45 ГБ виртуальной памяти и упал).
os.environ.setdefault("LRU_CACHE_CAPACITY", "1")
import threading
import time
import json
from fastapi import FastAPI, Request
from fastapi.concurrency import run_in_threadpool
from fastapi.responses import JSONResponse
import uvicorn
from sentence_transformers import SentenceTransformer
@ -27,6 +35,15 @@ print(f"Модель загружена. Размерность: {EMBED_DIM}")
app = FastAPI(title="Local Embedding Server")
# Одновременно считается только один батч: параллельные encode() в потоках
# не ускоряют CPU-инференс, но умножают потребление памяти.
_encode_lock = threading.Lock()
def _encode(texts):
with _encode_lock:
return model.encode(texts, normalize_embeddings=True)
@app.get("/health")
def health():
@ -41,7 +58,8 @@ async def _handle_embeddings(request: Request):
else:
texts = inp
vecs = model.encode(texts, normalize_embeddings=True).tolist()
# encode() блокирует; в пуле потоков сервер продолжает принимать запросы
vecs = (await run_in_threadpool(_encode, texts)).tolist()
data = [
{"object": "embedding", "index": i, "embedding": vec}

View File

@ -0,0 +1,135 @@
"""Build short LightRAG documents from selected ADC leaflet sections.
Instead of the full leaflet, each document contains only the parts relevant
for the agent: indications, contraindications, warnings and interactions.
Products are deduplicated by their active substance set (see
ingest_adc_structured.dedupe_key), keeping the product with the richest text.
The output is a JSON list in the same shape that ingest_adc_structured.py
expects (`urls`, `product`, `graph_hints`, `lightrag_text`), so it can be
uploaded with:
python scripts/lightrag_ingest/ingest_adc_structured.py \
--input outputs/lightrag_sections/adc_section_docs.json \
--progress outputs/lightrag_sections/progress.json --dedupe none --limit 100000
"""
from __future__ import annotations
import argparse
import json
import sys
from collections import Counter
from pathlib import Path
from typing import Any
sys.path.insert(0, str(Path(__file__).resolve().parent))
from ingest_adc_structured import DEFAULT_INPUT, dedupe_key, load_json, write_json # noqa: E402
DEFAULT_OUTPUT = Path("outputs/lightrag_sections/adc_section_docs.json")
# (heading in the document, where to read it in `pil`, fallback key in `graph_hints`)
SECTIONS = (
("Indikácie (na čo sa liek používa)", ("sections", "what_is_it"), "indications_text"),
("Kontraindikácie (kedy liek neužívať)", ("subsections", "contraindications"), "contraindications_text"),
("Upozornenia a opatrenia", ("subsections", "warnings"), "warnings_text"),
("Interakcie s inými liekmi", ("subsections", "interactions"), "interactions_text"),
)
if hasattr(sys.stdout, "reconfigure"):
sys.stdout.reconfigure(encoding="utf-8")
def section_text(record: dict[str, Any], location: tuple[str, str], fallback: str) -> str:
pil = record.get("pil") or {}
group = pil.get(location[0]) if isinstance(pil, dict) else None
text = group.get(location[1]) if isinstance(group, dict) else None
if not text:
text = (record.get("graph_hints") or {}).get(fallback)
return str(text or "").strip()
def build_document(record: dict[str, Any]) -> tuple[str, int]:
"""Return the document text and how many clinical sections it contains."""
product = record.get("product") or {}
hints = record.get("graph_hints") or {}
urls = record.get("urls") or {}
substances = hints.get("active_substances") or []
if isinstance(substances, str):
substances = [substances]
classes = [
str(item.get("name"))
for item in hints.get("classification_codes") or []
if isinstance(item, dict) and item.get("name")
]
lines = [f"Liek: {hints.get('drug') or product.get('name') or ''}"]
if substances:
lines.append(f"Účinná látka: {', '.join(str(s) for s in substances)}")
if hints.get("dosage_form"):
lines.append(f"Lieková forma: {hints['dosage_form']}")
if classes:
lines.append(f"Skupina: {' > '.join(classes)}")
if urls.get("pil"):
lines.append(f"Zdroj: {urls['pil']}")
found = 0
for heading, location, fallback in SECTIONS:
text = section_text(record, location, fallback)
if text:
found += 1
lines.append(f"\n## {heading}\n{text}")
return "\n".join(lines), found
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
parser.add_argument("--input", type=Path, default=DEFAULT_INPUT)
parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
parser.add_argument("--dedupe", choices=("none", "text", "substance"), default="substance")
args = parser.parse_args()
print(f"Loading {args.input} ...")
records = load_json(args.input)
stats: Counter = Counter(input=len(records))
groups: dict[str, tuple[int, dict[str, Any]]] = {}
for record in records:
text, found = build_document(record)
# Keep only products with at least one of contraindications / warnings / interactions.
if found == 0 or not any(
section_text(record, location, fallback) for _, location, fallback in SECTIONS[1:]
):
stats["no_clinical_sections"] += 1
continue
doc = {
"urls": record.get("urls") or {},
"product": {"name": (record.get("product") or {}).get("name")},
"graph_hints": {
key: (record.get("graph_hints") or {}).get(key)
for key in ("drug", "active_substances", "contraindications_text",
"warnings_text", "interactions_text", "side_effects_text")
},
"lightrag_text": text,
}
key = dedupe_key(record, args.dedupe)
if key in groups:
stats["duplicate"] += 1
if len(text) <= groups[key][0]:
continue
groups[key] = (len(text), doc)
docs = [doc for _, doc in groups.values()]
write_json(args.output, docs)
total_chars = sum(len(doc["lightrag_text"]) for doc in docs)
stats["documents"] = len(docs)
print(f"Stats: {dict(stats)}")
print(f"Average document length: {total_chars // max(len(docs), 1)} chars, total {total_chars / 1e6:.1f} M chars")
print(f"Saved: {args.output}")
if __name__ == "__main__":
main()

View File

@ -98,28 +98,76 @@ def is_clinical(record: dict[str, Any]) -> bool:
)
CLINICAL_KEYS = (
"contraindications_text",
"interactions_text",
"warnings_text",
"side_effects_text",
)
def clinical_text_length(record: dict[str, Any]) -> int:
hints = record.get("graph_hints") or {}
return sum(len(str(hints.get(key) or "")) for key in CLINICAL_KEYS)
def dedupe_key(record: dict[str, Any], mode: str) -> str:
"""Group key for records that describe the same medicine.
`substance` groups products by their set of active substances (different
packages and brands of the same substance share one leaflet content);
products without known substances fall back to the clinical text itself.
`text` groups only products whose clinical texts are identical.
"""
hints = record.get("graph_hints") or {}
if mode == "substance":
substances = hints.get("active_substances") or []
if isinstance(substances, str):
substances = [substances]
names = sorted({str(item).strip().lower() for item in substances if str(item).strip()})
if names:
return "substance:" + "|".join(names)
if mode in {"substance", "text"}:
texts = [str(hints.get(key) or "").strip() for key in CLINICAL_KEYS]
return "text:" + json.dumps(texts, ensure_ascii=False)
return "source:" + source_id(record)
def iter_selected_records(
records: list[dict[str, Any]],
progress: dict[str, Any],
clinical_only: bool,
dedupe: str = "none",
retry_failed: bool = False,
) -> tuple[list[dict[str, Any]], Counter]:
done = set(progress.get("done", []))
failed = set(progress.get("failed", []))
failed = set() if retry_failed else set(progress.get("failed", []))
skipped: Counter = Counter()
selected: list[dict[str, Any]] = []
candidates: list[dict[str, Any]] = []
for record in records:
source = source_id(record)
if source in done or source in failed:
skipped["already_processed"] += 1
continue
if clinical_only and not is_clinical(record):
skipped["not_clinical"] += 1
continue
if not str(record.get("lightrag_text") or "").strip():
skipped["missing_lightrag_text"] += 1
continue
selected.append(record)
candidates.append(record)
# Group duplicates; a group already covered by an uploaded record is skipped.
groups: dict[str, list[dict[str, Any]]] = {}
for record in candidates:
groups.setdefault(dedupe_key(record, dedupe), []).append(record)
selected: list[dict[str, Any]] = []
for members in groups.values():
sources = {source_id(record) for record in members}
if sources & (done | failed):
skipped["already_processed"] += len(members)
continue
# Representative: the record with the richest clinical text.
selected.append(max(members, key=clinical_text_length))
skipped["duplicate"] += len(members) - 1
return selected, skipped
@ -150,7 +198,13 @@ def ingest(args: argparse.Namespace) -> None:
raise SystemExit(f"Input must be a JSON list: {args.input}")
progress = load_progress(args.progress, resume=args.resume)
selected, skipped = iter_selected_records(records, progress, clinical_only=not args.all_records)
selected, skipped = iter_selected_records(
records,
progress,
clinical_only=not args.all_records,
dedupe=args.dedupe,
retry_failed=args.retry_failed,
)
progress["skipped"] = Counter(progress.get("skipped", {})) + skipped
if args.limit is not None:
@ -229,6 +283,13 @@ def main() -> None:
parser.add_argument("--save-every", type=int, default=10)
parser.add_argument("--resume", action="store_true", help="Resume from the progress file.")
parser.add_argument("--all-records", action="store_true", help="Upload all records, not only clinical ones.")
parser.add_argument(
"--dedupe",
choices=("none", "text", "substance"),
default="substance",
help="Upload one representative per active substance set (default) or per identical clinical text.",
)
parser.add_argument("--retry-failed", action="store_true", help="Upload again records marked as failed.")
parser.add_argument("--dry-run", action="store_true", help="Show selected records without uploading.")
parser.add_argument("--status", action="store_true", help="Show local progress and LightRAG pipeline status.")
args = parser.parse_args()

View File

@ -0,0 +1,152 @@
"""Rebuild LightRAG vector stores from the knowledge graph and text chunks.
The vector files (vdb_entities.json, vdb_relationships.json, vdb_chunks.json)
are derived data: every record is an embedding of an entity or relation stored
in graph_chunk_entity_relation.graphml, or of a chunk in
kv_store_text_chunks.json. If they get corrupted (e.g. the process
ran out of memory while saving), they can be recomputed locally without the
LLM. Records use the same ids, fields and content format as LightRAG
(lightrag/operate.py), so the server loads them as if it wrote them itself.
Stop the LightRAG server before running this script.
python scripts/lightrag_ingest/rebuild_vdb.py --storage lightrag/rag_storage_qwen36
"""
from __future__ import annotations
import os
# Prevents the PyTorch CPU allocator cache from growing with varying input shapes.
os.environ.setdefault("LRU_CACHE_CAPACITY", "1")
import argparse
import base64
import json
import shutil
import sys
import time
import zlib
from pathlib import Path
import networkx as nx
import numpy as np
from nano_vectordb import NanoVectorDB
from sentence_transformers import SentenceTransformer
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "lightrag"))
from lightrag.utils import compute_mdhash_id # noqa: E402
MODEL_NAME = "paraphrase-multilingual-MiniLM-L12-v2" # same model as embedding_server.py
if hasattr(sys.stdout, "reconfigure"):
sys.stdout.reconfigure(encoding="utf-8")
def entity_records(graph: nx.Graph) -> dict[str, dict]:
records = {}
for name, data in graph.nodes(data=True):
records[compute_mdhash_id(name, prefix="ent-")] = {
"entity_name": name,
"content": f"{name}\n{data.get('description', '')}",
"source_id": data.get("source_id", ""),
"file_path": data.get("file_path", ""),
}
return records
def relation_records(graph: nx.Graph) -> dict[str, dict]:
records = {}
for u, v, data in graph.edges(data=True):
src, tgt = (u, v) if u <= v else (v, u)
keywords = data.get("keywords", "")
records[compute_mdhash_id(src + tgt, prefix="rel-")] = {
"src_id": src,
"tgt_id": tgt,
"content": f"{keywords}\t{src}\n{tgt}\n{data.get('description', '')}",
"source_id": data.get("source_id", ""),
"file_path": data.get("file_path", ""),
}
return records
def chunk_records(storage: Path) -> dict[str, dict]:
chunks = json.loads((storage / "kv_store_text_chunks.json").read_text(encoding="utf-8"))
return {
chunk_id: {
"full_doc_id": chunk.get("full_doc_id", ""),
"content": chunk.get("content", ""),
"file_path": chunk.get("file_path", ""),
}
for chunk_id, chunk in chunks.items()
}
def build_store(path: Path, records: dict[str, dict], model: SentenceTransformer, batch_size: int) -> None:
ids = list(records)
contents = [records[i]["content"] for i in ids]
print(f"Embedding {len(ids)} records for {path.name} ...")
started = time.time()
vectors = model.encode(
contents,
batch_size=batch_size,
normalize_embeddings=True,
show_progress_bar=True,
convert_to_numpy=True,
).astype(np.float32)
print(f" done in {time.time() - started:.0f}s")
now = int(time.time())
datas = []
for record_id, vector in zip(ids, vectors):
encoded = base64.b64encode(zlib.compress(vector.astype(np.float16).tobytes())).decode("utf-8")
datas.append({
"__id__": record_id,
"__created_at__": now,
**records[record_id],
"vector": encoded,
"__vector__": vector,
})
tmp = path.with_suffix(".rebuilt.json")
if tmp.exists():
tmp.unlink()
db = NanoVectorDB(vectors.shape[1], storage_file=str(tmp))
db.upsert(datas=datas)
db.save()
if path.exists():
backup = path.with_suffix(".broken.json")
shutil.move(str(path), str(backup))
print(f" old file moved to {backup.name}")
shutil.move(str(tmp), str(path))
print(f" saved {path}")
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
parser.add_argument("--storage", type=Path, default=ROOT / "lightrag" / "rag_storage_qwen36")
parser.add_argument("--batch-size", type=int, default=128)
parser.add_argument(
"--only", choices=("entities", "relationships", "chunks"), help="Rebuild only one store."
)
args = parser.parse_args()
model = SentenceTransformer(MODEL_NAME)
if args.only in (None, "entities", "relationships"):
graph_file = args.storage / "graph_chunk_entity_relation.graphml"
print(f"Loading graph {graph_file} ...")
graph = nx.read_graphml(graph_file)
print(f" {graph.number_of_nodes()} nodes, {graph.number_of_edges()} edges")
if args.only in (None, "entities"):
build_store(args.storage / "vdb_entities.json", entity_records(graph), model, args.batch_size)
if args.only in (None, "relationships"):
build_store(args.storage / "vdb_relationships.json", relation_records(graph), model, args.batch_size)
if args.only in (None, "chunks"):
build_store(args.storage / "vdb_chunks.json", chunk_records(args.storage), model, args.batch_size)
print("Finished.")
if __name__ == "__main__":
main()