Remove old folder name

src/hipporag_pipeline/benchmark_runner.py

@@ -1,737 +0,0 @@
-"""
-Entry point for running multihop benchmarks (MuSiQue, HotpotQA, 2WikiMultiHopQA)
-with the HippoRAG-style KG pipeline + QAFD inside QAFD-RAG.
-
-Usage::
-
-    python -m src.hipporag_pipeline.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.hipporag_pipeline.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.hipporag_pipeline"]:
-    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.hipporag_pipeline.config import HippoRAGConfig
-from src.hipporag_pipeline.embedding_store import EmbeddingModelWrapper
-from src.hipporag_pipeline.kg_builder import KGBuilder
-from src.hipporag_pipeline.openie import OpenIE
-from src.hipporag_pipeline.reranker import FactReranker
-from src.hipporag_pipeline.retriever import HippoRAGRetriever
-from src.hipporag_pipeline.prompts import make_qa_messages
-from src.hipporag_pipeline.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 HippoRAG 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="HippoRAG + QAFD 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 = HippoRAGConfig(
-        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 = HippoRAGRetriever(
-        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/hipporag_pipeline/config.py

@@ -1,94 +0,0 @@
-"""
-Configuration dataclass for the HippoRAG-style KG pipeline inside QAFD-RAG.
-
-Combines HippoRAG BaseConfig fields with QAFD algorithm parameters.
-"""
-
-from dataclasses import dataclass, field
-from typing import Optional, Literal
-
-
-@dataclass
-class HippoRAGConfig:
-    """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/hipporag_pipeline/embedding_store.py

@@ -1,236 +0,0 @@
-"""
-Parquet-backed embedding store, adapted from HippoRAG's 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 HippoRAG's 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/hipporag_pipeline/graph_adapter.py

@@ -1,518 +0,0 @@
-"""
-Bridge between igraph (HippoRAG's 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 HippoRAG's
-      ``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 HippoRAG 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 HippoRAG's 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 HippoRAG's ``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/hipporag_pipeline/kg_builder.py

@@ -1,438 +0,0 @@
-"""
-Knowledge graph builder following HippoRAG's ``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 HippoRAGConfig
-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 HippoRAG's 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 HippoRAG-style knowledge graph and persist it to disk."""
-
-    def __init__(
-        self,
-        config: HippoRAGConfig,
-        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/hipporag_pipeline/openie.py

@@ -1,231 +0,0 @@
-"""
-OpenIE extraction (NER + triple extraction) using QAFD-RAG's LLM functions.
-
-Follows HippoRAG's 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/hipporag_pipeline/prompts.py

@@ -1,288 +0,0 @@
-"""
-Inline prompt templates for NER, triple extraction, query NER,
-fact reranking, and RAG QA — adapted from HippoRAG's 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 HippoRAG)
-# ============================================================================
-
-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 HippoRAG's 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/hipporag_pipeline/reranker.py

@@ -1,217 +0,0 @@
-"""
-Fact reranker following HippoRAG's DSPy filter logic.
-
-Uses QAFD-RAG's async LLM functions (wrapped synchronously) to call the
-same prompt structure that HippoRAG's 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/hipporag_pipeline/retriever.py

@@ -1,414 +0,0 @@
-"""
-Full retrieval pipeline following HippoRAG's 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 HippoRAGConfig
-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 HippoRAGRetriever:
-    """End-to-end retriever:  query -> ranked passages.
-
-    Uses the pre-built KG (igraph), embedding stores, and QAFD flow diffusion.
-    """
-
-    def __init__(
-        self,
-        config: HippoRAGConfig,
-        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 HippoRAG 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 HippoRAG 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/hipporag_pipeline/utils.py

@@ -1,221 +0,0 @@
-"""
-Utility functions and data classes for the HippoRAG-style KG pipeline.
-
-Adapted from HippoRAG's 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 HippoRAG 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")