""" Full retrieval pipeline following the original retrieve + graph_search_with_fact_entities. Steps: 1. Encode query for fact matching and passage matching. 2. Score facts, rerank with LLM. 3. Compute entity seed weights and passage weights. 4. Run MultiSeedFlowDiffusionRetriever (QAFD) on the KG. 5. Return ranked passages. """ import json import logging import os import time from typing import Callable, Dict, List, Optional, Tuple import igraph as ig import numpy as np from tqdm import tqdm from .config import PassageEntityConfig from .embedding_store import EmbeddingStore, EmbeddingModelWrapper from .graph_adapter import run_igraph_qafd from .prompts import get_query_instruction from .reranker import FactReranker from .utils import ( QuerySolution, NerRawOutput, TripleRawOutput, compute_mdhash_id, text_processing, extract_entity_nodes, flatten_facts, min_max_normalize, reformat_openie_results, ) logger = logging.getLogger(__name__) class PassageEntityRetriever: """End-to-end retriever: query -> ranked passages. Uses the pre-built KG (igraph), embedding stores, and QAFD flow diffusion. """ def __init__( self, config: PassageEntityConfig, embedding_model: EmbeddingModelWrapper, reranker: FactReranker, graph: ig.Graph, chunk_embedding_store: EmbeddingStore, entity_embedding_store: EmbeddingStore, fact_embedding_store: EmbeddingStore, openie_results_path: str, ): self.config = config self.embedding_model = embedding_model self.reranker = reranker self.graph = graph self.chunk_store = chunk_embedding_store self.entity_store = entity_embedding_store self.fact_store = fact_embedding_store self.openie_results_path = openie_results_path # Filled by prepare() self._ready = False self.entity_node_keys: List[str] = [] self.passage_node_keys: List[str] = [] self.fact_node_keys: List[str] = [] self.entity_embeddings: np.ndarray = np.array([]) self.passage_embeddings: np.ndarray = np.array([]) self.fact_embeddings: np.ndarray = np.array([]) self.node_name_to_vertex_idx: Dict[str, int] = {} self.entity_node_idxs: List[int] = [] self.passage_node_idxs: List[int] = [] self.ent_node_to_chunk_ids: Dict[str, set] = {} # Cached query embeddings self._query_emb_fact: Dict[str, np.ndarray] = {} self._query_emb_pass: Dict[str, np.ndarray] = {} # Timing accumulators self.rerank_time = 0.0 self.qafd_time = 0.0 self.total_time = 0.0 # ------------------------------------------------------------------ # Preparation (mirrors the passage-entity pipeline prepare_retrieval_objects) # ------------------------------------------------------------------ def prepare(self): """Load embeddings, build lookup structures. Call once before retrieve().""" logger.info("Preparing retrieval objects ...") self.entity_node_keys = list(self.entity_store.get_all_ids()) self.passage_node_keys = list(self.chunk_store.get_all_ids()) self.fact_node_keys = list(self.fact_store.get_all_ids()) # Node index mapping try: name_to_idx = {v["name"]: idx for idx, v in enumerate(self.graph.vs)} self.node_name_to_vertex_idx = name_to_idx self.entity_node_idxs = [name_to_idx[k] for k in self.entity_node_keys] self.passage_node_idxs = [name_to_idx[k] for k in self.passage_node_keys] except Exception as e: logger.error(f"Graph index mapping failed: {e}") self.node_name_to_vertex_idx = {} self.entity_node_idxs = [] self.passage_node_idxs = [] # Embeddings self.entity_embeddings = np.array( self.entity_store.get_embeddings(self.entity_node_keys) ) if self.entity_node_keys else np.array([]) self.passage_embeddings = np.array( self.chunk_store.get_embeddings(self.passage_node_keys) ) if self.passage_node_keys else np.array([]) self.fact_embeddings = np.array( self.fact_store.get_embeddings(self.fact_node_keys) ) if self.fact_node_keys else np.array([]) # Build ent_node_to_chunk_ids from openie results self.ent_node_to_chunk_ids = {} if os.path.isfile(self.openie_results_path): all_info = json.load(open(self.openie_results_path)).get("docs", []) ner_dict, triple_dict = reformat_openie_results(all_info) for cid in self.passage_node_keys: if cid not in triple_dict: continue triples = [text_processing(t) for t in triple_dict[cid].triples] for triple in triples: if len(triple) == 3: for ent in [triple[0], triple[2]]: nk = compute_mdhash_id(ent, prefix="entity-") self.ent_node_to_chunk_ids.setdefault(nk, set()).add(cid) self._ready = True logger.info( f"Ready. entities={len(self.entity_node_keys)}, " f"passages={len(self.passage_node_keys)}, " f"facts={len(self.fact_node_keys)}" ) # ------------------------------------------------------------------ # Query embedding # ------------------------------------------------------------------ def _encode_queries(self, queries: List[str]): to_encode = [q for q in queries if q not in self._query_emb_fact] if not to_encode: return fact_embs = self.embedding_model.batch_encode( to_encode, instruction=get_query_instruction("query_to_fact"), norm=True ) pass_embs = self.embedding_model.batch_encode( to_encode, instruction=get_query_instruction("query_to_passage"), norm=True ) for q, fe, pe in zip(to_encode, fact_embs, pass_embs): self._query_emb_fact[q] = fe self._query_emb_pass[q] = pe # ------------------------------------------------------------------ # Fact scoring # ------------------------------------------------------------------ def _get_fact_scores(self, query: str) -> np.ndarray: qe = self._query_emb_fact.get(query) if qe is None: qe = self.embedding_model.batch_encode( query, instruction=get_query_instruction("query_to_fact"), norm=True ) if len(self.fact_embeddings) == 0: return np.array([]) scores = np.dot(self.fact_embeddings, qe.T) scores = np.squeeze(scores) if scores.ndim == 2 else scores return min_max_normalize(scores) # ------------------------------------------------------------------ # Dense passage retrieval (fallback) # ------------------------------------------------------------------ def _dense_passage_retrieval(self, query: str) -> Tuple[np.ndarray, np.ndarray]: qe = self._query_emb_pass.get(query) if qe is None: qe = self.embedding_model.batch_encode( query, instruction=get_query_instruction("query_to_passage"), norm=True ) scores = np.dot(self.passage_embeddings, qe.T) scores = np.squeeze(scores) if scores.ndim == 2 else scores scores = min_max_normalize(scores) sorted_ids = np.argsort(scores)[::-1] return sorted_ids, scores[sorted_ids] # ------------------------------------------------------------------ # Rerank facts # ------------------------------------------------------------------ def _rerank_facts( self, query: str, fact_scores: np.ndarray ) -> Tuple[List[int], List[tuple], dict]: link_top_k = self.config.linking_top_k if len(fact_scores) == 0 or len(self.fact_node_keys) == 0: return [], [], {} if len(fact_scores) <= link_top_k: cand_indices = np.argsort(fact_scores)[::-1].tolist() else: cand_indices = np.argsort(fact_scores)[-link_top_k:][::-1].tolist() real_ids = [self.fact_node_keys[i] for i in cand_indices] rows = self.fact_store.get_rows(real_ids) cand_facts = [eval(rows[rid]["content"]) for rid in real_ids] top_indices, top_facts, meta = self.reranker( query, cand_facts, cand_indices, len_after_rerank=link_top_k ) return top_indices, top_facts, meta # ------------------------------------------------------------------ # Graph search (core of passage-entity retrieval) # ------------------------------------------------------------------ def _graph_search( self, query: str, fact_scores: np.ndarray, top_k_facts: List[tuple], top_k_fact_indices: List[int], ) -> Tuple[np.ndarray, np.ndarray]: """Compute seed weights -> run QAFD -> return sorted passage ids + scores.""" link_top_k = self.config.linking_top_k n_nodes = self.graph.vcount() # --- entity seed weights --- linking_score_map: Dict[str, float] = {} phrase_scores: Dict[str, list] = {} phrase_weights = np.zeros(n_nodes) passage_weights = np.zeros(n_nodes) number_of_occurs = np.zeros(n_nodes) phrases_and_ids = set() for rank, f in enumerate(top_k_facts): subj = f[0].lower() obj = f[2].lower() fs = ( fact_scores[top_k_fact_indices[rank]] if fact_scores.ndim > 0 else float(fact_scores) ) for phrase in [subj, obj]: pk = compute_mdhash_id(phrase, prefix="entity-") pid = self.node_name_to_vertex_idx.get(pk) if pid is not None: wfs = fs num_chunks = len(self.ent_node_to_chunk_ids.get(pk, set())) if num_chunks > 0: wfs /= num_chunks phrase_weights[pid] += wfs number_of_occurs[pid] += 1 phrases_and_ids.add((phrase, pid)) # Normalise nonzero = number_of_occurs > 0 phrase_weights[nonzero] /= number_of_occurs[nonzero] for phrase, pid in phrases_and_ids: if pid is not None: phrase_scores.setdefault(phrase, []).append(phrase_weights[pid]) for phrase, scores in phrase_scores.items(): linking_score_map[phrase] = float(np.mean(scores)) # Keep only top-k entity seeds if link_top_k and linking_score_map: linking_score_map = dict( sorted(linking_score_map.items(), key=lambda x: x[1], reverse=True)[ :link_top_k ] ) top_phrases = { compute_mdhash_id(p, prefix="entity-") for p in linking_score_map } for nk in self.node_name_to_vertex_idx: if nk not in top_phrases: pid = self.node_name_to_vertex_idx.get(nk) if pid is not None: phrase_weights[pid] = 0.0 # --- passage seed weights --- dpr_ids, dpr_scores = self._dense_passage_retrieval(query) norm_dpr = min_max_normalize(dpr_scores) pw = self.config.passage_node_weight for i, did in enumerate(dpr_ids.tolist()): pk = self.passage_node_keys[did] pid = self.node_name_to_vertex_idx.get(pk) if pid is not None: passage_weights[pid] = norm_dpr[i] * pw node_weights = phrase_weights + passage_weights if np.sum(node_weights) == 0: logger.warning("All node weights are zero after seed selection, falling back to DPR") return dpr_ids, dpr_scores # --- Build node embeddings dict (cached) --- if not hasattr(self, '_node_emb_dict') or self._node_emb_dict is None: self._node_emb_dict = {} for i, nk in enumerate(self.entity_node_keys): if i < len(self.entity_embeddings): self._node_emb_dict[nk] = self.entity_embeddings[i] for i, nk in enumerate(self.passage_node_keys): if i < len(self.passage_embeddings): self._node_emb_dict[nk] = self.passage_embeddings[i] query_emb = self._query_emb_fact.get(query) # --- Run QAFD --- qafd_start = time.time() sorted_ids, sorted_scores = run_igraph_qafd( graph=self.graph, node_name_to_idx=self.node_name_to_vertex_idx, passage_node_idxs=self.passage_node_idxs, source_weights=node_weights, node_embeddings=self._node_emb_dict, query_embedding=query_emb, alpha=self.config.qafd_alpha, epsilon=self.config.qafd_epsilon, max_iterations=self.config.qafd_max_iterations, step_size=self.config.qafd_step_size, weight_scheme=self.config.qafd_weight_scheme, use_node_degree=self.config.qafd_use_node_degree, random_seed=self.config.qafd_random_seed, sim_mode=self.config.sim_mode, qa_sink_gamma=self.config.qa_sink_gamma, qa_warm_delta=self.config.qa_warm_delta, qa_warm_walk=self.config.qa_warm_walk, qa_warm_steps=self.config.qa_warm_steps, qa_accum_gamma=self.config.qa_accum_gamma, qa_post_lambda=self.config.qa_post_lambda, batch_push=self.config.batch_push, ) qafd_elapsed = time.time() - qafd_start self.qafd_time += qafd_elapsed logger.info(f"QAFD completed in {qafd_elapsed:.2f}s") return sorted_ids, sorted_scores # ------------------------------------------------------------------ # Public interface # ------------------------------------------------------------------ def retrieve( self, queries: List[str], num_to_retrieve: int = None, gold_docs: List[List[str]] = None, ) -> List[QuerySolution]: """Retrieve documents for a batch of queries. Returns a list of ``QuerySolution`` objects. """ if not self._ready: self.prepare() if num_to_retrieve is None: num_to_retrieve = self.config.retrieval_top_k self._encode_queries(queries) results = [] t0 = time.time() for q in tqdm(queries, desc="Retrieving"): rerank_t0 = time.time() fact_scores = self._get_fact_scores(q) top_indices, top_facts, _ = self._rerank_facts(q, fact_scores) self.rerank_time += time.time() - rerank_t0 if len(top_facts) == 0: # Reranker rejected all facts — use top facts by embedding score logger.info("No facts after reranking -> using top facts by embedding score") link_top_k = self.config.linking_top_k if len(fact_scores) > 0: top_indices = np.argsort(fact_scores)[-link_top_k:][::-1].tolist() real_ids = [self.fact_node_keys[i] for i in top_indices] rows = self.fact_store.get_rows(real_ids) top_facts = [eval(rows[rid]["content"]) for rid in real_ids] sorted_ids, sorted_scores = self._graph_search( q, fact_scores, top_facts, top_indices ) top_docs = [ self.chunk_store.get_row(self.passage_node_keys[idx])["content"] for idx in sorted_ids[:num_to_retrieve] ] results.append( QuerySolution( question=q, docs=top_docs, doc_scores=sorted_scores[:num_to_retrieve], ) ) self.total_time += time.time() - t0 logger.info( f"Retrieval done. total={self.total_time:.1f}s, " f"rerank={self.rerank_time:.1f}s, qafd={self.qafd_time:.1f}s" ) return results