| import json, os, numpy as np, faiss |
| from pathlib import Path |
| from typing import Dict, List, Tuple, Any |
|
|
| class ConceptNode: |
| def __init__(self, cid, label, embedding, confidence, edges=None): |
| self.cid, self.label, self.embedding, self.confidence, self.edges = cid, label, embedding, confidence, edges or [] |
| def to_dict(self): |
| return {"cid": self.cid, "label": self.label, "embedding": self.embedding.tolist(), "confidence": self.confidence, "edges": self.edges} |
| @staticmethod |
| def from_dict(d): |
| return ConceptNode(int(d["cid"]), str(d["label"]), np.array(d["embedding"], dtype=np.float32), float(d["confidence"]), [tuple(e) for e in d.get("edges", [])]) |
|
|
| class ConceptGraph: |
| def __init__(self, dim=768, persist_dir="data/concept_graph"): |
| self.dim, self.persist_dir = dim, Path(persist_dir) |
| self.persist_dir.mkdir(parents=True, exist_ok=True) |
| self.index = faiss.IndexFlatL2(dim) |
| self._nodes: Dict[int, ConceptNode] = {} |
| def add_node(self, label, embedding, confidence, edges=None): |
| vec = embedding.astype(np.float32) |
| vec /= np.linalg.norm(vec) |
| self.index.add(np.expand_dims(vec, 0)) |
| cid = self.index.ntotal - 1 |
| node = ConceptNode(cid, label, vec, confidence, edges) |
| self._nodes[cid] = node |
| return cid |
| def persist(self): |
| with (self.persist_dir / "concepts.json").open("w") as f: |
| json.dump([n.to_dict() for n in self._nodes.values()], f, indent=2) |
| faiss.write_index(self.index, str(self.persist_dir / "faiss.index")) |
|
|
