benchmarks/multihop/benchmark_multihop.py

#!/usr/bin/env python3
"""
Unified Multi-hop QA Benchmark for QAFD-RAG
Supports: MuSiQue, HotpotQA, 2WikiMultiHopQA
"""

import os
import sys
import asyncio
import time
import json
import re
import string
import logging
from typing import List, Dict
from dataclasses import dataclass, asdict
from datetime import datetime
from pathlib import Path
from collections import Counter
import numpy as np

# Suppress verbose logging
logging.getLogger("httpx").setLevel(logging.ERROR)
logging.getLogger("QAFD_RAG").setLevel(logging.ERROR)
logging.getLogger("nano-vectordb").setLevel(logging.ERROR)
logging.getLogger("OpenAI").setLevel(logging.ERROR)

# Add QAFD-RAG to path
QAFD_RAG_HOME = str(Path(__file__).parent.parent.parent)
sys.path.insert(0, QAFD_RAG_HOME)

import nest_asyncio
nest_asyncio.apply()

# Dataset configurations
DATASETS = {
    "musique": {
        "name": "MuSiQue",
        "data_file": "musique.json",
        "corpus_file": "musique_corpus.json",
        "kg_dir": "musique",
    },
    "hotpotqa": {
        "name": "HotpotQA",
        "data_file": "hotpotqa.json",
        "corpus_file": "hotpotqa_corpus.json",
        "kg_dir": "hotpotqa",
    },
    "2wikimultihopqa": {
        "name": "2WikiMultiHopQA",
        "data_file": "2wikimultihopqa.json",
        "corpus_file": "2wikimultihopqa_corpus.json",
        "kg_dir": "2wikimultihopqa",
    },
}


def print_header(title: str, width: int = 70):
    """Print a formatted header"""
    print(f"\n{'=' * width}")
    print(f"  {title}")
    print(f"{'=' * width}")


def print_config(items: Dict[str, str], width: int = 70):
    """Print configuration items"""
    print(f"{'─' * width}")
    for key, value in items.items():
        print(f"  {key:<20} {value}")
    print(f"{'─' * width}")


def print_progress(current: int, total: int, prefix: str = "", width: int = 40):
    """Print a progress bar"""
    percent = current / total
    filled = int(width * percent)
    bar = '█' * filled + '░' * (width - filled)
    print(f"\r  {prefix} [{bar}] {current}/{total} ({percent*100:.1f}%)", end='', flush=True)
    if current == total:
        print()


def print_metric(name: str, value: float, std: float = None, width: int = 25):
    """Print a metric with optional std"""
    if std is not None:
        print(f"  {name:<{width}} {value:.4f} ± {std:.4f}")
    else:
        print(f"  {name:<{width}} {value:.4f}")


@dataclass
class BenchmarkResult:
    """Benchmark result for multi-hop QA"""
    model_name: str
    dataset_name: str
    total_questions: int
    success_count: int
    total_time: float
    kg_build_time: float
    query_time: float
    avg_time_per_question: float
    f1_score_mean: float
    f1_score_std: float
    exact_match_mean: float
    exact_match_std: float
    f1_scores: List[float] = None
    exact_match_scores: List[float] = None
    responses: List[str] = None
    questions: List[str] = None
    gold_answers: List[List[str]] = None
    error_message: str = ""


def normalize_answer(answer: str) -> str:
    """Normalize answer for comparison"""
    def remove_articles(text):
        return re.sub(r"\b(a|an|the)\b", " ", text)
    def white_space_fix(text):
        return " ".join(text.split())
    def remove_punc(text):
        exclude = set(string.punctuation)
        return "".join(ch for ch in text if ch not in exclude)
    def lower(text):
        return text.lower()
    return white_space_fix(remove_articles(remove_punc(lower(answer))))


def compute_f1(gold: str, predicted: str) -> float:
    """Compute F1 score between gold and predicted answers"""
    gold_tokens = normalize_answer(gold).split()
    predicted_tokens = normalize_answer(predicted).split()
    common = Counter(predicted_tokens) & Counter(gold_tokens)
    num_same = sum(common.values())
    if num_same == 0:
        return 0.0
    precision = 1.0 * num_same / len(predicted_tokens) if predicted_tokens else 0.0
    recall = 1.0 * num_same / len(gold_tokens) if gold_tokens else 0.0
    if precision + recall == 0:
        return 0.0
    return 2 * (precision * recall) / (precision + recall)


def compute_exact_match(gold: str, predicted: str) -> float:
    """Compute exact match score"""
    return 1.0 if normalize_answer(gold) == normalize_answer(predicted) else 0.0


def get_gold_answers(samples):
    """Extract gold answers from samples"""
    gold_answers = []
    for sample in samples:
        if 'answer' in sample:
            gold_ans = sample['answer']
        elif 'reference' in sample:
            gold_ans = sample['reference']
        else:
            gold_ans = "Unknown"
        if isinstance(gold_ans, str):
            gold_ans = [gold_ans]
        elif not isinstance(gold_ans, list):
            gold_ans = [str(gold_ans)]
        gold_answers.append(gold_ans)
    return gold_answers


class MultiHopBenchmark:
    def __init__(self, dataset: str, api_key: str, embedding_model: str = "openai-small",
                 llm_model: str = "gpt-4o-mini"):
        if dataset not in DATASETS:
            raise ValueError(f"Unknown dataset: {dataset}. Choose from: {list(DATASETS.keys())}")

        self.dataset = dataset
        self.config = DATASETS[dataset]
        self.api_key = api_key
        self.embedding_model = embedding_model
        self.llm_model = llm_model
        os.environ["OPENAI_API_KEY"] = api_key
        os.environ["OPENAI_API_BASE"] = "https://api.openai.com/v1"

    def _get_working_dir(self) -> str:
        return os.path.join(QAFD_RAG_HOME, "kg", "multihop", f"{self.llm_model}_{self.embedding_model}_{self.config['kg_dir']}")

    def _kg_exists(self, working_dir: str) -> bool:
        kg_files = [
            os.path.join(working_dir, "vdb_entities.json"),
            os.path.join(working_dir, "vdb_chunks.json"),
            os.path.join(working_dir, "kv_store_full_docs.json"),
        ]
        return all(os.path.exists(f) for f in kg_files)

    def _get_llm_func(self):
        from src import llm
        llm_funcs = {
            "gpt-4o-mini": llm.gpt_4o_mini_complete,
            "gpt-4o": llm.gpt_4o_complete,
            "gpt-oss-120b": llm.gpt_oss_120b_complete,
            "gpt-5": llm.gpt_5_complete,
            "gpt-5-mini": llm.gpt_5_mini_complete,
            "gpt-5-nano": llm.gpt_5_nano_complete,
        }
        return llm_funcs.get(self.llm_model, llm.gpt_4o_mini_complete)

    def _ensure_data_file(self, filename: str) -> str:
        """Return path to data file, downloading from HuggingFace if missing."""
        data_dir = os.path.join(QAFD_RAG_HOME, "data", "multihop")
        filepath = os.path.join(data_dir, filename)
        if not os.path.exists(filepath):
            print(f"  Downloading {filename} from HuggingFace...", end=" ", flush=True)
            from huggingface_hub import hf_hub_download
            os.makedirs(data_dir, exist_ok=True)
            hf_hub_download(
                repo_id="osunlp/HippoRAG",  # Dataset source
                filename=filename,
                repo_type="dataset",
                local_dir=data_dir,
            )
            print("done")
        return filepath

    def _load_dataset(self) -> List[Dict]:
        dataset_path = self._ensure_data_file(self.config["data_file"])
        with open(dataset_path, 'r', encoding='utf-8') as f:
            samples = json.load(f)
        return samples

    def _load_corpus(self) -> List[str]:
        corpus_path = self._ensure_data_file(self.config["corpus_file"])
        with open(corpus_path, 'r', encoding='utf-8') as f:
            corpus = json.load(f)
        docs = [f"{doc['title']}\n{doc['text']}" for doc in corpus]
        return docs

    async def build_kg(self, max_documents: int = None) -> bool:
        """Build KG only (no benchmark)"""
        from src.QAFD_RAG import QAFD_RAG

        print_header(f"QAFD-RAG Knowledge Graph Builder")
        print_config({
            "Graph Type": "entity",
            "Dataset": self.config['name'],
            "Embedding": self.embedding_model,
            "LLM": self.llm_model,
            "Working Dir": self._get_working_dir()
        })

        working_dir = self._get_working_dir()
        os.makedirs(working_dir, exist_ok=True)

        llm_func = self._get_llm_func()
        rag = QAFD_RAG(
            working_dir=working_dir,
            llm_model_func=llm_func,
            llm_model_name=self.llm_model,
            embedding_model_key=self.embedding_model,
            enable_llm_cache=True,
        )

        print("\n  Loading corpus...", end=" ", flush=True)
        docs = self._load_corpus()
        docs_to_process = min(max_documents, len(docs)) if max_documents else len(docs)
        print(f"done ({len(docs)} documents available)")

        print(f"\n  Building KG from {docs_to_process} documents...")
        start_time = time.time()

        for i, doc in enumerate(docs[:docs_to_process]):
            print_progress(i + 1, docs_to_process, "Progress")
            await rag.ainsert(doc)

        build_time = time.time() - start_time

        print_header("Build Complete")
        print(f"  Documents processed:  {docs_to_process}")
        print(f"  Time elapsed:         {build_time:.2f}s")
        print(f"  Avg per document:     {build_time/docs_to_process:.2f}s")
        print(f"  Output directory:     {working_dir}")
        print()

        return True

    async def run_benchmark(self, question_count: int = 100, force_build: bool = False,
                           max_documents: int = None, mode: str = "hybrid",
                           max_source_nodes: int = 20, min_flow_threshold: float = 0.1,
                           alpha: float = 2.0) -> BenchmarkResult:
        """Run benchmark"""
        from src.QAFD_RAG import QAFD_RAG, QueryParam

        print_header(f"QAFD-RAG Multi-hop QA Benchmark")
        print_config({
            "Graph Type": "entity",
            "Dataset": self.config['name'],
            "Questions": str(question_count),
            "Embedding": self.embedding_model,
            "LLM": self.llm_model,
            "Mode": mode,
            "Max Nodes": str(max_source_nodes),
            "Threshold": str(min_flow_threshold),
            "Alpha": str(alpha)
        })

        working_dir = self._get_working_dir()
        os.makedirs(working_dir, exist_ok=True)

        llm_func = self._get_llm_func()
        rag = QAFD_RAG(
            working_dir=working_dir,
            llm_model_func=llm_func,
            llm_model_name=self.llm_model,
            embedding_model_key=self.embedding_model,
            enable_llm_cache=True,
        )

        # Check if KG exists or needs to be built
        kg_build_time = 0.0
        if self._kg_exists(working_dir) and not force_build:
            print(f"\n  Using existing KG: {working_dir}")
        else:
            print("\n  Loading corpus...", end=" ", flush=True)
            docs = self._load_corpus()
            docs_to_process = min(max_documents, len(docs)) if max_documents else len(docs)
            print(f"done ({docs_to_process} documents)")

            print(f"  Building KG...")
            start_time = time.time()
            for i, doc in enumerate(docs[:docs_to_process]):
                print_progress(i + 1, docs_to_process, "Progress")
                await rag.ainsert(doc)
            kg_build_time = time.time() - start_time
            print(f"  KG built in {kg_build_time:.2f}s")

        # Load dataset
        print("\n  Loading dataset...", end=" ", flush=True)
        samples = self._load_dataset()
        samples = samples[:question_count]
        questions = [s['question'] for s in samples]
        gold_answers = get_gold_answers(samples)
        print(f"done ({len(questions)} questions)")

        # Run queries
        print(f"\n  Running queries...")
        start_time = time.time()
        responses = []
        success_count = 0

        for i, question in enumerate(questions):
            try:
                print_progress(i + 1, len(questions), "Progress")
                query_param = QueryParam(
                    mode=mode,
                    max_source_nodes=max_source_nodes,
                    min_flow_threshold=min_flow_threshold,
                    alpha=alpha,
                    response_type="Brief, accurate answer (maximum 14 words)."
                )
                response = await rag.aquery(question, query_param)
                if response and len(response.split()) > 14:
                    response = " ".join(response.split()[:14])
                responses.append(response)
                success_count += 1
            except Exception as e:
                responses.append("")

        query_time = time.time() - start_time

        # Calculate metrics
        f1_scores = []
        em_scores = []
        for gold_list, predicted in zip(gold_answers, responses):
            if not predicted:
                f1_scores.append(0.0)
                em_scores.append(0.0)
                continue
            f1_scores.append(max(compute_f1(g, predicted) for g in gold_list))
            em_scores.append(max(compute_exact_match(g, predicted) for g in gold_list))

        result = BenchmarkResult(
            model_name="QAFD_RAG",
            dataset_name=self.dataset,
            total_questions=len(questions),
            success_count=success_count,
            total_time=kg_build_time + query_time,
            kg_build_time=kg_build_time,
            query_time=query_time,
            avg_time_per_question=query_time / len(questions) if questions else 0,
            f1_score_mean=float(np.mean(f1_scores)) if f1_scores else 0,
            f1_score_std=float(np.std(f1_scores)) if f1_scores else 0,
            exact_match_mean=float(np.mean(em_scores)) if em_scores else 0,
            exact_match_std=float(np.std(em_scores)) if em_scores else 0,
            f1_scores=f1_scores,
            exact_match_scores=em_scores,
            responses=responses,
            questions=questions,
            gold_answers=gold_answers,
        )

        self.print_results(result)
        return result

    def print_results(self, result: BenchmarkResult):
        """Print benchmark results"""
        print_header(f"Results: {self.config['name']}")

        print("\n  PERFORMANCE")
        print(f"  {'─' * 40}")
        print(f"  {'Questions':<25} {result.total_questions}")
        print(f"  {'Successful':<25} {result.success_count}/{result.total_questions}")
        print(f"  {'KG Build Time':<25} {result.kg_build_time:.2f}s")
        print(f"  {'Query Time':<25} {result.query_time:.2f}s")
        print(f"  {'Avg per Question':<25} {result.avg_time_per_question:.2f}s")

        print("\n  ACCURACY METRICS")
        print(f"  {'─' * 40}")
        print_metric("F1 Score", result.f1_score_mean, result.f1_score_std)
        print_metric("Exact Match", result.exact_match_mean, result.exact_match_std)
        print()

    def save_results(self, result: BenchmarkResult):
        """Save results as two separate files: eval metrics and generated responses"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        results_dir = os.path.join(QAFD_RAG_HOME, "results", "multihop", self.dataset)
        os.makedirs(results_dir, exist_ok=True)

        eval_file = os.path.join(results_dir, f"{self.dataset}_{timestamp}_eval.json")
        output_file = os.path.join(results_dir, f"{self.dataset}_{timestamp}_responses.json")

        # --- Eval file: metrics and timing ---
        eval_data = {
            "timestamp": datetime.now().isoformat(),
            "graph_type": "entity",
            "model": result.model_name,
            "llm": self.llm_model,
            "embedding": self.embedding_model,
            "dataset": result.dataset_name,
            "performance": {
                "total_questions": result.total_questions,
                "success_count": result.success_count,
                "kg_build_time": result.kg_build_time,
                "query_time": result.query_time,
                "total_time": result.total_time,
                "avg_time_per_question": result.avg_time_per_question,
            },
            "metrics": {
                "f1_score_mean": result.f1_score_mean,
                "f1_score_std": result.f1_score_std,
                "exact_match_mean": result.exact_match_mean,
                "exact_match_std": result.exact_match_std,
            },
            "per_question_f1": result.f1_scores,
            "per_question_em": result.exact_match_scores,
            "error": result.error_message,
        }

        with open(eval_file, 'w', encoding='utf-8') as f:
            json.dump(eval_data, f, indent=2, ensure_ascii=False)

        # --- Responses file: questions + generated answers + gold answers ---
        output_entries = []
        if result.responses:
            for i, response in enumerate(result.responses):
                entry = {
                    "id": i + 1,
                    "question": result.questions[i] if result.questions else "",
                    "generated_answer": response,
                    "gold_answers": result.gold_answers[i] if result.gold_answers else [],
                }
                output_entries.append(entry)

        output_data = {
            "timestamp": datetime.now().isoformat(),
            "model": result.model_name,
            "llm": self.llm_model,
            "embedding": self.embedding_model,
            "dataset": result.dataset_name,
            "num_responses": len(output_entries),
            "responses": output_entries,
        }

        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(output_data, f, indent=2, ensure_ascii=False)

        print(f"  Eval saved:      {eval_file}")
        print(f"  Responses saved: {output_file}\n")


async def main():
    import argparse

    parser = argparse.ArgumentParser(description="QAFD_RAG Multi-hop QA Benchmark")
    parser.add_argument("--dataset", type=str, required=True,
                        choices=["musique", "hotpotqa", "2wikimultihopqa"])
    parser.add_argument("--questions", type=int, default=100)
    parser.add_argument("--max-documents", type=int, default=None)
    parser.add_argument("--build", action="store_true")
    parser.add_argument("--force-build", action="store_true")
    parser.add_argument("--embedding", type=str, default="openai-small",
                        choices=["openai-small", "openai-large", "jina-v3", "gritlm", "nvidia-nv-embed-v2"])
    parser.add_argument("--llm", type=str, default="gpt-4o-mini",
                        choices=["gpt-4o-mini", "gpt-4o", "gpt-oss-120b", "gpt-5", "gpt-5-mini", "gpt-5-nano"])
    parser.add_argument("--mode", type=str, default="hybrid",
                        choices=["local", "global", "hybrid"])
    parser.add_argument("--max-source-nodes", type=int, default=20)
    parser.add_argument("--min-flow-threshold", type=float, default=0.1)
    parser.add_argument("--alpha", type=float, default=2.0)

    args = parser.parse_args()

    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        print("ERROR: Set OPENAI_API_KEY environment variable")
        return

    benchmark = MultiHopBenchmark(args.dataset, api_key, args.embedding, args.llm)

    if args.build:
        await benchmark.build_kg(max_documents=args.max_documents)
        return

    result = await benchmark.run_benchmark(
        question_count=args.questions,
        force_build=args.force_build,
        max_documents=args.max_documents,
        mode=args.mode,
        max_source_nodes=args.max_source_nodes,
        min_flow_threshold=args.min_flow_threshold,
        alpha=args.alpha
    )
    benchmark.save_results(result)


if __name__ == "__main__":
    asyncio.run(main())