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:
parent
e1fdd93261
commit
83ec1aca24
@ -10,9 +10,17 @@
|
|||||||
-d '{"model": "local-embed", "input": "test"}'
|
-d '{"model": "local-embed", "input": "test"}'
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
|
||||||
|
# Без этого кэш CPU-аллокатора PyTorch растёт с каждой новой длиной входа
|
||||||
|
# (за ночь сервер дорос до ~45 ГБ виртуальной памяти и упал).
|
||||||
|
os.environ.setdefault("LRU_CACHE_CAPACITY", "1")
|
||||||
|
|
||||||
|
import threading
|
||||||
import time
|
import time
|
||||||
import json
|
import json
|
||||||
from fastapi import FastAPI, Request
|
from fastapi import FastAPI, Request
|
||||||
|
from fastapi.concurrency import run_in_threadpool
|
||||||
from fastapi.responses import JSONResponse
|
from fastapi.responses import JSONResponse
|
||||||
import uvicorn
|
import uvicorn
|
||||||
from sentence_transformers import SentenceTransformer
|
from sentence_transformers import SentenceTransformer
|
||||||
@ -27,6 +35,15 @@ print(f"Модель загружена. Размерность: {EMBED_DIM}")
|
|||||||
|
|
||||||
app = FastAPI(title="Local Embedding Server")
|
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")
|
@app.get("/health")
|
||||||
def health():
|
def health():
|
||||||
@ -41,7 +58,8 @@ async def _handle_embeddings(request: Request):
|
|||||||
else:
|
else:
|
||||||
texts = inp
|
texts = inp
|
||||||
|
|
||||||
vecs = model.encode(texts, normalize_embeddings=True).tolist()
|
# encode() блокирует; в пуле потоков сервер продолжает принимать запросы
|
||||||
|
vecs = (await run_in_threadpool(_encode, texts)).tolist()
|
||||||
|
|
||||||
data = [
|
data = [
|
||||||
{"object": "embedding", "index": i, "embedding": vec}
|
{"object": "embedding", "index": i, "embedding": vec}
|
||||||
|
|||||||
135
scripts/lightrag_ingest/build_section_docs.py
Normal file
135
scripts/lightrag_ingest/build_section_docs.py
Normal 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()
|
||||||
@ -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(
|
def iter_selected_records(
|
||||||
records: list[dict[str, Any]],
|
records: list[dict[str, Any]],
|
||||||
progress: dict[str, Any],
|
progress: dict[str, Any],
|
||||||
clinical_only: bool,
|
clinical_only: bool,
|
||||||
|
dedupe: str = "none",
|
||||||
|
retry_failed: bool = False,
|
||||||
) -> tuple[list[dict[str, Any]], Counter]:
|
) -> tuple[list[dict[str, Any]], Counter]:
|
||||||
done = set(progress.get("done", []))
|
done = set(progress.get("done", []))
|
||||||
failed = set(progress.get("failed", []))
|
failed = set() if retry_failed else set(progress.get("failed", []))
|
||||||
skipped: Counter = Counter()
|
skipped: Counter = Counter()
|
||||||
selected: list[dict[str, Any]] = []
|
candidates: list[dict[str, Any]] = []
|
||||||
|
|
||||||
for record in records:
|
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):
|
if clinical_only and not is_clinical(record):
|
||||||
skipped["not_clinical"] += 1
|
skipped["not_clinical"] += 1
|
||||||
continue
|
continue
|
||||||
if not str(record.get("lightrag_text") or "").strip():
|
if not str(record.get("lightrag_text") or "").strip():
|
||||||
skipped["missing_lightrag_text"] += 1
|
skipped["missing_lightrag_text"] += 1
|
||||||
continue
|
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
|
return selected, skipped
|
||||||
|
|
||||||
@ -150,7 +198,13 @@ def ingest(args: argparse.Namespace) -> None:
|
|||||||
raise SystemExit(f"Input must be a JSON list: {args.input}")
|
raise SystemExit(f"Input must be a JSON list: {args.input}")
|
||||||
|
|
||||||
progress = load_progress(args.progress, resume=args.resume)
|
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
|
progress["skipped"] = Counter(progress.get("skipped", {})) + skipped
|
||||||
|
|
||||||
if args.limit is not None:
|
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("--save-every", type=int, default=10)
|
||||||
parser.add_argument("--resume", action="store_true", help="Resume from the progress file.")
|
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("--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("--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.")
|
parser.add_argument("--status", action="store_true", help="Show local progress and LightRAG pipeline status.")
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|||||||
152
scripts/lightrag_ingest/rebuild_vdb.py
Normal file
152
scripts/lightrag_ingest/rebuild_vdb.py
Normal 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()
|
||||||
Loading…
Reference in New Issue
Block a user