""" Knowledge graph builder following the original ``index()`` method. Steps: 1. Insert docs into chunk embedding store. 2. Run OpenIE (NER + triple extraction). 3. Build igraph with entity nodes, passage nodes, fact edges, passage-to-entity edges, and synonymy edges. 4. Save to ``graph.pickle``. """ import json import logging import os import re from collections import defaultdict from typing import Dict, List, Set, Tuple import igraph as ig import numpy as np import torch from tqdm import tqdm from .config import PassageEntityConfig from .embedding_store import EmbeddingStore, EmbeddingModelWrapper from .openie import OpenIE from .utils import ( NerRawOutput, TripleRawOutput, compute_mdhash_id, text_processing, extract_entity_nodes, flatten_facts, reformat_openie_results, filter_invalid_triples, ) logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # KNN helper (simplified from the original embed_utils.py) # --------------------------------------------------------------------------- def retrieve_knn( query_ids: List[str], key_ids: List[str], query_vecs: np.ndarray, key_vecs: np.ndarray, k: int = 2047, query_batch_size: int = 1000, key_batch_size: int = 10000, ) -> Dict: """Batched top-k cosine nearest-neighbour search using PyTorch.""" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if len(key_vecs) == 0: return {} q = torch.tensor(np.array(query_vecs), dtype=torch.float32) q = torch.nn.functional.normalize(q, dim=1) keys = torch.tensor(np.array(key_vecs), dtype=torch.float32) keys = torch.nn.functional.normalize(keys, dim=1) results = {} def _batches(vecs, bs): for i in range(0, len(vecs), bs): yield vecs[i : i + bs], i for qb, qstart in tqdm( _batches(q, query_batch_size), total=(len(q) + query_batch_size - 1) // query_batch_size, desc="KNN", ): qb = qb.to(device) batch_scores, batch_indices = [], [] offset = 0 for kb, _ in _batches(keys, key_batch_size): kb = kb.to(device) actual_kb_size = kb.size(0) sim = torch.mm(qb, kb.T) topk_s, topk_i = torch.topk(sim, min(k, actual_kb_size), dim=1, largest=True, sorted=True) topk_i += offset batch_scores.append(topk_s) batch_indices.append(topk_i) del sim kb = kb.cpu() torch.cuda.empty_cache() offset += actual_kb_size batch_scores = torch.cat(batch_scores, dim=1) batch_indices = torch.cat(batch_indices, dim=1) final_s, final_i = torch.topk( batch_scores, min(k, batch_scores.size(1)), dim=1, largest=True, sorted=True ) final_i = final_i.cpu() final_s = final_s.cpu() for i in range(final_i.size(0)): qi = qstart + i topk_rel = batch_indices[i][final_i[i]].cpu() topk_keys = [key_ids[idx] for idx in topk_rel.numpy()] results[query_ids[qi]] = (topk_keys, final_s[i].cpu().numpy().tolist()) qb = qb.cpu() torch.cuda.empty_cache() return results # =========================================================================== # KGBuilder # =========================================================================== class KGBuilder: """Build a passage-entity knowledge graph and persist it to disk.""" def __init__( self, config: PassageEntityConfig, embedding_model: EmbeddingModelWrapper, openie: OpenIE, ): self.config = config self.embedding_model = embedding_model self.openie = openie wd = config.working_dir os.makedirs(wd, exist_ok=True) self.chunk_embedding_store = EmbeddingStore( embedding_model, os.path.join(wd, "chunk_embeddings"), config.embedding_batch_size, "chunk", ) self.entity_embedding_store = EmbeddingStore( embedding_model, os.path.join(wd, "entity_embeddings"), config.embedding_batch_size, "entity", ) self.fact_embedding_store = EmbeddingStore( embedding_model, os.path.join(wd, "fact_embeddings"), config.embedding_batch_size, "fact", ) self._graph_pickle_path = os.path.join(wd, "graph.pickle") self.openie_results_path = os.path.join( config.save_dir, f"openie_results_ner_{config.llm_model.replace('/', '_')}.json", ) self.graph: ig.Graph = self._load_or_create_graph() self.node_to_node_stats: Dict[Tuple[str, str], float] = {} self.ent_node_to_chunk_ids: Dict[str, set] = {} # ------------------------------------------------------------------ # Graph init # ------------------------------------------------------------------ def _load_or_create_graph(self) -> ig.Graph: if ( not self.config.force_index_from_scratch and os.path.exists(self._graph_pickle_path) ): g = ig.Graph.Read_Pickle(self._graph_pickle_path) logger.info( f"Loaded graph from {self._graph_pickle_path}: " f"{g.vcount()} nodes, {g.ecount()} edges" ) return g return ig.Graph(directed=self.config.is_directed_graph) # ------------------------------------------------------------------ # index() # ------------------------------------------------------------------ def index(self, docs: List[str]): """Index documents: embed chunks, run OpenIE, build KG, save.""" logger.info("=== Indexing documents ===") # 1) Insert chunks into embedding store self.chunk_embedding_store.insert_strings(docs) chunk_to_rows = self.chunk_embedding_store.get_all_id_to_rows() # 2) Run OpenIE (or load cached) all_openie_info, chunk_keys_to_process = self._load_existing_openie( chunk_to_rows.keys() ) new_openie_rows = {k: chunk_to_rows[k] for k in chunk_keys_to_process} if len(chunk_keys_to_process) > 0: logger.info(f"Running OpenIE on {len(chunk_keys_to_process)} new chunks") ner_dict, triple_dict = self.openie.batch_openie(new_openie_rows) self._merge_openie_results( all_openie_info, new_openie_rows, ner_dict, triple_dict ) if self.config.save_openie: self._save_openie_results(all_openie_info) ner_results, triple_results = reformat_openie_results(all_openie_info) # Sanity check — fill missing entries for cid in chunk_to_rows: if cid not in ner_results: ner_results[cid] = NerRawOutput(cid, None, [], {}) if cid not in triple_results: triple_results[cid] = TripleRawOutput(cid, None, [], {}) chunk_ids = list(chunk_to_rows.keys()) chunk_triples = [ [text_processing(t) for t in triple_results[cid].triples] for cid in chunk_ids ] entity_nodes, chunk_triple_entities = extract_entity_nodes(chunk_triples) facts = flatten_facts(chunk_triples) # 3) Encode entities + facts logger.info("Encoding entities") self.entity_embedding_store.insert_strings(entity_nodes) logger.info("Encoding facts") self.fact_embedding_store.insert_strings([str(f) for f in facts]) # 4) Build graph edges logger.info("Building graph edges") self.node_to_node_stats = {} self.ent_node_to_chunk_ids = {} self._add_fact_edges(chunk_ids, chunk_triples) num_new = self._add_passage_edges(chunk_ids, chunk_triple_entities) if num_new > 0: logger.info(f"{num_new} new chunks → adding synonymy edges") self._add_synonymy_edges() self._augment_graph() self._save_graph() logger.info("=== Indexing complete ===") # ------------------------------------------------------------------ # Edge builders # ------------------------------------------------------------------ def _add_fact_edges(self, chunk_ids: List[str], chunk_triples: List[list]): current_nodes = set(self.graph.vs["name"]) if "name" in self.graph.vs.attribute_names() else set() for chunk_key, triples in tqdm( zip(chunk_ids, chunk_triples), desc="Fact edges", total=len(chunk_ids) ): entities_in_chunk: set = set() if chunk_key not in current_nodes: for triple in triples: triple = tuple(triple) nk1 = compute_mdhash_id(triple[0], prefix="entity-") nk2 = compute_mdhash_id(triple[2], prefix="entity-") self.node_to_node_stats[(nk1, nk2)] = ( self.node_to_node_stats.get((nk1, nk2), 0.0) + 1 ) self.node_to_node_stats[(nk2, nk1)] = ( self.node_to_node_stats.get((nk2, nk1), 0.0) + 1 ) entities_in_chunk.update([nk1, nk2]) for node in entities_in_chunk: self.ent_node_to_chunk_ids[node] = ( self.ent_node_to_chunk_ids.get(node, set()) | {chunk_key} ) def _add_passage_edges( self, chunk_ids: List[str], chunk_triple_entities: List[List[str]] ) -> int: current_nodes = set(self.graph.vs["name"]) if "name" in self.graph.vs.attribute_names() else set() num_new = 0 for idx, chunk_key in tqdm( enumerate(chunk_ids), desc="Passage edges", total=len(chunk_ids) ): if chunk_key not in current_nodes: for ent in chunk_triple_entities[idx]: nk = compute_mdhash_id(ent, prefix="entity-") self.node_to_node_stats[(chunk_key, nk)] = 1.0 num_new += 1 return num_new def _add_synonymy_edges(self): logger.info("Expanding graph with synonymy edges") entity_id_to_row = self.entity_embedding_store.get_all_id_to_rows() entity_node_keys = list(entity_id_to_row.keys()) entity_embs = self.entity_embedding_store.get_embeddings(entity_node_keys) knn = retrieve_knn( query_ids=entity_node_keys, key_ids=entity_node_keys, query_vecs=entity_embs, key_vecs=entity_embs, k=self.config.synonymy_edge_topk, query_batch_size=self.config.synonymy_edge_query_batch_size, key_batch_size=self.config.synonymy_edge_key_batch_size, ) for nk in tqdm(knn, desc="Synonymy edges"): entity = entity_id_to_row[nk]["content"] if len(re.sub('[^A-Za-z0-9]', '', entity)) <= 2: continue nns_keys, nns_scores = knn[nk] num_nns = 0 for nn, score in zip(nns_keys, nns_scores): if score < self.config.synonymy_edge_sim_threshold or num_nns > 100: break nn_phrase = entity_id_to_row.get(nn, {}).get("content", "") if nn != nk and nn_phrase: self.node_to_node_stats[(nk, nn)] = score num_nns += 1 # ------------------------------------------------------------------ # Graph augmentation # ------------------------------------------------------------------ def _augment_graph(self): self._add_new_nodes() self._add_new_edges() info = self._get_graph_info() logger.info(f"Graph info: {info}") def _add_new_nodes(self): existing = {v["name"]: v for v in self.graph.vs if "name" in v.attributes()} entity_rows = self.entity_embedding_store.get_all_id_to_rows() passage_rows = self.chunk_embedding_store.get_all_id_to_rows() all_rows = {**entity_rows, **passage_rows} new_nodes: Dict[str, list] = {} for nid, node in all_rows.items(): node["name"] = nid if nid not in existing: for k, v in node.items(): new_nodes.setdefault(k, []).append(v) if new_nodes: self.graph.add_vertices( n=len(next(iter(new_nodes.values()))), attributes=new_nodes ) def _add_new_edges(self): edge_src, edge_tgt, weights = [], [], [] for (s, t), w in self.node_to_node_stats.items(): if s == t: continue edge_src.append(s) edge_tgt.append(t) weights.append(w) current_ids = set(self.graph.vs["name"]) valid_edges, valid_w = [], [] for s, t, w in zip(edge_src, edge_tgt, weights): if s in current_ids and t in current_ids: valid_edges.append((s, t)) valid_w.append(w) else: logger.warning(f"Skipping invalid edge {s} -> {t}") self.graph.add_edges(valid_edges, attributes={"weight": valid_w}) def _save_graph(self): logger.info( f"Writing graph: {self.graph.vcount()} nodes, {self.graph.ecount()} edges" ) self.graph.write_pickle(self._graph_pickle_path) def _get_graph_info(self) -> Dict: ent_keys = set(self.entity_embedding_store.get_all_ids()) pass_keys = set(self.chunk_embedding_store.get_all_ids()) return { "num_entity_nodes": len(ent_keys), "num_passage_nodes": len(pass_keys), "num_total_nodes": len(ent_keys) + len(pass_keys), "num_facts": len(self.fact_embedding_store.get_all_ids()), "num_edges": len(self.node_to_node_stats), } # ------------------------------------------------------------------ # OpenIE persistence # ------------------------------------------------------------------ def _load_existing_openie(self, chunk_keys) -> Tuple[list, set]: chunk_keys_to_save: set = set() if ( not self.config.force_openie_from_scratch and os.path.isfile(self.openie_results_path) ): data = json.load(open(self.openie_results_path)) all_info = data.get("docs", []) # Standardise indices for item in all_info: item["idx"] = compute_mdhash_id(item["passage"], "chunk-") # Only consider a doc "done" if it has at least some extracted content existing_keys = set() for info in all_info: has_content = ( len(info.get("extracted_entities", [])) > 0 or len(info.get("extracted_triples", [])) > 0 ) if has_content: existing_keys.add(info["idx"]) for ck in chunk_keys: if ck not in existing_keys: chunk_keys_to_save.add(ck) else: all_info = [] chunk_keys_to_save = set(chunk_keys) return all_info, chunk_keys_to_save def _merge_openie_results(self, all_info, chunks, ner_dict, triple_dict): for ck, row in chunks.items(): passage = row["content"] try: info = { "idx": ck, "passage": passage, "extracted_entities": ner_dict[ck].unique_entities, "extracted_triples": triple_dict[ck].triples, } except Exception as e: logger.error(f"Error merging chunk {ck}: {e}") info = { "idx": ck, "passage": passage, "extracted_entities": [], "extracted_triples": [], } all_info.append(info) def _save_openie_results(self, all_info: list): num_phrases = sum(len(c["extracted_entities"]) for c in all_info) if num_phrases > 0: avg_chars = round( sum(len(e) for c in all_info for e in c["extracted_entities"]) / num_phrases, 4 ) avg_words = round( sum(len(e.split()) for c in all_info for e in c["extracted_entities"]) / num_phrases, 4 ) else: avg_chars, avg_words = 0, 0 os.makedirs(os.path.dirname(self.openie_results_path), exist_ok=True) with open(self.openie_results_path, "w") as f: json.dump( {"docs": all_info, "avg_ent_chars": avg_chars, "avg_ent_words": avg_words}, f, ) logger.info(f"OpenIE results saved to {self.openie_results_path}")