Rename: passage_entity pipeline (was hipporag_pipeline)

src/passage_entity/benchmark_runner.py

@@ -0,0 +1,737 @@
+"""
+Entry point for running multihop benchmarks (MuSiQue, HotpotQA, 2WikiMultiHopQA)
+Passage-entity KG pipeline benchmark runner for QAFD-RAG.
+
+Usage::
+
+    python -m src.passage_entity.benchmark_runner \\
+        --dataset musique \\
+        --llm_model gpt-4o-mini \\
+        --embedding_model nvidia-nv-embed-v2 \\
+        --num_queries 100 \\
+        --qafd_alpha 10.0
+
+The script will:
+    1. Load corpus and questions from ``data/multihop/``.
+    2. Build (or load) the knowledge graph.
+    3. Run retrieval + QA.
+    4. Evaluate Recall@K, Exact Match, and F1.
+"""
+
+import argparse
+import asyncio
+import collections
+import json
+import logging
+import os
+import re
+import string
+import sys
+import time
+from typing import Dict, List, Optional, Set, Tuple
+
+import numpy as np
+
+# ---------------------------------------------------------------------------
+# Ensure the QAFD-RAG root is on the path so ``src.*`` imports work
+# when this file is executed as ``python -m src.passage_entity.benchmark_runner``
+# ---------------------------------------------------------------------------
+_project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+if _project_root not in sys.path:
+    sys.path.insert(0, _project_root)
+
+# ---------------------------------------------------------------------------
+# Bypass src/__init__.py (which imports heavy AWS deps) by registering
+# src as a plain namespace package before any sub-package imports.
+# ---------------------------------------------------------------------------
+import types as _types
+for _pkg_path in ["src", "src.retrievers", "src.passage_entity"]:
+    if _pkg_path not in sys.modules:
+        _m = _types.ModuleType(_pkg_path)
+        _m.__path__ = [os.path.join(_project_root, *_pkg_path.split("."))]
+        _m.__package__ = _pkg_path
+        sys.modules[_pkg_path] = _m
+
+# Load only the modules we actually need (no aioboto3, no AWS, no SAPIEN)
+import importlib.util as _ilu
+def _load_mod(fqn, filepath):
+    spec = _ilu.spec_from_file_location(fqn, filepath)
+    mod = _ilu.module_from_spec(spec)
+    sys.modules[fqn] = mod
+    spec.loader.exec_module(mod)
+    return mod
+
+_src = os.path.join(_project_root, "src")
+_load_mod("src.retrievers.base", os.path.join(_src, "retrievers", "base.py"))
+_load_mod("src.retrievers.flow_diffusion", os.path.join(_src, "retrievers", "flow_diffusion.py"))
+
+from src.passage_entity.config import PassageEntityConfig
+from src.passage_entity.embedding_store import EmbeddingModelWrapper
+from src.passage_entity.kg_builder import KGBuilder
+from src.passage_entity.openie import OpenIE
+from src.passage_entity.reranker import FactReranker
+from src.passage_entity.retriever import PassageEntityRetriever
+from src.passage_entity.prompts import make_qa_messages
+from src.passage_entity.utils import QuerySolution
+
+# ---------------------------------------------------------------------------
+# Minimal OpenAI LLM + Embedding (no AWS deps, no src/llm.py)
+# ---------------------------------------------------------------------------
+from openai import AsyncOpenAI
+
+# Shared client — avoids "Event loop is closed" errors from abandoned clients
+_openai_clients: dict = {}
+
+def _get_client(base_url="https://api.openai.com/v1", api_key=""):
+    key = (base_url, api_key)
+    if key not in _openai_clients:
+        _openai_clients[key] = AsyncOpenAI(
+            base_url=base_url,
+            api_key=api_key or os.environ.get("OPENAI_API_KEY", ""),
+        )
+    return _openai_clients[key]
+
+async def _openai_complete(model, prompt, system_prompt=None, history_messages=[],
+                           base_url="https://api.openai.com/v1", api_key="", **kwargs):
+    client = _get_client(base_url, api_key)
+    kwargs.pop("hashing_kv", None)
+    kwargs.pop("keyword_extraction", None)
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+    messages.extend(history_messages)
+    messages.append({"role": "user", "content": prompt})
+    response = await client.chat.completions.create(model=model, messages=messages, **kwargs)
+    return response.choices[0].message.content
+
+async def _openai_embed(texts, model="text-embedding-3-small", api_key=""):
+    client = _get_client(api_key=api_key)
+    cleaned = [t if t.strip() else " " for t in texts]
+    response = await client.embeddings.create(model=model, input=cleaned, encoding_format="float")
+    return np.array([dp.embedding for dp in response.data])
+
+logger = logging.getLogger(__name__)
+
+# ===========================================================================
+# Gold extraction helpers (from the original pipeline main_qafd.py)
+# ===========================================================================
+
+def get_gold_docs(samples: List[dict], dataset_name: str = None) -> List[List[str]]:
+    gold_docs = []
+    for sample in samples:
+        if "supporting_facts" in sample:
+            gold_titles = {item[0] for item in sample["supporting_facts"]}
+            pairs = [item for item in sample["context"] if item[0] in gold_titles]
+            if dataset_name and dataset_name.startswith("hotpotqa"):
+                gd = [item[0] + "\n" + "".join(item[1]) for item in pairs]
+            else:
+                gd = [item[0] + "\n" + " ".join(item[1]) for item in pairs]
+        elif "contexts" in sample:
+            gd = [
+                item["title"] + "\n" + item["text"]
+                for item in sample["contexts"]
+                if item["is_supporting"]
+            ]
+        elif "paragraphs" in sample:
+            paras = [
+                p for p in sample["paragraphs"]
+                if p.get("is_supporting", True)
+            ]
+            gd = [
+                p["title"] + "\n" + p.get("text", p.get("paragraph_text", ""))
+                for p in paras
+            ]
+        else:
+            gd = []
+        gold_docs.append(list(set(gd)))
+    return gold_docs
+
+
+def get_gold_answers(samples: List[dict]) -> List[Set[str]]:
+    answers = []
+    for s in samples:
+        ans = s.get("answer") or s.get("gold_ans") or s.get("reference")
+        if ans is None and "obj" in s:
+            ans = list(
+                {s["obj"], s.get("possible_answers", ""), s.get("o_wiki_title", ""), s.get("o_aliases", "")}
+            )
+        if ans is None:
+            ans = ""
+        if isinstance(ans, str):
+            ans = [ans]
+        ans_set = set(ans)
+        if "answer_aliases" in s:
+            ans_set.update(s["answer_aliases"])
+        answers.append(ans_set)
+    return answers
+
+
+# ===========================================================================
+# Evaluation metrics
+# ===========================================================================
+
+def _normalize_answer(s: str) -> str:
+    """Lower-case, remove articles, punctuation, extra whitespace."""
+    s = s.lower()
+    s = re.sub(r"\b(a|an|the)\b", " ", s)
+    s = "".join(ch for ch in s if ch not in string.punctuation)
+    return " ".join(s.split())
+
+
+def exact_match(prediction: str, gold_answers: Set[str]) -> float:
+    pred_norm = _normalize_answer(prediction)
+    return float(any(_normalize_answer(g) == pred_norm for g in gold_answers))
+
+
+def f1_score(prediction: str, gold_answers: Set[str]) -> float:
+    pred_tokens = _normalize_answer(prediction).split()
+    best_f1 = 0.0
+    for gold in gold_answers:
+        gold_tokens = _normalize_answer(gold).split()
+        common = collections.Counter(pred_tokens) & collections.Counter(gold_tokens)
+        num_same = sum(common.values())
+        if num_same == 0:
+            continue
+        precision = num_same / len(pred_tokens)
+        recall = num_same / len(gold_tokens)
+        f1 = 2 * precision * recall / (precision + recall)
+        best_f1 = max(best_f1, f1)
+    return best_f1
+
+
+def recall_at_k(
+    gold_docs: List[List[str]], retrieved_docs: List[List[str]], k_list: List[int]
+) -> Dict[str, float]:
+    """Compute Recall@K across all queries."""
+    results = {}
+    for k in k_list:
+        recalls = []
+        for gd, rd in zip(gold_docs, retrieved_docs):
+            if not gd:
+                continue
+            retrieved_set = set(rd[:k])
+            found = sum(1 for g in gd if g in retrieved_set)
+            recalls.append(found / len(gd))
+        results[f"Recall@{k}"] = round(np.mean(recalls), 4) if recalls else 0.0
+    return results
+
+
+# ===========================================================================
+# QA (reading comprehension)
+# ===========================================================================
+
+def _run_sync(coro):
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = None
+    if loop is not None and loop.is_running():
+        import concurrent.futures
+        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
+            return pool.submit(asyncio.run, coro).result()
+    else:
+        return asyncio.run(coro)
+
+
+def run_qa(
+    queries: List[QuerySolution],
+    llm_func,
+    qa_top_k: int = 5,
+) -> List[QuerySolution]:
+    """Run reading-comprehension QA over retrieved passages."""
+    for qs in queries:
+        passages = qs.docs[:qa_top_k]
+        msgs = make_qa_messages(passages, qs.question)
+        # Convert messages to single call
+        system_prompt = None
+        history = []
+        user_prompt = ""
+        for msg in msgs:
+            if msg["role"] == "system":
+                system_prompt = msg["content"]
+            elif msg["role"] == "assistant":
+                history.append(msg)
+            elif msg["role"] == "user":
+                if user_prompt:
+                    history.append({"role": "user", "content": user_prompt})
+                user_prompt = msg["content"]
+
+        try:
+            response = _run_sync(
+                llm_func(
+                    prompt=user_prompt,
+                    system_prompt=system_prompt,
+                    history_messages=history,
+                    max_tokens=512,
+                )
+            )
+            # Extract answer
+            if "Answer:" in response:
+                qs.answer = response.split("Answer:")[-1].strip()
+            else:
+                qs.answer = response.strip()
+        except Exception as e:
+            logger.error(f"QA error: {e}")
+            qs.answer = ""
+    return queries
+
+
+def run_qa_ultradomain(
+    queries: List[QuerySolution],
+    llm_func,
+    qa_top_k: int = 5,
+) -> List[QuerySolution]:
+    """Generate full responses for UltraDomain (not short answers)."""
+    for qs in queries:
+        passages = qs.docs[:qa_top_k]
+        context = "\n\n".join(passages)
+        prompt = (
+            f"Based on the following context, provide a comprehensive and detailed "
+            f"answer to the question.\n\n"
+            f"Context:\n{context}\n\n"
+            f"Question: {qs.question}\n\n"
+            f"Answer:"
+        )
+        try:
+            response = _run_sync(
+                llm_func(prompt=prompt, max_tokens=1024)
+            )
+            qs.answer = response.strip()
+        except Exception as e:
+            logger.error(f"QA error: {e}")
+            qs.answer = ""
+    return queries
+
+
+def run_quality_eval(
+    queries: List[str],
+    responses: List[str],
+    llm_func,
+    num_eval_rounds: int = 5,
+) -> Dict[str, List[float]]:
+    """Evaluate response quality using LLM scoring (same as entity-graph pipeline).
+
+    Each response is evaluated num_eval_rounds times on 5 criteria.
+    Returns dict of criterion -> list of per-query average scores.
+    """
+    criteria = ["comprehensiveness", "diversity", "logicality", "relevance", "coherence"]
+    result = {c: [] for c in criteria}
+
+    for i, (query, response) in enumerate(zip(queries, responses)):
+        if not response:
+            for c in criteria:
+                result[c].append(0.0)
+            continue
+
+        criterion_scores = {c: [] for c in criteria}
+        for _ in range(num_eval_rounds):
+            prompt = f"""Evaluate the following response to a question based on five criteria. Rate each criterion from 0-100.
+
+Question: {query}
+Response: {response}
+
+Please evaluate based on these criteria:
+- Comprehensiveness: How much detail does the answer provide to cover all aspects and details of the question?
+- Diversity: How varied and rich is the answer in providing different perspectives and insights on the question?
+- Logicality: How logically does the answer respond to all parts of the question?
+- Relevance: How relevant is the answer to the question, staying focused and addressing the intended topic or issue?
+- Coherence: How well does the answer maintain internal logical connections between its parts, ensuring a smooth and consistent structure?
+
+Provide scores in JSON format:
+{{
+    "comprehensiveness": [score],
+    "diversity": [score],
+    "logicality": [score],
+    "relevance": [score],
+    "coherence": [score]
+}}"""
+            try:
+                eval_response = _run_sync(
+                    llm_func(prompt=prompt, max_tokens=200)
+                )
+                import re as _re
+                json_match = _re.search(r'\{.*\}', eval_response, _re.DOTALL)
+                if json_match:
+                    scores = json.loads(json_match.group())
+                    for c in criteria:
+                        if c in scores:
+                            val = float(scores[c])
+                            if 0 <= val <= 100:
+                                criterion_scores[c].append(val)
+            except Exception:
+                continue
+
+        for c in criteria:
+            if criterion_scores[c]:
+                result[c].append(np.mean(criterion_scores[c]))
+            else:
+                result[c].append(0.0)
+
+    return result
+
+
+# ===========================================================================
+# Main
+# ===========================================================================
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="QAFD-RAG passage-entity benchmark runner",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("--dataset", type=str, default="musique",
+                        help="Dataset name (e.g. musique, hotpotqa, 2wikimultihopqa, mix)")
+    parser.add_argument("--task", type=str, default="multihop",
+                        choices=["multihop", "ultradomain"],
+                        help="Task type (determines data loading)")
+    parser.add_argument("--num_queries", type=int, default=-1,
+                        help="Number of queries (-1 = all)")
+    parser.add_argument("--data_dir", type=str, default="data/multihop",
+                        help="Directory with corpus/question JSON files (multihop only)")
+    parser.add_argument("--save_dir", type=str, default="outputs",
+                        help="Output directory")
+
+    # LLM
+    parser.add_argument("--llm_model", type=str, default="gpt-4o-mini")
+    parser.add_argument("--llm_base_url", type=str, default="https://api.openai.com/v1")
+    parser.add_argument("--llm_api_key", type=str, default="")
+
+    # Embedding
+    parser.add_argument("--embedding_model", type=str, default="nvidia-nv-embed-v2",
+                        help="Key from QAFD-RAG embedding registry")
+
+    # Indexing
+    parser.add_argument("--force_index", action="store_true")
+    parser.add_argument("--force_openie", action="store_true")
+
+    # QAFD
+    parser.add_argument("--qafd_alpha", type=float, default=2.0)
+    parser.add_argument("--qafd_epsilon", type=float, default=0.01)
+    parser.add_argument("--qafd_max_iterations", type=int, default=500)
+    parser.add_argument("--qafd_weight_scheme", type=str, default="original")
+    parser.add_argument("--qafd_step_size", type=float, default=0.2)
+
+    # Retrieval
+    parser.add_argument("--linking_top_k", type=int, default=5)
+    parser.add_argument("--retrieval_top_k", type=int, default=200)
+    parser.add_argument("--passage_node_weight", type=float, default=0.05)
+
+    # QA
+    parser.add_argument("--qa_top_k", type=int, default=5)
+    parser.add_argument("--skip_qa", action="store_true",
+                        help="Only run retrieval, skip QA step")
+
+    # Query-aware enhancements (all default = original behaviour)
+    parser.add_argument("--sim_mode", type=str, default="normalized",
+                        choices=["normalized", "relu", "relu_sq"],
+                        help="Similarity contrast function (default=normalized)")
+    parser.add_argument("--qa_sink_gamma", type=float, default=0.0,
+                        help="Query-aware sink capacity (0=off)")
+    parser.add_argument("--qa_warm_delta", type=float, default=0.0,
+                        help="Query-aware seed bias (0=off)")
+    parser.add_argument("--qa_warm_walk", action="store_true",
+                        help="Use QA edge weights in warm-start random walk")
+    parser.add_argument("--qa_warm_steps", type=int, default=2,
+                        help="Number of warm-start steps (default 2)")
+    parser.add_argument("--qa_accum_gamma", type=float, default=0.0,
+                        help="Query-aware x accumulation boost (0=off)")
+    parser.add_argument("--qa_post_lambda", type=float, default=0.0,
+                        help="Post-diffusion query reranking (0=off)")
+    parser.add_argument("--batch_push", action="store_true",
+                        help="Use batch push-relabel (all excess nodes per iter)")
+
+    # Reranker
+    parser.add_argument("--rerank_dspy_path", type=str, default=None)
+
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
+    )
+
+    # ----------------------------------------------------------------
+    # Config
+    # ----------------------------------------------------------------
+    config = PassageEntityConfig(
+        llm_model=args.llm_model,
+        llm_base_url=args.llm_base_url,
+        llm_api_key=args.llm_api_key,
+        embedding_model_key=args.embedding_model,
+        dataset=args.dataset,
+        save_dir=args.save_dir,
+        force_index_from_scratch=args.force_index,
+        force_openie_from_scratch=args.force_openie,
+        linking_top_k=args.linking_top_k,
+        retrieval_top_k=args.retrieval_top_k,
+        passage_node_weight=args.passage_node_weight,
+        qa_top_k=args.qa_top_k,
+        use_qafd=True,
+        qafd_alpha=args.qafd_alpha,
+        qafd_epsilon=args.qafd_epsilon,
+        qafd_max_iterations=args.qafd_max_iterations,
+        qafd_weight_scheme=args.qafd_weight_scheme,
+        qafd_step_size=args.qafd_step_size,
+        sim_mode=args.sim_mode,
+        qa_sink_gamma=args.qa_sink_gamma,
+        qa_warm_walk=args.qa_warm_walk,
+        qa_warm_steps=args.qa_warm_steps,
+        qa_accum_gamma=args.qa_accum_gamma,
+        batch_push=args.batch_push,
+        qa_warm_delta=args.qa_warm_delta,
+        qa_post_lambda=args.qa_post_lambda,
+        rerank_dspy_file_path=args.rerank_dspy_path,
+    )
+
+    # ----------------------------------------------------------------
+    # LLM function
+    # ----------------------------------------------------------------
+    _api_key = config.llm_api_key or os.environ.get("OPENAI_API_KEY", "")
+
+    async def llm_func(prompt, system_prompt=None, history_messages=[], **kwargs):
+        return await _openai_complete(
+            model=config.llm_model,
+            prompt=prompt,
+            system_prompt=system_prompt,
+            history_messages=history_messages,
+            base_url=config.llm_base_url,
+            api_key=_api_key,
+            **kwargs,
+        )
+
+    # ----------------------------------------------------------------
+    # Embedding function (must match the model used to build the KG)
+    # ----------------------------------------------------------------
+    emb_key = config.embedding_model_key
+    if emb_key in ("openai-small", "openai-large"):
+        openai_model = "text-embedding-3-small" if emb_key == "openai-small" else "text-embedding-3-large"
+        async def embed_func(texts):
+            return await _openai_embed(texts, model=openai_model, api_key=_api_key)
+    else:
+        # Local embedding model — use QAFD-RAG's embedding registry
+        logger.info(f"Loading local embedding model: {emb_key}")
+        _emb_cfg = type("Cfg", (), {
+            "embedding_model_name": {
+                "nvidia-nv-embed-v2": "nvidia/NV-Embed-v2",
+                "jina-v3": "jinaai/jina-embeddings-v3",
+                "gritlm": "GritLM/GritLM-7B",
+            }.get(emb_key, emb_key),
+            "embedding_batch_size": config.embedding_batch_size,
+        })()
+        _emb_src = os.path.join(_project_root, "src", "embedding_models")
+        if emb_key == "nvidia-nv-embed-v2":
+            _mod = _load_mod("src.embedding_models.NVEmbedV2", os.path.join(_emb_src, "NVEmbedV2.py"))
+            _local_model = _mod.NVEmbedV2EmbeddingModel(_emb_cfg)
+        elif emb_key == "jina-v3":
+            _mod = _load_mod("src.embedding_models.JinaV3", os.path.join(_emb_src, "JinaV3.py"))
+            _local_model = _mod.JinaV3EmbeddingModel(_emb_cfg)
+        elif emb_key == "gritlm":
+            _mod = _load_mod("src.embedding_models.GritLM", os.path.join(_emb_src, "GritLM.py"))
+            _local_model = _mod.GritLMEmbeddingModel(_emb_cfg)
+        else:
+            raise ValueError(f"Unknown embedding model: {emb_key}")
+        async def embed_func(texts):
+            return np.array(_local_model.batch_encode(texts))
+    embedding_model = EmbeddingModelWrapper(embed_func, batch_size=config.embedding_batch_size)
+
+    # ----------------------------------------------------------------
+    # Load data
+    # ----------------------------------------------------------------
+    if args.task == "ultradomain":
+        # UltraDomain: load from HuggingFace, each record has context + input
+        from datasets import load_dataset as hf_load_dataset
+
+        dataset_file = f"{args.dataset}.jsonl"
+        logger.info(f"Loading UltraDomain dataset: {dataset_file}")
+        hf_dataset = hf_load_dataset(
+            "TommyChien/UltraDomain", data_files=dataset_file, split="train"
+        )
+
+        num_q = args.num_queries if args.num_queries > 0 else len(hf_dataset)
+        num_q = min(num_q, len(hf_dataset))
+
+        # Each record's context becomes the corpus.
+        # UltraDomain contexts can be very long (30K+ chars), so we chunk them
+        # into ~500-word passages to fit embedding model token limits.
+        docs = []
+        chunk_size = 500  # words per chunk
+        chunk_overlap = 50  # word overlap between chunks
+        for i in range(num_q):
+            ctx = hf_dataset[i].get("context", "")
+            if not ctx:
+                continue
+            words = ctx.split()
+            if len(words) <= chunk_size:
+                docs.append(ctx)
+            else:
+                for start in range(0, len(words), chunk_size - chunk_overlap):
+                    chunk = " ".join(words[start : start + chunk_size])
+                    if chunk.strip():
+                        docs.append(chunk)
+
+        all_queries = [hf_dataset[i]["input"] for i in range(num_q)]
+        samples = [dict(hf_dataset[i]) for i in range(num_q)]
+        gold_answers = [
+            set(s.get("answers", [s.get("label", "")])) for s in samples
+        ]
+        gold_docs = None  # UltraDomain has no gold supporting docs
+
+    else:
+        # Multihop: load from local JSON files
+        corpus_path = os.path.join(args.data_dir, f"{args.dataset}_corpus.json")
+        questions_path = os.path.join(args.data_dir, f"{args.dataset}.json")
+
+        logger.info(f"Loading corpus from {corpus_path}")
+        with open(corpus_path) as f:
+            corpus = json.load(f)
+        docs = [f"{d['title']}\n{d['text']}" for d in corpus]
+
+        logger.info(f"Loading questions from {questions_path}")
+        with open(questions_path) as f:
+            samples = json.load(f)
+
+        all_queries = [s["question"] for s in samples]
+        if args.num_queries > 0:
+            all_queries = all_queries[: args.num_queries]
+            samples = samples[: args.num_queries]
+
+        gold_answers = get_gold_answers(samples)
+        try:
+            gold_docs = get_gold_docs(samples, args.dataset)
+        except Exception:
+            gold_docs = None
+
+    print("=" * 70)
+    print(f"  Graph type:  passage-entity")
+    print(f"  Task:        {args.task}")
+    print(f"  Dataset:     {args.dataset}")
+    print(f"  Queries:     {len(all_queries)}")
+    print(f"  Corpus:      {len(docs)} documents")
+    print(f"  LLM:         {config.llm_model}")
+    print(f"  Embedding:   {config.embedding_model_key}")
+    print(f"  QAFD alpha:  {config.qafd_alpha}")
+    print("=" * 70)
+
+    # ----------------------------------------------------------------
+    # Build / load KG
+    # ----------------------------------------------------------------
+    openie = OpenIE(llm_func)
+    builder = KGBuilder(config, embedding_model, openie)
+
+    if builder.graph.vcount() > 0 and not config.force_index_from_scratch:
+        logger.info(f"Using existing KG: {config.working_dir} "
+                     f"({builder.graph.vcount()} nodes, {builder.graph.ecount()} edges)")
+    else:
+        logger.info("Building KG from scratch ...")
+        builder.index(docs)
+
+    # ----------------------------------------------------------------
+    # Retriever
+    # ----------------------------------------------------------------
+    reranker = FactReranker(llm_func, dspy_file_path=config.rerank_dspy_file_path)
+    retriever = PassageEntityRetriever(
+        config=config,
+        embedding_model=embedding_model,
+        reranker=reranker,
+        graph=builder.graph,
+        chunk_embedding_store=builder.chunk_embedding_store,
+        entity_embedding_store=builder.entity_embedding_store,
+        fact_embedding_store=builder.fact_embedding_store,
+        openie_results_path=builder.openie_results_path,
+    )
+
+    logger.info("Running retrieval ...")
+    retrieval_results = retriever.retrieve(
+        queries=all_queries, num_to_retrieve=config.retrieval_top_k
+    )
+
+    # ----------------------------------------------------------------
+    # Retrieval evaluation
+    # ----------------------------------------------------------------
+    if gold_docs is not None:
+        k_list = [1, 2, 5, 10, 20, 50, 100, 200]
+        retrieved = [r.docs for r in retrieval_results]
+        retrieval_metrics = recall_at_k(gold_docs, retrieved, k_list)
+        print("\n--- Retrieval Metrics ---")
+        for metric, val in retrieval_metrics.items():
+            print(f"  {metric}: {val}")
+    else:
+        retrieval_metrics = {}
+
+    # ----------------------------------------------------------------
+    # QA + Evaluation (task-aware)
+    # ----------------------------------------------------------------
+    avg_em, avg_f1 = None, None
+    quality_scores = None
+
+    if not args.skip_qa:
+        logger.info("Running QA ...")
+
+        if args.task == "ultradomain":
+            # UltraDomain: generate full responses, evaluate with quality scores
+            retrieval_results = run_qa_ultradomain(
+                retrieval_results, llm_func, qa_top_k=config.qa_top_k
+            )
+            # Quality evaluation (same as entity-graph pipeline)
+            quality_scores = run_quality_eval(
+                all_queries, [qs.answer for qs in retrieval_results], llm_func
+            )
+            if quality_scores:
+                print("\n--- Quality Metrics ---")
+                overall = []
+                for criterion, scores in quality_scores.items():
+                    avg = np.mean(scores)
+                    std = np.std(scores)
+                    print(f"  {criterion:<25} {avg:.2f} +/- {std:.2f}")
+                    overall.append(avg)
+                print(f"  {'Overall Average':<25} {np.mean(overall):.2f}")
+        else:
+            # Multihop: generate short answers, evaluate with F1/EM
+            retrieval_results = run_qa(retrieval_results, llm_func, qa_top_k=config.qa_top_k)
+
+            em_scores, f1_scores = [], []
+            for qs, ga in zip(retrieval_results, gold_answers):
+                qs.gold_answers = list(ga)
+                em_scores.append(exact_match(qs.answer or "", ga))
+                f1_scores.append(f1_score(qs.answer or "", ga))
+
+            avg_em = round(np.mean(em_scores), 4)
+            avg_f1 = round(np.mean(f1_scores), 4)
+            print("\n--- QA Metrics ---")
+            print(f"  Exact Match: {avg_em}")
+            print(f"  F1 Score:    {avg_f1}")
+
+    # ----------------------------------------------------------------
+    # Save results
+    # ----------------------------------------------------------------
+    os.makedirs(config.working_dir, exist_ok=True)
+    results_path = os.path.join(config.working_dir, f"results_{args.dataset}.json")
+    output = {
+        "graph_type": "passage-entity",
+        "dataset": args.dataset,
+        "task": args.task,
+        "num_queries": len(all_queries),
+        "retrieval_metrics": retrieval_metrics,
+        "qa_em": avg_em,
+        "qa_f1": avg_f1,
+        "quality_scores": quality_scores,
+        "config": {
+            "llm_model": config.llm_model,
+            "embedding_model_key": config.embedding_model_key,
+            "qafd_alpha": config.qafd_alpha,
+            "qafd_epsilon": config.qafd_epsilon,
+            "qafd_max_iterations": config.qafd_max_iterations,
+            "qafd_weight_scheme": config.qafd_weight_scheme,
+            "linking_top_k": config.linking_top_k,
+            "retrieval_top_k": config.retrieval_top_k,
+        },
+        "per_query": [qs.to_dict() for qs in retrieval_results],
+    }
+    with open(results_path, "w") as f:
+        json.dump(output, f, indent=2, default=str)
+    print(f"\nResults saved to {results_path}")
+    print("=" * 70)
+
+
+if __name__ == "__main__":
+    main()

src/passage_entity/config.py

@@ -0,0 +1,94 @@
+"""
+Configuration dataclass for the passage-entity KG pipeline in QAFD-RAG.
+
+Combines passage-entity BaseConfig fields with QAFD algorithm parameters.
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional, Literal
+
+
+@dataclass
+class PassageEntityConfig:
+    """Unified configuration for indexing, retrieval, and QAFD."""
+
+    # ── LLM ────────────────────────────────────────────────────────────────
+    llm_model: str = "gpt-4o-mini"
+    llm_base_url: str = "https://api.openai.com/v1"
+    llm_api_key: str = ""                    # falls back to OPENAI_API_KEY env
+    max_new_tokens: Optional[int] = 2048
+    temperature: float = 0.0
+
+    # ── Embedding ──────────────────────────────────────────────────────────
+    embedding_model_key: str = "nvidia-nv-embed-v2"   # key in QAFD-RAG registry
+    embedding_batch_size: int = 16
+
+    # ── Dataset / paths ────────────────────────────────────────────────────
+    dataset: Optional[str] = None                     # musique, hotpotqa, 2wikimultihopqa
+    save_dir: str = "outputs"
+    force_index_from_scratch: bool = False
+    force_openie_from_scratch: bool = False
+    save_openie: bool = True
+
+    # ── Graph construction ─────────────────────────────────────────────────
+    is_directed_graph: bool = False
+    synonymy_edge_topk: int = 2047
+    synonymy_edge_query_batch_size: int = 1000
+    synonymy_edge_key_batch_size: int = 10000
+    synonymy_edge_sim_threshold: float = 0.8
+
+    # ── Retrieval ──────────────────────────────────────────────────────────
+    linking_top_k: int = 5
+    retrieval_top_k: int = 200
+    passage_node_weight: float = 0.05
+    damping: float = 0.5
+
+    # ── QA ─────────────────────────────────────────────────────────────────
+    qa_top_k: int = 5
+
+    # ── QAFD algorithm parameters ──────────────────────────────────────────
+    use_qafd: bool = True
+    qafd_alpha: float = 2.0
+    qafd_epsilon: float = 0.01
+    qafd_max_iterations: int = 500
+    qafd_weight_scheme: str = "original"     # "multiply", "add", "original"
+    qafd_use_node_degree: bool = True
+    qafd_step_size: float = 0.2
+    qafd_random_seed: int = 42
+
+    # ── Query-aware enhancements (all default = original behaviour) ────────
+    sim_mode: str = "normalized"     # Similarity contrast: "normalized", "relu", "relu_sq"
+    qa_sink_gamma: float = 0.0       # Query-aware sink capacity (0=off)
+    qa_warm_delta: float = 0.0       # Query-aware seed bias (0=off)
+    qa_warm_walk: bool = False       # Use QA edge weights in warm-start walk
+    qa_warm_steps: int = 2           # Number of warm-start steps (default 2)
+    qa_accum_gamma: float = 0.0      # Query-aware x accumulation boost (0=off)
+    qa_post_lambda: float = 0.0      # Post-diffusion reranking (0=off)
+    batch_push: bool = False         # Batch push-relabel (process all excess nodes per iter)
+
+    # ── Reranker ───────────────────────────────────────────────────────────
+    rerank_dspy_file_path: Optional[str] = None      # path to DSPy JSON; None → built-in prompt
+
+    def __post_init__(self):
+        if self.save_dir == "outputs" and self.dataset:
+            self.save_dir = f"outputs/{self.dataset}"
+
+    @property
+    def working_dir(self) -> str:
+        """Model-specific sub-directory under save_dir.
+
+        Also checks kg/multihop/ for pre-downloaded KGs from HuggingFace.
+        If found there, uses that path instead of outputs/.
+        """
+        import os
+        llm_label = self.llm_model.replace("/", "_")
+        emb_label = self.embedding_model_key.replace("/", "_")
+
+        # Check HuggingFace download location (kg/multihop/{llm}_{emb}_{dataset}/)
+        if self.dataset:
+            for task_dir in ["multihop", "ultradomain"]:
+                hf_path = os.path.join("kg", task_dir, f"{llm_label}_{emb_label}_{self.dataset}")
+                if os.path.isdir(hf_path) and os.path.exists(os.path.join(hf_path, "graph.pickle")):
+                    return hf_path
+
+        return f"{self.save_dir}/{llm_label}_{emb_label}"

src/passage_entity/embedding_store.py

@@ -0,0 +1,236 @@
+"""
+Parquet-backed embedding store, adapted from the original EmbeddingStore.
+
+Uses QAFD-RAG's async embedding functions (wrapped synchronously) so we
+can share models / GPU memory with the rest of the QAFD-RAG system.
+"""
+
+import asyncio
+import logging
+import os
+from copy import deepcopy
+from typing import List, Dict, Optional, Callable, Any
+
+import numpy as np
+import pandas as pd
+
+from .utils import compute_mdhash_id
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingModelWrapper:
+    """Thin sync wrapper around a QAFD-RAG *async* embedding function.
+
+    The wrapped function must have the signature::
+
+        async def embed(texts: list[str], **kwargs) -> np.ndarray
+
+    Parameters
+    ----------
+    embed_func : callable
+        An async embedding function from ``QAFD-RAG/src/llm.py``.
+    batch_size : int
+        Max texts per call.
+    """
+
+    def __init__(self, embed_func: Callable, batch_size: int = 16):
+        self._embed_func = embed_func
+        self.batch_size = batch_size
+
+    # ------------------------------------------------------------------
+    def batch_encode(self, texts, instruction: str = None, norm: bool = True) -> np.ndarray:
+        """Synchronously encode *texts* into embeddings."""
+        if isinstance(texts, str):
+            texts = [texts]
+        all_embeddings = []
+        # Use larger batch for API-based embeddings (OpenAI supports up to 2048)
+        effective_batch = max(self.batch_size, 512)
+        for start in range(0, len(texts), effective_batch):
+            batch = texts[start : start + effective_batch]
+            if instruction:
+                batch = [f"{instruction} {t}" for t in batch]
+            embs = self._run_async(self._embed_func(batch))
+            if not isinstance(embs, np.ndarray):
+                embs = np.array(embs)
+            if norm:
+                norms = np.linalg.norm(embs, axis=1, keepdims=True)
+                norms = np.where(norms == 0, 1, norms)
+                embs = embs / norms
+            all_embeddings.append(embs)
+        return np.vstack(all_embeddings)
+
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _run_async(coro):
+        """Run an async coroutine synchronously."""
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = None
+
+        if loop is not None and loop.is_running():
+            # We are inside an already-running event loop (e.g. Jupyter).
+            import concurrent.futures
+            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
+                return pool.submit(asyncio.run, coro).result()
+        else:
+            return asyncio.run(coro)
+
+
+class EmbeddingStore:
+    """Parquet-backed vector store.
+
+    Mirrors the original EmbeddingStore but uses ``EmbeddingModelWrapper``
+    (which calls QAFD-RAG's async embedding functions under the hood).
+    """
+
+    def __init__(
+        self,
+        embedding_model: EmbeddingModelWrapper,
+        db_filename: str,
+        batch_size: int,
+        namespace: str,
+    ):
+        self.embedding_model = embedding_model
+        self.batch_size = batch_size
+        self.namespace = namespace
+
+        if not os.path.exists(db_filename):
+            logger.info(f"Creating directory: {db_filename}")
+            os.makedirs(db_filename, exist_ok=True)
+
+        self.filename = os.path.join(db_filename, f"vdb_{self.namespace}.parquet")
+        self._load_data()
+
+    # ------------------------------------------------------------------
+    # Data persistence
+    # ------------------------------------------------------------------
+
+    def _load_data(self):
+        if os.path.exists(self.filename):
+            df = pd.read_parquet(self.filename)
+            self.hash_ids = df["hash_id"].values.tolist()
+            self.texts = df["content"].values.tolist()
+            self.embeddings = df["embedding"].values.tolist()
+            self._rebuild_indices()
+            assert len(self.hash_ids) == len(self.texts) == len(self.embeddings)
+            logger.info(f"Loaded {len(self.hash_ids)} records from {self.filename}")
+        else:
+            self.hash_ids, self.texts, self.embeddings = [], [], []
+            self.hash_id_to_idx: Dict[str, int] = {}
+            self.hash_id_to_row: Dict[str, dict] = {}
+            self.hash_id_to_text: Dict[str, str] = {}
+            self.text_to_hash_id: Dict[str, str] = {}
+
+    def _rebuild_indices(self):
+        self.hash_id_to_idx = {h: idx for idx, h in enumerate(self.hash_ids)}
+        self.hash_id_to_row = {
+            h: {"hash_id": h, "content": t} for h, t in zip(self.hash_ids, self.texts)
+        }
+        self.hash_id_to_text = {h: self.texts[idx] for idx, h in enumerate(self.hash_ids)}
+        self.text_to_hash_id = {self.texts[idx]: h for idx, h in enumerate(self.hash_ids)}
+
+    def _save_data(self):
+        data = pd.DataFrame({
+            "hash_id": self.hash_ids,
+            "content": self.texts,
+            "embedding": self.embeddings,
+        })
+        data.to_parquet(self.filename, index=False)
+        self._rebuild_indices()
+        logger.info(f"Saved {len(self.hash_ids)} records to {self.filename}")
+
+    def _upsert(self, hash_ids, texts, embeddings):
+        self.embeddings.extend(embeddings)
+        self.hash_ids.extend(hash_ids)
+        self.texts.extend(texts)
+        self._save_data()
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def get_missing_string_hash_ids(self, texts: List[str]) -> Dict[str, dict]:
+        nodes_dict = {}
+        for text in texts:
+            hid = compute_mdhash_id(text, prefix=self.namespace + "-")
+            nodes_dict[hid] = {"content": text}
+
+        if not nodes_dict:
+            return {}
+
+        existing = set(self.hash_id_to_row.keys())
+        missing = {h: {"hash_id": h, "content": v["content"]}
+                   for h, v in nodes_dict.items() if h not in existing}
+        return missing
+
+    def insert_strings(self, texts: List[str]):
+        nodes_dict = {}
+        for text in texts:
+            if not text or not text.strip():
+                continue
+            hid = compute_mdhash_id(text, prefix=self.namespace + "-")
+            nodes_dict[hid] = {"content": text}
+
+        all_ids = list(nodes_dict.keys())
+        if not all_ids:
+            return
+
+        existing = set(self.hash_id_to_row.keys())
+        missing_ids = [h for h in all_ids if h not in existing]
+
+        logger.info(
+            f"Inserting {len(missing_ids)} new records, "
+            f"{len(all_ids) - len(missing_ids)} already exist."
+        )
+        if not missing_ids:
+            return
+
+        texts_to_encode = [nodes_dict[h]["content"] for h in missing_ids]
+        missing_embeddings = self.embedding_model.batch_encode(texts_to_encode)
+        # Convert ndarray rows to list of lists for parquet storage
+        if isinstance(missing_embeddings, np.ndarray):
+            missing_embeddings = missing_embeddings.tolist()
+        self._upsert(missing_ids, texts_to_encode, missing_embeddings)
+
+    def delete(self, hash_ids):
+        indices = sorted(
+            [self.hash_id_to_idx[h] for h in hash_ids], reverse=True
+        )
+        for idx in indices:
+            self.hash_ids.pop(idx)
+            self.texts.pop(idx)
+            self.embeddings.pop(idx)
+        self._save_data()
+
+    # Lookups
+    def get_row(self, hash_id: str) -> dict:
+        return self.hash_id_to_row[hash_id]
+
+    def get_hash_id(self, text: str) -> str:
+        return self.text_to_hash_id[text]
+
+    def get_rows(self, hash_ids: List[str], dtype=np.float32) -> Dict[str, dict]:
+        if not hash_ids:
+            return {}
+        return {hid: self.hash_id_to_row[hid] for hid in hash_ids}
+
+    def get_all_ids(self) -> List[str]:
+        return deepcopy(self.hash_ids)
+
+    def get_all_id_to_rows(self) -> Dict[str, dict]:
+        return deepcopy(self.hash_id_to_row)
+
+    def get_all_texts(self) -> set:
+        return set(row["content"] for row in self.hash_id_to_row.values())
+
+    def get_embedding(self, hash_id: str, dtype=np.float32) -> np.ndarray:
+        return np.array(self.embeddings[self.hash_id_to_idx[hash_id]], dtype=dtype)
+
+    def get_embeddings(self, hash_ids: List[str], dtype=np.float32) -> np.ndarray:
+        if not hash_ids:
+            return np.array([])
+        indices = np.array([self.hash_id_to_idx[h] for h in hash_ids], dtype=np.intp)
+        all_embs = np.array(self.embeddings, dtype=dtype)
+        return all_embs[indices]

src/passage_entity/graph_adapter.py

@@ -0,0 +1,518 @@
+"""
+Bridge between igraph (the original graph format) and QAFD-RAG's flow diffusion.
+
+Provides:
+    - ``igraph_to_networkx``: convert an igraph.Graph to NetworkX (kept for
+      compatibility, but no longer used in the main retrieval path).
+    - ``IGraphQAFD``: igraph-native QAFD that matches the original
+      ``QueryAwareFlowDiffusion`` exactly — numpy arrays, C-based neighbor
+      lookups, no NetworkX conversion overhead.
+"""
+
+import logging
+import random
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+# ===========================================================================
+# igraph  -->  NetworkX  (kept for compatibility; not used in hot path)
+# ===========================================================================
+
+def igraph_to_networkx(ig_graph):
+    """Convert an igraph.Graph to a NetworkX (undirected) graph."""
+    import networkx as nx
+
+    G = nx.Graph()
+    name_attr = ig_graph.vs.attribute_names()
+    has_name = "name" in name_attr
+
+    for v in ig_graph.vs:
+        node_id = v["name"] if has_name else v.index
+        G.add_node(node_id)
+
+    has_weight = "weight" in ig_graph.es.attribute_names()
+
+    for e in ig_graph.es:
+        src = ig_graph.vs[e.source]["name"] if has_name else e.source
+        tgt = ig_graph.vs[e.target]["name"] if has_name else e.target
+        w = e["weight"] if has_weight else 1.0
+        G.add_edge(src, tgt, weight=w)
+
+    return G
+
+
+# ===========================================================================
+# igraph-native Query-Aware Flow Diffusion
+# ===========================================================================
+
+def _cosine_similarity(vec1: np.ndarray, vec2: np.ndarray, mode: str = "normalized") -> float:
+    """Cosine similarity with configurable contrast.
+
+    Modes:
+        "normalized": (cos+1)/2 → [0, 1] (original, low contrast)
+        "relu":       max(0, cos) → [0, 1] (natural contrast)
+        "relu_sq":    max(0, cos)² → [0, 1] (sharpest contrast)
+    """
+    if len(vec1) == 0 or len(vec2) == 0:
+        return 0.0
+    dot = np.dot(vec1, vec2)
+    m1 = np.linalg.norm(vec1)
+    m2 = np.linalg.norm(vec2)
+    if m1 == 0 or m2 == 0:
+        return 0.0
+    raw = dot / (m1 * m2)
+    if mode == "relu":
+        return max(0.0, raw)
+    elif mode == "relu_sq":
+        r = max(0.0, raw)
+        return r * r
+    else:  # "normalized" — original
+        return max(0.0, (raw + 1.0) / 2.0)
+
+
+class IGraphQAFD:
+    """Query-Aware Flow Diffusion directly on igraph — matches the original exactly.
+
+    Uses numpy arrays for mass/x/sink_capacity and igraph's C-based
+    ``graph.neighbors()`` for fast neighbour lookups.
+
+    Parameters
+    ----------
+    graph : igraph.Graph
+    node_name_to_idx : dict
+        Mapping from node name (str) -> vertex index (int).
+    source_weights : np.ndarray
+        Per-node seed weights (length = number of nodes). Will be normalised.
+    node_embeddings : dict
+        Mapping node_name -> np.ndarray embedding.
+    query_embedding : np.ndarray
+        Query embedding vector.
+    alpha, epsilon, max_iterations, step_size : float / int
+        Algorithm parameters.
+    weight_scheme : str
+        "original", "multiply", or "add".
+    random_seed : int
+    """
+
+    def __init__(
+        self,
+        graph,
+        node_name_to_idx: Dict[str, int],
+        source_weights: np.ndarray,
+        node_embeddings: Dict[str, np.ndarray],
+        query_embedding: Optional[np.ndarray],
+        alpha: float = 10.0,
+        epsilon: float = 1e-6,
+        max_iterations: int = 10000,
+        step_size: float = 0.2,
+        weight_scheme: str = "original",
+        hybrid_a: float = 1.0,
+        hybrid_b: float = 0.5,
+        use_node_degree: bool = True,
+        random_seed: int = 42,
+        threshold: float = 1e-5,
+        # ── Query-aware enhancements (all default OFF = original behaviour) ──
+        sim_mode: str = "normalized",   # Similarity contrast: "normalized", "relu", "relu_sq"
+        qa_sink_gamma: float = 0.0,     # query-aware sink capacity
+        qa_warm_delta: float = 0.0,     # query-aware seed bias
+        qa_warm_walk: bool = False,     # query-aware warm-start random walk (uses edge weights)
+        qa_warm_steps: int = 2,         # number of warm-start steps (default 2)
+        qa_accum_gamma: float = 0.0,    # query-aware x accumulation boost
+    ):
+        self.graph = graph
+        self.node_name_to_idx = node_name_to_idx
+        self.idx_to_node_name = {v: k for k, v in node_name_to_idx.items()}
+        self.node_embeddings = node_embeddings or {}
+        self.query_embedding = query_embedding
+        self.alpha = alpha
+        self.epsilon = epsilon
+        self.max_iterations = max_iterations
+        self.step_size = step_size
+        self.weight_scheme = weight_scheme
+        self.hybrid_a = hybrid_a
+        self.hybrid_b = hybrid_b
+        self.use_node_degree = use_node_degree
+        self.sim_mode = sim_mode
+        self.qa_sink_gamma = qa_sink_gamma
+        self.qa_warm_delta = qa_warm_delta
+        self.qa_warm_walk = qa_warm_walk
+        self.qa_accum_gamma = qa_accum_gamma
+
+        n = len(node_name_to_idx)
+
+        # Precompute per-node query similarity (used by sink/warm QA)
+        self._node_query_sim = np.zeros(n)
+        if (qa_sink_gamma > 0 or qa_warm_delta > 0) and query_embedding is not None:
+            for i in range(n):
+                name = self.idx_to_node_name.get(i)
+                if name:
+                    emb = self.node_embeddings.get(name)
+                    if emb is not None:
+                        self._node_query_sim[i] = _cosine_similarity(emb, query_embedding, mode=sim_mode)
+
+        # Normalise source weights (threshold small values, then normalise)
+        sw = np.copy(source_weights).astype(np.float64)
+        sw[sw < threshold] = 0.0
+        sw_sum = np.sum(sw)
+        if sw_sum > 0:
+            sw /= sw_sum
+        else:
+            sw = np.ones(n) / n
+        self.source_weights = sw
+
+        # State arrays
+        self.mass = np.zeros(n)
+        self.sink_capacity = np.zeros(n)
+        self.x = np.zeros(n)
+
+        # Edge weight cache
+        self._edge_weight_cache: Dict[Tuple[int, int], float] = {}
+
+        random.seed(random_seed)
+
+        # Warm-start x: multi-step lazy random walk from seed distribution
+        if qa_warm_delta > 0:
+            x = self.source_weights * (1.0 + qa_warm_delta * self._node_query_sim)
+            x_sum = np.sum(x)
+            if x_sum > 0:
+                x /= x_sum
+        else:
+            x = self.source_weights.copy()
+
+        for _ in range(qa_warm_steps):
+            x_new = np.zeros(n)
+            for i in range(n):
+                if x[i] > 0:
+                    neighbors = self.graph.neighbors(i)
+                    if not neighbors:
+                        continue
+                    if qa_warm_walk and query_embedding is not None:
+                        # Query-aware walk: spread proportional to edge weights
+                        weights = []
+                        for j in neighbors:
+                            w = self._get_edge_weight(i, j)
+                            weights.append(w)
+                        total_w = sum(weights)
+                        if total_w > 0:
+                            for j, w in zip(neighbors, weights):
+                                x_new[j] += x[i] * w / total_w
+                        else:
+                            spread = x[i] / len(neighbors)
+                            for j in neighbors:
+                                x_new[j] += spread
+                    else:
+                        # Original: uniform spread
+                        spread = x[i] / len(neighbors)
+                        for j in neighbors:
+                            x_new[j] += spread
+            x = (self.source_weights + x_new) / 2.0
+        self.x = x
+
+    # ------------------------------------------------------------------
+    def _get_edge_weight(self, i: int, j: int) -> float:
+        """Get (cached) query-aware edge weight between node indices i and j."""
+        key = (i, j)
+        if key in self._edge_weight_cache:
+            return self._edge_weight_cache[key]
+
+        try:
+            eid = self.graph.get_eid(i, j)
+            attrs = self.graph.es[eid].attributes()
+            w = attrs.get("weight", 1.0)
+        except Exception:
+            self._edge_weight_cache[key] = 0.0
+            return 0.0
+
+        if w <= 0:
+            self._edge_weight_cache[key] = 0.0
+            return 0.0
+
+        # Query-aware modulation
+        if self.weight_scheme == "none" or not self.node_embeddings or self.query_embedding is None:
+            self._edge_weight_cache[key] = w
+            return w
+
+        n1 = self.idx_to_node_name.get(i)
+        n2 = self.idx_to_node_name.get(j)
+        if n1 is None or n2 is None:
+            self._edge_weight_cache[key] = w
+            return w
+
+        e1 = self.node_embeddings.get(n1)
+        e2 = self.node_embeddings.get(n2)
+        if e1 is None and e2 is None:
+            self._edge_weight_cache[key] = w
+            return w
+
+        zero = np.zeros_like(self.query_embedding)
+        s1 = _cosine_similarity(e1 if e1 is not None else zero, self.query_embedding, mode=self.sim_mode)
+        s2 = _cosine_similarity(e2 if e2 is not None else zero, self.query_embedding, mode=self.sim_mode)
+
+        if self.weight_scheme == "multiply":
+            # Product (Eq. 5b): w * sim(u,q) * sim(v,q)
+            qw = w * s1 * s2
+        elif self.weight_scheme == "add":
+            # Mean (Eq. 5a): (w + sim(u,q) + sim(v,q)) / 3
+            qw = (w + s1 + s2) / 3.0
+        else:  # "original" = Hybrid (Eq. 5c)
+            # w * (a + b * avg_query_sim)
+            qf = (s1 + s2) / 2.0
+            qw = w * (self.hybrid_a + self.hybrid_b * qf)
+
+        self._edge_weight_cache[key] = qw
+        return qw
+
+    # ------------------------------------------------------------------
+    def _initialize(self):
+        """Set sink capacities and inject mass at seeds."""
+        n = len(self.source_weights)
+
+        if self.use_node_degree:
+            for i in range(n):
+                self.sink_capacity[i] = max(self.graph.degree(i), 1.0)
+        else:
+            self.sink_capacity[:] = 1.0
+
+        total_sink = np.sum(self.sink_capacity)
+        self.sink_capacity = 10.0 * self.sink_capacity / total_sink
+
+        # Phase 1: query-aware sink capacity — relevant nodes absorb more
+        if self.qa_sink_gamma > 0:
+            self.sink_capacity *= (1.0 + self.qa_sink_gamma * self._node_query_sim)
+
+        total_sink = np.sum(self.sink_capacity)
+
+        # Inject mass at seeds
+        self.mass[:] = 0.0
+        for i in range(n):
+            if self.source_weights[i] > 0:
+                self.mass[i] = self.alpha * total_sink * self.source_weights[i]
+
+    # ------------------------------------------------------------------
+    def _get_structural_weight(self, i: int, j: int) -> float:
+        """Get original (non-query-aware) edge weight."""
+        try:
+            eid = self.graph.get_eid(i, j)
+            return self.graph.es[eid].attributes().get("weight", 1.0)
+        except Exception:
+            return 0.0
+
+    def _push(self, node_idx: int) -> bool:
+        """Push excess mass from node to neighbours.
+
+        Decoupled accumulation/routing: x accumulates by structural degree
+        (independent of query), mass routes by query-aware edge weights.
+        This ensures query-aware modulation steers flow without penalising
+        the accumulation rate at query-relevant nodes.
+        """
+        neighbors = self.graph.neighbors(node_idx)
+        if not neighbors:
+            return False
+
+        # Query-aware weights (for routing)
+        w_qa = 0.0
+        for j in neighbors:
+            w_qa += self._get_edge_weight(node_idx, j)
+
+        if w_qa == 0:
+            return False
+
+        excess = self.mass[node_idx] - self.sink_capacity[node_idx]
+        if excess <= 0:
+            return False
+
+        # Structural weights (for accumulation) — decoupled from QA
+        w_struct = 0.0
+        for j in neighbors:
+            w_struct += self._get_structural_weight(node_idx, j)
+        if w_struct == 0:
+            w_struct = w_qa  # fallback
+
+        # Accumulate importance based on STRUCTURAL degree (not QA)
+        accum = self.step_size * excess / (w_struct + 1e-8)
+        if self.qa_accum_gamma > 0:
+            accum *= (1.0 + self.qa_accum_gamma * self._node_query_sim[node_idx])
+        self.x[node_idx] += accum
+        self.mass[node_idx] = self.sink_capacity[node_idx]
+
+        # Route mass using QUERY-AWARE weights
+        for j in neighbors:
+            w_ij = self._get_edge_weight(node_idx, j)
+            if w_ij > 0:
+                self.mass[j] += excess * w_ij / (w_qa + 1e-8)
+
+        return True
+
+    # ------------------------------------------------------------------
+    def run(self, batch_push: bool = False) -> np.ndarray:
+        """Run push-relabel flow diffusion. Returns per-node scores (np.ndarray).
+
+        batch_push: If True, process ALL excess nodes per iteration (parallel
+        push-relabel). This makes edge weights effective because each iteration
+        touches all excess nodes' edges, not just one random node's.
+        """
+        self._initialize()
+
+        iterations = 0
+        pushes = 0
+
+        while iterations < self.max_iterations:
+            iterations += 1
+
+            # Find nodes with excess mass (vectorised)
+            excess_mask = self.mass > (self.sink_capacity + self.epsilon)
+            excess_indices = np.nonzero(excess_mask)[0]
+
+            if len(excess_indices) == 0:
+                logger.info(f"QAFD converged in {iterations} iters ({pushes} pushes)")
+                break
+
+            if batch_push:
+                # Batch push: process ALL excess nodes in this iteration
+                for node_idx in excess_indices:
+                    if self._push(int(node_idx)):
+                        pushes += 1
+            else:
+                # Single push: process one random excess node (original)
+                node_idx = int(random.choice(excess_indices))
+                if self._push(node_idx):
+                    pushes += 1
+
+            if iterations % 10 == 0:
+                remaining = np.sum(np.maximum(0, self.mass - self.sink_capacity))
+                if remaining < self.epsilon:
+                    logger.info(f"QAFD converged in {iterations} iters ({pushes} pushes)")
+                    break
+
+        if iterations >= self.max_iterations:
+            logger.warning(f"QAFD did not converge after {self.max_iterations} iterations")
+
+        logger.info(f"QAFD: {iterations} iters, {pushes} pushes, batch={batch_push}")
+        return self.x
+
+
+# ===========================================================================
+# Convenience wrapper matching the interface used by retriever.py
+# ===========================================================================
+
+def run_igraph_qafd(
+    graph,
+    node_name_to_idx: Dict[str, int],
+    passage_node_idxs: List[int],
+    source_weights: np.ndarray,
+    node_embeddings: Dict[str, np.ndarray],
+    query_embedding: Optional[np.ndarray],
+    alpha: float = 10.0,
+    epsilon: float = 1e-6,
+    max_iterations: int = 10000,
+    step_size: float = 0.2,
+    weight_scheme: str = "original",
+    hybrid_a: float = 1.0,
+    hybrid_b: float = 0.5,
+    use_node_degree: bool = True,
+    random_seed: int = 42,
+    sim_mode: str = "normalized",
+    qa_sink_gamma: float = 0.0,
+    qa_warm_delta: float = 0.0,
+    qa_warm_walk: bool = False,
+    qa_warm_steps: int = 2,
+    qa_accum_gamma: float = 0.0,
+    qa_post_lambda: float = 0.0,
+    batch_push: bool = False,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Run QAFD on igraph and return (sorted_doc_ids, sorted_doc_scores).
+
+    sim_mode: Similarity contrast function ("normalized", "relu", "relu_sq")
+    Query-aware enhancement flags (all default 0.0 = original behaviour):
+        qa_sink_gamma:  Scale sink capacity by (1 + gamma * sim(node, query))
+        qa_warm_delta:  Bias warm-start x toward query-relevant seeds
+        qa_post_lambda: Rerank output by (1 + lambda * sim(passage, query))
+    """
+    qafd = IGraphQAFD(
+        graph=graph,
+        node_name_to_idx=node_name_to_idx,
+        source_weights=source_weights,
+        node_embeddings=node_embeddings,
+        query_embedding=query_embedding,
+        alpha=alpha,
+        epsilon=epsilon,
+        max_iterations=max_iterations,
+        step_size=step_size,
+        weight_scheme=weight_scheme,
+        hybrid_a=hybrid_a,
+        hybrid_b=hybrid_b,
+        use_node_degree=use_node_degree,
+        random_seed=random_seed,
+        sim_mode=sim_mode,
+        qa_sink_gamma=qa_sink_gamma,
+        qa_warm_delta=qa_warm_delta,
+        qa_warm_walk=qa_warm_walk,
+        qa_warm_steps=qa_warm_steps,
+        qa_accum_gamma=qa_accum_gamma,
+    )
+
+    node_scores = qafd.run(batch_push=batch_push)
+
+    # Extract passage scores
+    doc_scores = np.array([node_scores[idx] for idx in passage_node_idxs])
+
+    # Phase 3: post-diffusion query-aware reranking
+    if qa_post_lambda > 0 and query_embedding is not None and node_embeddings:
+        idx_to_name = qafd.idx_to_node_name
+        for pi, pidx in enumerate(passage_node_idxs):
+            name = idx_to_name.get(pidx)
+            if name:
+                emb = node_embeddings.get(name)
+                if emb is not None:
+                    sim = _cosine_similarity(emb, query_embedding, mode=sim_mode)
+                    doc_scores[pi] *= (1.0 + qa_post_lambda * sim)
+
+    total = np.sum(doc_scores)
+    if total > 0:
+        doc_scores = doc_scores / total
+    else:
+        doc_scores = np.ones(len(doc_scores)) / max(len(doc_scores), 1)
+
+    sorted_ids = np.argsort(doc_scores)[::-1]
+    sorted_scores = doc_scores[sorted_ids]
+
+    return sorted_ids, sorted_scores
+
+
+# ===========================================================================
+# Fast PPR via igraph (matches the original actual benchmark method)
+# ===========================================================================
+
+def run_ppr(
+    graph,
+    node_name_to_idx: Dict[str, int],
+    passage_node_idxs: List[int],
+    reset_prob: np.ndarray,
+    damping: float = 0.5,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Run Personalized PageRank on igraph and return (sorted_doc_ids, sorted_doc_scores).
+
+    This matches the original ``run_ppr()`` with ``use_qafd=False``.
+    Uses igraph's C-based prpack implementation — converges instantly.
+    """
+    reset_prob = np.where(np.isnan(reset_prob) | (reset_prob < 0), 0, reset_prob)
+
+    pagerank_scores = graph.personalized_pagerank(
+        vertices=range(len(node_name_to_idx)),
+        damping=damping,
+        directed=False,
+        weights="weight",
+        reset=reset_prob,
+        implementation="prpack",
+    )
+
+    doc_scores = np.array([pagerank_scores[idx] for idx in passage_node_idxs])
+
+    sorted_ids = np.argsort(doc_scores)[::-1]
+    sorted_scores = doc_scores[sorted_ids]
+
+    return sorted_ids, sorted_scores

src/passage_entity/kg_builder.py

@@ -0,0 +1,438 @@
+"""
+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-")
+            existing_keys = {info["idx"] for info in all_info}
+            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}")

src/passage_entity/openie.py

@@ -0,0 +1,231 @@
+"""
+OpenIE extraction (NER + triple extraction) using QAFD-RAG's LLM functions.
+
+Follows the original openie_openai.py logic but calls the async LLM wrappers
+from ``QAFD-RAG/src/llm.py`` synchronously via ``asyncio.run``.
+"""
+
+import asyncio
+import json
+import logging
+import re
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
+from typing import Dict, Any, List, Tuple, TypedDict, Callable
+
+from tqdm import tqdm
+
+from .prompts import make_ner_messages, make_triple_messages
+from .utils import (
+    NerRawOutput,
+    TripleRawOutput,
+    fix_broken_generated_json,
+    filter_invalid_triples,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ChunkInfo(TypedDict):
+    num_tokens: int
+    content: str
+
+
+def _run_sync(coro):
+    """Run async coroutine from sync context."""
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = None
+
+    if loop is not None and loop.is_running():
+        import concurrent.futures
+        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
+            return pool.submit(asyncio.run, coro).result()
+    else:
+        return asyncio.run(coro)
+
+
+def _extract_ner_from_response(response_text: str) -> List[str]:
+    pattern = r'\{[^{}]*"named_entities"\s*:\s*\[[^\]]*\][^{}]*\}'
+    match = re.search(pattern, response_text, re.DOTALL)
+    if match is None:
+        return []
+    try:
+        return eval(match.group())["named_entities"]
+    except Exception:
+        return []
+
+
+def _extract_triples_from_response(response_text: str) -> List[List[str]]:
+    pattern = r'\{[^{}]*"triples"\s*:\s*\[[^\]]*\][^{}]*\}'
+    match = re.search(pattern, response_text, re.DOTALL)
+    if match is None:
+        return []
+    try:
+        return eval(match.group())["triples"]
+    except Exception:
+        return []
+
+
+class OpenIE:
+    """Synchronous OpenIE using QAFD-RAG's async LLM function.
+
+    Parameters
+    ----------
+    llm_func : callable
+        An async function with the signature::
+
+            async def llm_func(prompt, system_prompt=None,
+                               history_messages=[], **kwargs) -> str
+
+        Typically one of the ``gpt_*_complete`` helpers from ``src/llm.py``.
+    """
+
+    def __init__(self, llm_func: Callable):
+        self.llm_func = llm_func
+
+    def _call_llm(self, messages: List[Dict[str, str]]) -> str:
+        """Convert chat messages to a single LLM call."""
+        system_prompt = None
+        history = []
+        user_prompt = ""
+        for msg in messages:
+            if msg["role"] == "system":
+                system_prompt = msg["content"]
+            elif msg["role"] == "assistant":
+                history.append(msg)
+            elif msg["role"] == "user":
+                # All user messages except the last go into history
+                if user_prompt:
+                    history.append({"role": "user", "content": user_prompt})
+                user_prompt = msg["content"]
+
+        return _run_sync(
+            self.llm_func(
+                prompt=user_prompt,
+                system_prompt=system_prompt,
+                history_messages=history,
+                max_tokens=2048,
+            )
+        )
+
+    # ------------------------------------------------------------------
+    def ner(self, chunk_key: str, passage: str) -> NerRawOutput:
+        messages = make_ner_messages(passage)
+        raw_response = ""
+        metadata: Dict[str, Any] = {}
+        try:
+            raw_response = self._call_llm(messages)
+            real_response = fix_broken_generated_json(raw_response)
+            extracted = _extract_ner_from_response(real_response)
+            unique_entities = list(dict.fromkeys(extracted))
+        except Exception as e:
+            logger.warning(f"NER error for chunk {chunk_key}: {e}")
+            metadata["error"] = str(e)
+            return NerRawOutput(
+                chunk_id=chunk_key,
+                response=raw_response,
+                unique_entities=[],
+                metadata=metadata,
+            )
+
+        return NerRawOutput(
+            chunk_id=chunk_key,
+            response=raw_response,
+            unique_entities=unique_entities,
+            metadata=metadata,
+        )
+
+    # ------------------------------------------------------------------
+    def triple_extraction(
+        self, chunk_key: str, passage: str, named_entities: List[str]
+    ) -> TripleRawOutput:
+        messages = make_triple_messages(passage, named_entities)
+        raw_response = ""
+        metadata: Dict[str, Any] = {}
+        try:
+            raw_response = self._call_llm(messages)
+            real_response = fix_broken_generated_json(raw_response)
+            extracted = _extract_triples_from_response(real_response)
+            triplets = filter_invalid_triples(triples=extracted)
+        except Exception as e:
+            logger.warning(f"Triple extraction error for chunk {chunk_key}: {e}")
+            metadata["error"] = str(e)
+            return TripleRawOutput(
+                chunk_id=chunk_key,
+                response=raw_response,
+                metadata=metadata,
+                triples=[],
+            )
+
+        return TripleRawOutput(
+            chunk_id=chunk_key,
+            response=raw_response,
+            metadata=metadata,
+            triples=triplets,
+        )
+
+    # ------------------------------------------------------------------
+    def openie(self, chunk_key: str, passage: str) -> Dict[str, Any]:
+        ner_output = self.ner(chunk_key=chunk_key, passage=passage)
+        triple_output = self.triple_extraction(
+            chunk_key=chunk_key,
+            passage=passage,
+            named_entities=ner_output.unique_entities,
+        )
+        return {"ner": ner_output, "triplets": triple_output}
+
+    # ------------------------------------------------------------------
+    def batch_openie(
+        self, chunks: Dict[str, dict]
+    ) -> Tuple[Dict[str, NerRawOutput], Dict[str, TripleRawOutput]]:
+        """Run NER + triple extraction over all chunks using multithreading.
+
+        Parameters
+        ----------
+        chunks : dict
+            Mapping ``chunk_hash_id -> {"content": text, ...}``.
+
+        Returns
+        -------
+        (ner_dict, triple_dict)
+        """
+        chunk_passages = {k: v["content"] for k, v in chunks.items()}
+
+        # ---- NER pass ----
+        ner_results: List[NerRawOutput] = []
+
+        with ThreadPoolExecutor() as executor:
+            ner_futures = {
+                executor.submit(self.ner, ckey, passage): ckey
+                for ckey, passage in chunk_passages.items()
+            }
+            for future in tqdm(
+                as_completed(ner_futures), total=len(ner_futures), desc="NER"
+            ):
+                ner_results.append(future.result())
+
+        # ---- Triple extraction pass ----
+        triple_results: List[TripleRawOutput] = []
+
+        with ThreadPoolExecutor() as executor:
+            re_futures = {
+                executor.submit(
+                    self.triple_extraction,
+                    nr.chunk_id,
+                    chunk_passages[nr.chunk_id],
+                    nr.unique_entities,
+                ): nr.chunk_id
+                for nr in ner_results
+            }
+            for future in tqdm(
+                as_completed(re_futures),
+                total=len(re_futures),
+                desc="Triple extraction",
+            ):
+                triple_results.append(future.result())
+
+        ner_dict = {r.chunk_id: r for r in ner_results}
+        triple_dict = {r.chunk_id: r for r in triple_results}
+        return ner_dict, triple_dict

src/passage_entity/prompts.py

@@ -0,0 +1,288 @@
+"""
+Inline prompt templates for NER, triple extraction, query NER,
+fact reranking, and RAG QA — adapted from the original template files.
+"""
+
+from string import Template
+from typing import List, Dict
+
+# ============================================================================
+# NER (passage → named entities)
+# ============================================================================
+
+NER_SYSTEM = (
+    "Your task is to extract named entities from the given paragraph. "
+    "Respond with a JSON list of entities."
+)
+
+NER_ONE_SHOT_INPUT = (
+    "Radio City\n"
+    "Radio City is India's first private FM radio station and was started on 3 July 2001.\n"
+    "It plays Hindi, English and regional songs.\n"
+    "Radio City recently forayed into New Media in May 2008 with the launch of a music "
+    "portal - PlanetRadiocity.com that offers music related news, videos, songs, and "
+    "other music-related features."
+)
+
+NER_ONE_SHOT_OUTPUT = (
+    '{"named_entities":\n'
+    '    ["Radio City", "India", "3 July 2001", "Hindi", "English", '
+    '"May 2008", "PlanetRadiocity.com"]\n'
+    '}'
+)
+
+def make_ner_messages(passage: str) -> List[Dict[str, str]]:
+    return [
+        {"role": "system", "content": NER_SYSTEM},
+        {"role": "user",   "content": NER_ONE_SHOT_INPUT},
+        {"role": "assistant", "content": NER_ONE_SHOT_OUTPUT},
+        {"role": "user",   "content": passage},
+    ]
+
+# ============================================================================
+# Query NER (question → named entities)
+# ============================================================================
+
+QUERY_NER_SYSTEM = "You're a very effective entity extraction system."
+
+QUERY_NER_ONE_SHOT_INPUT = (
+    "Please extract all named entities that are important for solving the questions below.\n"
+    "Place the named entities in json format.\n\n"
+    "Question: Which magazine was started first Arthur's Magazine or First for Women?\n"
+)
+
+QUERY_NER_ONE_SHOT_OUTPUT = (
+    '\n{"named_entities": ["First for Women", "Arthur\'s Magazine"]}\n'
+)
+
+def make_query_ner_messages(query: str) -> List[Dict[str, str]]:
+    return [
+        {"role": "system", "content": QUERY_NER_SYSTEM},
+        {"role": "user",   "content": QUERY_NER_ONE_SHOT_INPUT},
+        {"role": "assistant", "content": QUERY_NER_ONE_SHOT_OUTPUT},
+        {"role": "user",   "content": f"Question: {query}"},
+    ]
+
+# ============================================================================
+# Triple extraction (passage + entities → RDF triples)
+# ============================================================================
+
+TRIPLE_SYSTEM = (
+    "Your task is to construct an RDF (Resource Description Framework) graph from "
+    "the given passages and named entity lists. "
+    "Respond with a JSON list of triples, with each triple representing a relationship "
+    "in the RDF graph. \n\n"
+    "Pay attention to the following requirements:\n"
+    "- Each triple should contain at least one, but preferably two, of the named entities "
+    "in the list for each passage.\n"
+    "- Clearly resolve pronouns to their specific names to maintain clarity.\n"
+)
+
+_TRIPLE_FRAME = (
+    "Convert the paragraph into a JSON dict, it has a named entity list and a triple list.\n"
+    "Paragraph:\n```\n{passage}\n```\n\n{named_entity_json}\n"
+)
+
+TRIPLE_ONE_SHOT_INPUT = _TRIPLE_FRAME.format(
+    passage=NER_ONE_SHOT_INPUT,
+    named_entity_json=NER_ONE_SHOT_OUTPUT,
+)
+
+TRIPLE_ONE_SHOT_OUTPUT = (
+    '{"triples": [\n'
+    '            ["Radio City", "located in", "India"],\n'
+    '            ["Radio City", "is", "private FM radio station"],\n'
+    '            ["Radio City", "started on", "3 July 2001"],\n'
+    '            ["Radio City", "plays songs in", "Hindi"],\n'
+    '            ["Radio City", "plays songs in", "English"],\n'
+    '            ["Radio City", "forayed into", "New Media"],\n'
+    '            ["Radio City", "launched", "PlanetRadiocity.com"],\n'
+    '            ["PlanetRadiocity.com", "launched in", "May 2008"],\n'
+    '            ["PlanetRadiocity.com", "is", "music portal"],\n'
+    '            ["PlanetRadiocity.com", "offers", "news"],\n'
+    '            ["PlanetRadiocity.com", "offers", "videos"],\n'
+    '            ["PlanetRadiocity.com", "offers", "songs"]\n'
+    '    ]\n'
+    '}'
+)
+
+
+def make_triple_messages(passage: str, named_entities: List[str]) -> List[Dict[str, str]]:
+    import json as _json
+    named_entity_json = _json.dumps({"named_entities": named_entities})
+    user_content = _TRIPLE_FRAME.format(passage=passage, named_entity_json=named_entity_json)
+    return [
+        {"role": "system", "content": TRIPLE_SYSTEM},
+        {"role": "user",   "content": TRIPLE_ONE_SHOT_INPUT},
+        {"role": "assistant", "content": TRIPLE_ONE_SHOT_OUTPUT},
+        {"role": "user",   "content": user_content},
+    ]
+
+# ============================================================================
+# Fact reranker / filter  (DSPy-style prompt from the original pipeline)
+# ============================================================================
+
+RERANKER_SYSTEM = (
+    "Your input fields are:\n"
+    "1. `question` (str): Query for retrieval\n"
+    "2. `fact_before_filter` (str): Candidate facts to be filtered\n\n"
+    "Your output fields are:\n"
+    '1. `fact_after_filter` (Fact): Filtered facts in JSON format\n\n'
+    "All interactions will be structured in the following way, with the appropriate "
+    "values filled in.\n\n"
+    "[[ ## question ## ]]\n{question}\n\n"
+    "[[ ## fact_before_filter ## ]]\n{fact_before_filter}\n\n"
+    "[[ ## fact_after_filter ## ]]\n{fact_after_filter}        "
+    '# note: the value you produce must be pareseable according to the following JSON schema: '
+    '{"type": "object", "properties": {"fact": {"type": "array", '
+    '"description": "A list of facts, each fact is a list of 3 strings: [subject, predicate, object]", '
+    '"items": {"type": "array", "items": {"type": "string"}}, '
+    '"title": "Fact"}}, "required": ["fact"], "title": "Fact"}\n\n'
+    "[[ ## completed ## ]]\n\n"
+    "In adhering to this structure, your objective is: \n"
+    "        You are a critical component of a high-stakes question-answering system used by "
+    "top researchers and decision-makers worldwide. Your task is to filter facts based on their "
+    "relevance to a given query, ensuring that the most crucial information is presented to "
+    "these stakeholders. The query requires careful analysis and possibly multi-hop reasoning "
+    "to connect different pieces of information. You must select up to 4 relevant facts from "
+    "the provided candidate list that have a strong connection to the query, aiding in reasoning "
+    "and providing an accurate answer. The output should be in JSON format, e.g., "
+    '{"fact": [["s1", "p1", "o1"], ["s2", "p2", "o2"]]}, and if no facts are relevant, '
+    'return an empty list, {"fact": []}. The accuracy of your response is paramount, as it '
+    "will directly impact the decisions made by these high-level stakeholders. You must only "
+    "use facts from the candidate list and not generate new facts. The future of critical "
+    "decision-making relies on your ability to accurately filter and present relevant information."
+)
+
+RERANKER_INPUT_TEMPLATE = (
+    "[[ ## question ## ]]\n{question}\n\n"
+    "[[ ## fact_before_filter ## ]]\n{fact_before_filter}\n\n"
+    "Respond with the corresponding output fields, starting with the field "
+    "`[[ ## fact_after_filter ## ]]` (must be formatted as a valid Python Fact), "
+    "and then ending with the marker for `[[ ## completed ## ]]`."
+)
+
+RERANKER_OUTPUT_TEMPLATE = (
+    "[[ ## fact_after_filter ## ]]\n{fact_after_filter}\n\n"
+    "[[ ## completed ## ]]"
+)
+
+# Built-in few-shot demos (from the original filter_default_prompt.py)
+RERANKER_DEMOS = [
+    {
+        "question": "Are Imperial River (Florida) and Amaradia (Dolj) both located in the same country?",
+        "fact_before_filter": '{"fact": [["imperial river", "is located in", "florida"], ["imperial river", "is a river in", "united states"], ["imperial river", "may refer to", "south america"], ["amaradia", "flows through", "ro ia de amaradia"], ["imperial river", "may refer to", "united states"]]}',
+        "fact_after_filter": '{"fact":[["imperial river","is located in","florida"],["imperial river","is a river in","united states"],["amaradia","flows through","ro ia de amaradia"]]}',
+    },
+    {
+        "question": "When is the director of film The Ancestor 's birthday?",
+        "fact_before_filter": '{"fact": [["jean jacques annaud", "born on", "1 october 1943"], ["tsui hark", "born on", "15 february 1950"], ["pablo trapero", "born on", "4 october 1971"], ["the ancestor", "directed by", "guido brignone"], ["benh zeitlin", "born on", "october 14  1982"]]}',
+        "fact_after_filter": '{"fact":[["the ancestor","directed by","guido brignone"]]}',
+    },
+    {
+        "question": "In what geographic region is the country where Teafuone is located?",
+        "fact_before_filter": '{"fact": [["teafuaniua", "is on the", "east"], ["motuloa", "lies between", "teafuaniua"], ["motuloa", "lies between", "teafuanonu"], ["teafuone", "is", "islet"], ["teafuone", "located in", "nukufetau"]]}',
+        "fact_after_filter": '{"fact":[["teafuone","is","islet"],["teafuone","located in","nukufetau"]]}',
+    },
+    {
+        "question": "When did the director of film S.O.B. (Film) die?",
+        "fact_before_filter": '{"fact": [["allan dwan", "died on", "28 december 1981"], ["s o b", "written and directed by", "blake edwards"], ["robert aldrich", "died on", "december 5  1983"], ["robert siodmak", "died on", "10 march 1973"], ["bernardo bertolucci", "died on", "26 november 2018"]]}',
+        "fact_after_filter": '{"fact":[["s o b","written and directed by","blake edwards"]]}',
+    },
+]
+
+
+def make_reranker_messages(question: str, fact_before_filter_json: str) -> List[Dict[str, str]]:
+    """Build the full chat history for fact reranking."""
+    messages = [{"role": "system", "content": RERANKER_SYSTEM}]
+    for demo in RERANKER_DEMOS:
+        messages.append({
+            "role": "user",
+            "content": RERANKER_INPUT_TEMPLATE.format(
+                question=demo["question"],
+                fact_before_filter=demo["fact_before_filter"],
+            ),
+        })
+        messages.append({
+            "role": "assistant",
+            "content": RERANKER_OUTPUT_TEMPLATE.format(
+                fact_after_filter=demo["fact_after_filter"],
+            ),
+        })
+    messages.append({
+        "role": "user",
+        "content": RERANKER_INPUT_TEMPLATE.format(
+            question=question,
+            fact_before_filter=fact_before_filter_json,
+        ),
+    })
+    return messages
+
+# ============================================================================
+# RAG QA prompt (MuSiQue-style, also used for HotpotQA / 2Wiki)
+# ============================================================================
+
+RAG_QA_SYSTEM = (
+    "As an advanced reading comprehension assistant, your task is to analyze text passages "
+    "and corresponding questions meticulously. "
+    "Your response start after \"Thought: \", where you will methodically break down the "
+    "reasoning process, illustrating how you arrive at conclusions. "
+    "Conclude with \"Answer: \" to present a concise, definitive response, devoid of "
+    "additional elaborations."
+)
+
+_RAG_QA_ONE_SHOT_DOCS = (
+    "Wikipedia Title: The Last Horse\n"
+    "The Last Horse (Spanish:El último caballo) is a 1950 Spanish comedy film directed "
+    "by Edgar Neville starring Fernando Fernán Gómez.\n\n"
+    "Wikipedia Title: Southampton\n"
+    "The University of Southampton, which was founded in 1862 and received its Royal "
+    "Charter as a university in 1952, has over 22,000 students. The university is ranked "
+    "in the top 100 research universities in the world in the Academic Ranking of World "
+    "Universities 2010.\n\n"
+    "Wikipedia Title: Neville A. Stanton\n"
+    "Neville A. Stanton is a British Professor of Human Factors and Ergonomics at the "
+    "University of Southampton. Prof Stanton is a Chartered Engineer (C.Eng), Chartered "
+    "Psychologist (C.Psychol) and Chartered Ergonomist (C.ErgHF).\n"
+)
+
+RAG_QA_ONE_SHOT_INPUT = (
+    f"{_RAG_QA_ONE_SHOT_DOCS}\n\n"
+    "Question: When was Neville A. Stanton's employer founded?\nThought: "
+)
+
+RAG_QA_ONE_SHOT_OUTPUT = (
+    "The employer of Neville A. Stanton is University of Southampton. "
+    "The University of Southampton was founded in 1862. "
+    "\nAnswer: 1862."
+)
+
+
+def make_qa_messages(passages: List[str], question: str) -> List[Dict[str, str]]:
+    """Build QA chat messages from retrieved passages and question."""
+    prompt_user = ""
+    for p in passages:
+        prompt_user += f"Wikipedia Title: {p}\n\n"
+    prompt_user += f"Question: {question}\nThought: "
+    return [
+        {"role": "system", "content": RAG_QA_SYSTEM},
+        {"role": "user",   "content": RAG_QA_ONE_SHOT_INPUT},
+        {"role": "assistant", "content": RAG_QA_ONE_SHOT_OUTPUT},
+        {"role": "user",   "content": prompt_user},
+    ]
+
+# ============================================================================
+# Query instruction strings (for embedding model)
+# ============================================================================
+
+QUERY_INSTRUCTIONS = {
+    "ner_to_node": "Given a phrase, retrieve synonymous or relevant phrases that best match this phrase.",
+    "query_to_node": "Given a question, retrieve relevant phrases that are mentioned in this question.",
+    "query_to_fact": "Given a question, retrieve relevant triplet facts that matches this question.",
+    "query_to_sentence": "Given a question, retrieve relevant sentences that best answer the question.",
+    "query_to_passage": "Given a question, retrieve relevant documents that best answer the question.",
+}
+
+def get_query_instruction(linking_method: str) -> str:
+    default = "Given a question, retrieve relevant documents that best answer the question."
+    return QUERY_INSTRUCTIONS.get(linking_method, default)

src/passage_entity/reranker.py

@@ -0,0 +1,217 @@
+"""
+Fact reranker following the original DSPy filter logic.
+
+Uses QAFD-RAG's async LLM functions (wrapped synchronously) to call the
+same prompt structure that the original DSPyFilter uses.
+"""
+
+import ast
+import asyncio
+import difflib
+import json
+import logging
+import re
+from copy import deepcopy
+from typing import Callable, Dict, List, Tuple, Any
+
+from .prompts import make_reranker_messages
+
+logger = logging.getLogger(__name__)
+
+
+def _run_sync(coro):
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = None
+    if loop is not None and loop.is_running():
+        import concurrent.futures
+        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
+            return pool.submit(asyncio.run, coro).result()
+    else:
+        return asyncio.run(coro)
+
+
+class FactReranker:
+    """Rerank candidate fact triples using an LLM (DSPy-style filtering).
+
+    Parameters
+    ----------
+    llm_func : callable
+        Async LLM function from ``src/llm.py``.
+    dspy_file_path : str or None
+        Path to a DSPy-saved JSON file with custom demos/system prompt.
+        If ``None``, uses built-in demos from ``prompts.py``.
+    """
+
+    def __init__(self, llm_func: Callable, dspy_file_path: str = None):
+        self.llm_func = llm_func
+        self.dspy_file_path = dspy_file_path
+
+        if dspy_file_path is not None:
+            self._custom_template = self._load_dspy_template(dspy_file_path)
+        else:
+            self._custom_template = None
+
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _load_dspy_template(path: str) -> List[Dict[str, str]]:
+        """Load a DSPy-saved JSON and convert to chat messages."""
+        data = json.load(open(path))
+        system_prompt = data["prog"]["system"]
+        demos = data["prog"]["demos"]
+
+        one_in = (
+            "[[ ## question ## ]]\n{question}\n\n"
+            "[[ ## fact_before_filter ## ]]\n{fact_before_filter}\n\n"
+            "Respond with the corresponding output fields, starting with the field "
+            "`[[ ## fact_after_filter ## ]]` (must be formatted as a valid Python Fact), "
+            "and then ending with the marker for `[[ ## completed ## ]]`."
+        )
+        one_out = (
+            "[[ ## fact_after_filter ## ]]\n{fact_after_filter}\n\n"
+            "[[ ## completed ## ]]"
+        )
+
+        msgs = [{"role": "system", "content": system_prompt}]
+        for demo in demos:
+            msgs.append({
+                "role": "user",
+                "content": one_in.format(
+                    question=demo["question"],
+                    fact_before_filter=demo["fact_before_filter"],
+                ),
+            })
+            if "fact_after_filter" in demo:
+                msgs.append({
+                    "role": "assistant",
+                    "content": one_out.format(
+                        fact_after_filter=demo["fact_after_filter"],
+                    ),
+                })
+        return msgs
+
+    # ------------------------------------------------------------------
+    def _build_messages(
+        self, question: str, fact_before_filter_json: str
+    ) -> List[Dict[str, str]]:
+        if self._custom_template is not None:
+            msgs = deepcopy(self._custom_template)
+            one_in = (
+                "[[ ## question ## ]]\n{question}\n\n"
+                "[[ ## fact_before_filter ## ]]\n{fact_before_filter}\n\n"
+                "Respond with the corresponding output fields, starting with the field "
+                "`[[ ## fact_after_filter ## ]]` (must be formatted as a valid Python Fact), "
+                "and then ending with the marker for `[[ ## completed ## ]]`."
+            )
+            msgs.append({
+                "role": "user",
+                "content": one_in.format(
+                    question=question,
+                    fact_before_filter=fact_before_filter_json,
+                ),
+            })
+            return msgs
+        else:
+            return make_reranker_messages(question, fact_before_filter_json)
+
+    # ------------------------------------------------------------------
+    def _call_llm(self, messages: List[Dict[str, str]]) -> str:
+        system_prompt = None
+        history = []
+        user_prompt = ""
+        for msg in messages:
+            if msg["role"] == "system":
+                system_prompt = msg["content"]
+            elif msg["role"] == "assistant":
+                history.append(msg)
+            elif msg["role"] == "user":
+                if user_prompt:
+                    history.append({"role": "user", "content": user_prompt})
+                user_prompt = msg["content"]
+
+        return _run_sync(
+            self.llm_func(
+                prompt=user_prompt,
+                system_prompt=system_prompt,
+                history_messages=history,
+                max_tokens=512,
+            )
+        )
+
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _parse_filter(response: str) -> List[List[str]]:
+        """Extract fact_after_filter from the DSPy-style response."""
+        sections = [(None, [])]
+        header_re = re.compile(r'\[\[ ## (\w+) ## \]\]')
+        for line in response.splitlines():
+            m = header_re.match(line.strip())
+            if m:
+                sections.append((m.group(1), []))
+            else:
+                sections[-1][1].append(line)
+        sections = [(k, "\n".join(v).strip()) for k, v in sections]
+
+        parsed: List[List[str]] = []
+        for k, value in sections:
+            if k == "fact_after_filter":
+                try:
+                    try:
+                        pv = json.loads(value)
+                    except json.JSONDecodeError:
+                        try:
+                            pv = ast.literal_eval(value)
+                        except (ValueError, SyntaxError):
+                            pv = value
+                    if isinstance(pv, dict) and "fact" in pv:
+                        parsed = pv["fact"]
+                except Exception as e:
+                    logger.warning(f"Error parsing reranker output: {e}")
+        return parsed
+
+    # ------------------------------------------------------------------
+    def rerank(
+        self,
+        query: str,
+        candidate_items: List[Tuple],
+        candidate_indices: List[int],
+        len_after_rerank: int = None,
+    ) -> Tuple[List[int], List[Tuple], dict]:
+        """Rerank candidate facts by LLM-based filtering.
+
+        Returns
+        -------
+        (sorted_indices, sorted_facts, metadata_dict)
+        """
+        fact_json = json.dumps({"fact": [list(c) for c in candidate_items]})
+        try:
+            msgs = self._build_messages(query, fact_json)
+            response = self._call_llm(msgs)
+            generated_facts = self._parse_filter(response)
+        except Exception as e:
+            logger.warning(f"Reranker exception: {e}")
+            generated_facts = []
+
+        result_indices = []
+        for gf in generated_facts:
+            matches = difflib.get_close_matches(
+                str(gf), [str(i) for i in candidate_items], n=1, cutoff=0.0
+            )
+            if matches:
+                try:
+                    result_indices.append(candidate_items.index(eval(matches[0])))
+                except Exception as e:
+                    logger.warning(f"Index matching error: {e}")
+
+        sorted_indices = [candidate_indices[i] for i in result_indices]
+        sorted_items = [candidate_items[i] for i in result_indices]
+
+        if len_after_rerank is not None:
+            sorted_indices = sorted_indices[:len_after_rerank]
+            sorted_items = sorted_items[:len_after_rerank]
+
+        return sorted_indices, sorted_items, {"confidence": None}
+
+    def __call__(self, *args, **kwargs):
+        return self.rerank(*args, **kwargs)

src/passage_entity/retriever.py

@@ -0,0 +1,414 @@
+"""
+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:
+                logger.info("No facts after reranking -> fallback to DPR")
+                sorted_ids, sorted_scores = self._dense_passage_retrieval(q)
+            else:
+                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

src/passage_entity/utils.py

@@ -0,0 +1,221 @@
+"""
+Utility functions and data classes for the passage-entity KG pipeline.
+
+Adapted from the original misc_utils.py and llm_utils.py.
+"""
+
+import json
+import re
+import logging
+from dataclasses import dataclass
+from hashlib import md5
+from typing import Dict, Any, List, Tuple, Literal, Union, Optional
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Data classes
+# ---------------------------------------------------------------------------
+
+@dataclass
+class NerRawOutput:
+    chunk_id: str
+    response: str
+    unique_entities: List[str]
+    metadata: Dict[str, Any]
+
+
+@dataclass
+class TripleRawOutput:
+    chunk_id: str
+    response: str
+    triples: List[List[str]]
+    metadata: Dict[str, Any]
+
+
+@dataclass
+class QuerySolution:
+    question: str
+    docs: List[str]
+    doc_scores: np.ndarray = None
+    answer: str = None
+    gold_answers: List[str] = None
+    gold_docs: Optional[List[str]] = None
+
+    def to_dict(self):
+        return {
+            "question": self.question,
+            "answer": self.answer,
+            "gold_answers": self.gold_answers,
+            "docs": self.docs[:5],
+            "doc_scores": (
+                [round(v, 4) for v in self.doc_scores.tolist()[:5]]
+                if self.doc_scores is not None
+                else None
+            ),
+            "gold_docs": self.gold_docs,
+        }
+
+
+Triple = Union[List[str], Tuple[str, str, str]]
+
+# ---------------------------------------------------------------------------
+# Hashing
+# ---------------------------------------------------------------------------
+
+def compute_mdhash_id(content: str, prefix: str = "") -> str:
+    """Compute the MD5 hash of *content* and optionally prepend *prefix*."""
+    return prefix + md5(content.encode()).hexdigest()
+
+# ---------------------------------------------------------------------------
+# Text processing
+# ---------------------------------------------------------------------------
+
+def text_processing(text):
+    """Lower-case, strip non-alphanumeric characters (except spaces)."""
+    if isinstance(text, list):
+        return [text_processing(t) for t in text]
+    if not isinstance(text, str):
+        text = str(text)
+    return re.sub('[^A-Za-z0-9 ]', ' ', text.lower()).strip()
+
+# ---------------------------------------------------------------------------
+# OpenIE helpers
+# ---------------------------------------------------------------------------
+
+def extract_entity_nodes(chunk_triples: List[List[Triple]]) -> Tuple[List[str], List[List[str]]]:
+    """Extract unique entity nodes from chunk triples.
+
+    Returns:
+        graph_nodes: globally unique list of entity strings.
+        chunk_triple_entities: per-chunk list of entity strings.
+    """
+    chunk_triple_entities = []
+    for triples in chunk_triples:
+        triple_entities = set()
+        for t in triples:
+            if len(t) == 3:
+                triple_entities.update([t[0], t[2]])
+            else:
+                logger.warning(f"Invalid triple during graph construction: {t}")
+        chunk_triple_entities.append(list(triple_entities))
+    graph_nodes = list(np.unique([ent for ents in chunk_triple_entities for ent in ents]))
+    return graph_nodes, chunk_triple_entities
+
+
+def flatten_facts(chunk_triples: List[List[Triple]]) -> List[Tuple]:
+    """Flatten per-chunk triples into a unique list of tuples."""
+    graph_triples = []
+    for triples in chunk_triples:
+        graph_triples.extend([tuple(t) for t in triples])
+    return list(set(graph_triples))
+
+
+def reformat_openie_results(corpus_openie_results):
+    """Convert saved openie JSON list into (ner_dict, triple_dict)."""
+    ner_output_dict = {
+        chunk_item['idx']: NerRawOutput(
+            chunk_id=chunk_item['idx'],
+            response=None,
+            metadata={},
+            unique_entities=list(np.unique(chunk_item['extracted_entities']))
+        )
+        for chunk_item in corpus_openie_results
+    }
+    triple_output_dict = {
+        chunk_item['idx']: TripleRawOutput(
+            chunk_id=chunk_item['idx'],
+            response=None,
+            metadata={},
+            triples=filter_invalid_triples(triples=chunk_item['extracted_triples'])
+        )
+        for chunk_item in corpus_openie_results
+    }
+    return ner_output_dict, triple_output_dict
+
+# ---------------------------------------------------------------------------
+# Normalization
+# ---------------------------------------------------------------------------
+
+def min_max_normalize(x: np.ndarray) -> np.ndarray:
+    min_val = np.min(x)
+    max_val = np.max(x)
+    range_val = max_val - min_val
+    if range_val == 0:
+        return np.ones_like(x)
+    return (x - min_val) / range_val
+
+# ---------------------------------------------------------------------------
+# JSON repair helpers (from the original pipeline llm_utils)
+# ---------------------------------------------------------------------------
+
+def fix_broken_generated_json(json_str: str) -> str:
+    """Attempt to fix truncated JSON by closing open brackets/braces."""
+    def find_unclosed(s):
+        unclosed = []
+        inside_string = False
+        escape_next = False
+        for char in s:
+            if inside_string:
+                if escape_next:
+                    escape_next = False
+                elif char == '\\':
+                    escape_next = True
+                elif char == '"':
+                    inside_string = False
+            else:
+                if char == '"':
+                    inside_string = True
+                elif char in '{[':
+                    unclosed.append(char)
+                elif char in '}]':
+                    if unclosed and (
+                        (char == '}' and unclosed[-1] == '{') or
+                        (char == ']' and unclosed[-1] == '[')
+                    ):
+                        unclosed.pop()
+        return unclosed
+
+    try:
+        json.loads(json_str)
+        return json_str
+    except json.JSONDecodeError:
+        pass
+
+    last_comma_index = json_str.rfind(',')
+    if last_comma_index != -1:
+        json_str = json_str[:last_comma_index]
+
+    unclosed = find_unclosed(json_str)
+    closing_map = {'{': '}', '[': ']'}
+    for open_char in reversed(unclosed):
+        json_str += closing_map[open_char]
+    return json_str
+
+
+def filter_invalid_triples(triples: List[List[str]]) -> List[List[str]]:
+    """Keep only unique triples with exactly 3 elements."""
+    unique_triples = set()
+    valid_triples = []
+    for triple in triples:
+        if len(triple) != 3:
+            continue
+        valid_triple = [str(item) for item in triple]
+        key = tuple(valid_triple)
+        if key not in unique_triples:
+            unique_triples.add(key)
+            valid_triples.append(valid_triple)
+    return valid_triples
+
+
+def string_to_bool(v) -> bool:
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ("yes", "true", "t", "y", "1"):
+        return True
+    elif v.lower() in ("no", "false", "f", "n", "0"):
+        return False
+    else:
+        raise ValueError(f"Cannot convert {v!r} to bool")