diff --git a/embedding_server.py b/embedding_server.py index 70d07d1..b2b1ae6 100644 --- a/embedding_server.py +++ b/embedding_server.py @@ -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} diff --git a/scripts/lightrag_ingest/build_section_docs.py b/scripts/lightrag_ingest/build_section_docs.py new file mode 100644 index 0000000..a6ff084 --- /dev/null +++ b/scripts/lightrag_ingest/build_section_docs.py @@ -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() diff --git a/scripts/lightrag_ingest/ingest_adc_structured.py b/scripts/lightrag_ingest/ingest_adc_structured.py index dd31318..2a93289 100644 --- a/scripts/lightrag_ingest/ingest_adc_structured.py +++ b/scripts/lightrag_ingest/ingest_adc_structured.py @@ -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() diff --git a/scripts/lightrag_ingest/rebuild_vdb.py b/scripts/lightrag_ingest/rebuild_vdb.py new file mode 100644 index 0000000..e7f407f --- /dev/null +++ b/scripts/lightrag_ingest/rebuild_vdb.py @@ -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()