benchmarks/summarization/benchmark_summarization.py

#!/usr/bin/env python3
"""
SQuALITY Benchmark for QAFD_RAG
Multi-reference BLEU/ROUGE/METEOR metrics + LLM quality evaluation
"""

import os
import sys
import asyncio
import time
import json
import logging
import nest_asyncio
from typing import Dict, List
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path

# 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)
logging.getLogger("absl").setLevel(logging.ERROR)

QAFD_RAG_HOME = str(Path(__file__).parent.parent.parent)
sys.path.insert(0, QAFD_RAG_HOME)

nest_asyncio.apply()


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()

@dataclass
class BenchmarkResult:
    """Benchmark results"""
    model_name: str
    dataset_name: str
    success: bool = False
    total_time: float = 0.0
    kg_build_time: float = 0.0
    query_time: float = 0.0
    num_questions: int = 0
    metrics: Dict[str, float] = None  # BLEU, ROUGE, METEOR
    quality_scores: Dict[str, List[float]] = None  # Comprehensiveness, Diversity, etc.
    error_message: str = ""
    responses: List[str] = None  # Generated summaries
    questions: List[str] = None  # Original questions
    reference_answers: List[List[str]] = None  # Reference answers per question

class SQuALITYBenchmark:
    def __init__(self, api_key: str, embedding_model: str = "openai-small", llm_model: str = "gpt-4o-mini"):
        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_llm_func(self):
        """Get LLM function based on model name"""
        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 _get_working_dir(self, dataset_name: str) -> str:
        """Get working directory for dataset"""
        return os.path.join(QAFD_RAG_HOME, "kg", "summarization", f"{self.llm_model}_{self.embedding_model}_{dataset_name}")

    def _kg_exists(self, working_dir: str) -> bool:
        """Check if KG already exists"""
        # Check for key files that indicate a built KG
        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)

    async def _benchmark_qafd(self, dataset: List[Dict], dataset_name: str,
                            question_count: int, max_documents: int = None,
                            mode: str = "hybrid", max_source_nodes: int = 40,
                            min_flow_threshold: float = 0.01, alpha: float = 5.0,
                            evaluator: str = "mini", force_build: bool = False) -> BenchmarkResult:
        """Benchmark QAFD_RAG on SQuALITY"""
        try:
            from src.QAFD_RAG import QAFD_RAG, QueryParam

            working_dir = self._get_working_dir(dataset_name)

            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,
            )

            # Get unique documents and selected questions
            unique_docs = {}
            selected_questions = []

            for item in dataset:
                passage_id = item["passage_id"]

                if passage_id not in unique_docs:
                    if max_documents and len(unique_docs) >= max_documents:
                        continue
                    unique_docs[passage_id] = item["document"]

                if passage_id in unique_docs:
                    selected_questions.append(item)
                    if len(selected_questions) >= question_count:
                        break

            dataset = selected_questions

            # 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(f"\n  Building KG from {len(unique_docs)} documents...")
                start_time = time.time()

                for i, (passage_id, document) in enumerate(unique_docs.items()):
                    print_progress(i + 1, len(unique_docs), "Progress")
                    await rag.ainsert(document)

                kg_build_time = time.time() - start_time
                print(f"  KG built in {kg_build_time:.2f}s")

            # Query phase
            print(f"\n  Running queries...")
            start_time = time.time()
            responses = []

            for i in range(min(question_count, len(dataset))):
                print_progress(i + 1, min(question_count, len(dataset)), "Progress")
                question = dataset[i]["question"]

                query_param = QueryParam(
                    mode=mode,
                    max_source_nodes=max_source_nodes,
                    min_flow_threshold=min_flow_threshold,
                    alpha=alpha
                )
                response = await rag.aquery(question, query_param)
                responses.append(response)

            query_time = time.time() - start_time

            # Calculate metrics with MULTIPLE REFERENCES
            print(f"\n  Calculating reference metrics...", end=" ", flush=True)
            metrics = self._calculate_metrics(responses, dataset[:len(responses)])
            print("done")

            # Calculate quality scores
            print(f"  Evaluating quality ({len(responses)} responses)...")
            quality_scores = await self._evaluate_quality(dataset, responses, evaluator)
            
            # Collect questions and reference answers for output
            questions_list = [dataset[i]["question"] for i in range(len(responses))]
            references_list = [dataset[i].get("all_answers", [dataset[i]["answer"]]) for i in range(len(responses))]

            return BenchmarkResult(
                model_name="QAFD_RAG",
                dataset_name=dataset_name,
                success=True,
                total_time=kg_build_time + query_time,
                kg_build_time=kg_build_time,
                query_time=query_time,
                num_questions=len(responses),
                metrics=metrics,
                quality_scores=quality_scores,
                responses=responses,
                questions=questions_list,
                reference_answers=references_list,
            )

        except Exception as e:
            print(f"  ERROR: {str(e)}")
            import traceback
            traceback.print_exc()
            return BenchmarkResult(
                model_name="QAFD_RAG",
                dataset_name=dataset_name,
                success=False,
                error_message=str(e)
            )
    
    def _calculate_metrics(self, responses: List[str], dataset: List[Dict]) -> Dict[str, float]:
        """Calculate BLEU, ROUGE, METEOR metrics with MULTIPLE REFERENCES"""
        try:
            from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
            from nltk.translate.meteor_score import meteor_score
            from nltk.tokenize import word_tokenize
            from rouge_score import rouge_scorer
            import nltk
            
            nltk.download('wordnet', quiet=True)
            nltk.download('omw-1.4', quiet=True)
            nltk.download('punkt', quiet=True)
            
            scorer = rouge_scorer.RougeScorer(['rouge1', 'rouge2'], use_stemmer=True)
            smoothing = SmoothingFunction().method1
            
            bleu1_scores = []
            bleu2_scores = []
            rouge1_scores = []
            rouge2_scores = []
            meteor_scores = []
            
            for i, response in enumerate(responses):
                if not response:
                    continue
                
                # Get ALL reference answers for this question
                all_refs = dataset[i].get("all_answers", [dataset[i]["answer"]])
                
                # Tokenize all references
                try:
                    ref_tokens_list = [word_tokenize(ref.lower()) for ref in all_refs]
                    resp_tokens = word_tokenize(response.lower())
                except:
                    ref_tokens_list = [ref.lower().split() for ref in all_refs]
                    resp_tokens = response.lower().split()
                
                # BLEU with multiple references
                bleu1 = sentence_bleu(ref_tokens_list, resp_tokens, weights=(1,0,0,0), smoothing_function=smoothing)
                bleu2 = sentence_bleu(ref_tokens_list, resp_tokens, weights=(0.5,0.5,0,0), smoothing_function=smoothing)
                bleu1_scores.append(bleu1)
                bleu2_scores.append(bleu2)
                
                # ROUGE uses first reference (standard)
                primary_ref = all_refs[0]
                rouge_result = scorer.score(primary_ref, response)
                rouge1_scores.append(rouge_result['rouge1'].fmeasure)
                rouge2_scores.append(rouge_result['rouge2'].fmeasure)
                
                # METEOR with multiple references
                try:
                    meteor = meteor_score(ref_tokens_list, resp_tokens)
                    meteor_scores.append(meteor)
                except:
                    meteor_scores.append(0.0)
            
            return {
                "bleu_1": sum(bleu1_scores) / len(bleu1_scores) if bleu1_scores else 0,
                "bleu_2": sum(bleu2_scores) / len(bleu2_scores) if bleu2_scores else 0,
                "rouge_1_f1": sum(rouge1_scores) / len(rouge1_scores) if rouge1_scores else 0,
                "rouge_2_f1": sum(rouge2_scores) / len(rouge2_scores) if rouge2_scores else 0,
                "meteor": sum(meteor_scores) / len(meteor_scores) if meteor_scores else 0,
            }
            
        except Exception as e:
            print(f"  Warning: Metrics calculation failed: {e}")
            import traceback
            traceback.print_exc()
            return {}
    
    async def _evaluate_quality(self, dataset: List[Dict], responses: List[str],
                               evaluator: str = "mini") -> Dict[str, List[float]]:
        """Evaluate response quality using GPT-4o or GPT-4o-mini"""
        if not responses:
            return {}

        try:
            if evaluator == "4o":
                from src.llm import gpt_4o_complete as eval_complete
            else:
                from src.llm import gpt_4o_mini_complete as eval_complete

            criteria = ["comprehensiveness", "diversity", "logicality", "relevance", "coherence"]
            response_scores = {criterion: [] for criterion in criteria}

            for i in range(len(responses)):
                print_progress(i + 1, len(responses), "Evaluating")
                
                question = dataset[i]["question"]
                response = responses[i]
                
                if not response:
                    for criterion in criteria:
                        response_scores[criterion].append(0.0)
                    continue
                
                response_criterion_scores = {criterion: [] for criterion in criteria}
                for _ in range(5):
                    prompt = f"""Evaluate the following response based on five criteria. Rate each from 0-100.

Question: {question}
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?

JSON format:
{{
    "comprehensiveness": [score],
    "diversity": [score],
    "logicality": [score],
    "relevance": [score],
    "coherence": [score]
}}"""

                    evaluation_result = await eval_complete(prompt, max_tokens=200)
                    
                    import re
                    json_match = re.search(r'\{.*\}', evaluation_result, re.DOTALL)
                    if json_match:
                        try:
                            scores = json.loads(json_match.group())
                            for criterion in criteria:
                                if criterion in scores:
                                    score_value = float(scores[criterion])
                                    if 0 <= score_value <= 100:
                                        response_criterion_scores[criterion].append(score_value)
                        except:
                            pass
                
                for criterion in criteria:
                    if response_criterion_scores[criterion]:
                        avg = sum(response_criterion_scores[criterion]) / len(response_criterion_scores[criterion])
                        response_scores[criterion].append(avg)
                    else:
                        response_scores[criterion].append(0.0)
            
            return response_scores
        except Exception as e:
            print(f"  Warning: Quality evaluation failed: {e}")
            return {}
    
    def _download_dataset(self) -> List[Dict]:
        """Download SQuALITY dataset - COLLECTS ALL REFERENCE ANSWERS"""
        try:
            from datasets import load_dataset
            import warnings
            warnings.filterwarnings("ignore", category=FutureWarning)
            logging.getLogger("datasets").setLevel(logging.WARNING)
        except ImportError:
            print("ERROR: Install datasets with: pip install datasets")
            sys.exit(1)

        dataset = load_dataset("pszemraj/SQuALITY-v1.3", split="train")
        
        prepared_data = []
        for i, item in enumerate(dataset):
            document = item.get("document", "")
            questions = item.get("questions", [])
            metadata = item.get("metadata", {})
            passage_id = metadata.get("passage_id", f"passage_{i}") if isinstance(metadata, dict) else f"passage_{i}"
            
            if not document or not questions:
                continue
            
            for q in questions:
                question_text = q.get("question_text", "")
                responses = q.get("responses", [])
                
                if responses and len(responses) > 0:
                    # COLLECT ALL REFERENCE ANSWERS (KEY CHANGE!)
                    answer_texts = [r.get("response_text", "") for r in responses if r.get("response_text", "")]
                    
                    if question_text and answer_texts:
                        prepared_data.append({
                            "document": document,
                            "question": question_text,
                            "answer": answer_texts[0],  # Primary answer
                            "all_answers": answer_texts,  # All references for BLEU
                            "passage_id": passage_id
                        })
        
        return prepared_data
    
    async def build_kg(self, dataset_name: str = "squality", max_documents: int = None) -> bool:
        """Build KG only (no benchmark)"""
        print_header("QAFD-RAG Knowledge Graph Builder")
        print_config({
            "Dataset": "SQuALITY",
            "Embedding": self.embedding_model,
            "LLM": self.llm_model,
            "Max Documents": str(max_documents) if max_documents else "all"
        })

        print("\n  Downloading dataset...", end=" ", flush=True)
        dataset = self._download_dataset()
        print("done")

        try:
            from src.QAFD_RAG import QAFD_RAG

            working_dir = self._get_working_dir(dataset_name)
            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,
            )

            # Get unique documents
            unique_docs = {}
            for item in dataset:
                passage_id = item["passage_id"]
                if passage_id not in unique_docs:
                    if max_documents and len(unique_docs) >= max_documents:
                        break
                    unique_docs[passage_id] = item["document"]

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

            for i, (passage_id, document) in enumerate(unique_docs.items()):
                print_progress(i + 1, len(unique_docs), "Progress")
                await rag.ainsert(document)

            build_time = time.time() - start_time

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

        except Exception as e:
            print(f"\n  ERROR: {e}")
            return False

    async def run_benchmark(self, dataset_name: str = "squality", question_count: int = 250, max_documents: int = None,
                           mode: str = "hybrid", max_source_nodes: int = 40,
                           min_flow_threshold: float = 0.01, alpha: float = 10.0,
                           evaluator: str = "mini", force_build: bool = False) -> BenchmarkResult:
        """Run benchmark"""
        print_header("QAFD-RAG SQuALITY Benchmark")
        print_config({
            "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),
            "Evaluator": "GPT-4o" if evaluator == "4o" else "GPT-4o-mini"
        })

        print("\n  Downloading dataset...", end=" ", flush=True)
        dataset = self._download_dataset()
        print("done")

        if len(dataset) < question_count:
            question_count = len(dataset)

        result = await self._benchmark_qafd(dataset, dataset_name, question_count, max_documents,
                                           mode, max_source_nodes, min_flow_threshold, alpha,
                                           evaluator, force_build)
        return result
    
    def print_results(self, result: BenchmarkResult):
        """Print results"""
        print_header("Results: SQuALITY")

        if not result.success:
            print(f"\n  Status: FAILED")
            print(f"  Error: {result.error_message}")
            return

        print("\n  PERFORMANCE")
        print(f"  {'─' * 40}")
        print(f"  {'Questions':<25} {result.num_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"  {'Total Time':<25} {result.total_time:.2f}s")
        if result.num_questions > 0:
            print(f"  {'Avg per Question':<25} {result.query_time/result.num_questions:.2f}s")

        if result.metrics:
            print("\n  REFERENCE METRICS")
            print(f"  {'─' * 40}")
            for k, v in result.metrics.items():
                print(f"  {k:<25} {v*100:.2f}%")

        if result.quality_scores:
            print("\n  QUALITY METRICS")
            print(f"  {'─' * 40}")
            import statistics
            for k, v in result.quality_scores.items():
                if v:
                    avg = sum(v) / len(v)
                    std = statistics.stdev(v) if len(v) > 1 else 0
                    print(f"  {k:<25} {avg:.2f} ± {std:.2f}")

            all_scores = []
            for scores in result.quality_scores.values():
                all_scores.extend(scores)
            if all_scores:
                overall = sum(all_scores) / len(all_scores)
                print(f"\n  {'Overall Average':<25} {overall:.2f}")
        print()
    
    def save_results(self, result: BenchmarkResult, dataset_name: str = "squality"):
        """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", "summarization", dataset_name)
        os.makedirs(results_dir, exist_ok=True)

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

        # --- Eval file: metrics, quality scores, timing ---
        quality_averages = {}
        quality_stds = {}
        if result.quality_scores:
            import statistics
            for criterion, scores in result.quality_scores.items():
                if scores:
                    quality_averages[criterion] = sum(scores) / len(scores)
                    quality_stds[criterion] = statistics.stdev(scores) if len(scores) > 1 else 0.0

        eval_data = {
            "timestamp": datetime.now().isoformat(),
            "model": result.model_name,
            "llm": self.llm_model,
            "embedding": self.embedding_model,
            "success": result.success,
            "performance": {
                "num_questions": result.num_questions,
                "kg_build_time": result.kg_build_time,
                "query_time": result.query_time,
                "total_time": result.total_time,
                "avg_per_question": result.query_time / result.num_questions if result.num_questions > 0 else 0
            },
            "reference_metrics": result.metrics,
            "quality_scores": quality_averages,
            "quality_stds": quality_stds,
            "quality_raw": result.quality_scores,
            "error": result.error_message
        }

        with open(eval_file, 'w') as f:
            json.dump(eval_data, f, indent=2)

        # --- Responses file: questions + generated summaries + references ---
        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_summary": response,
                    "reference_answers": result.reference_answers[i] if result.reference_answers else [],
                }
                output_entries.append(entry)

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

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

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

async def main():
    """Main entry point"""
    import argparse
    
    parser = argparse.ArgumentParser(description="SQuALITY Benchmark for QAFD_RAG")
    parser.add_argument("--questions", type=int, default=250,
                       help="Number of questions (default: 250)")
    parser.add_argument("--max-documents", type=int, default=None,
                       help="Max documents (default: unlimited)")
    parser.add_argument("--mode", type=str, default="hybrid",
                       choices=["local", "global", "hybrid"],
                       help="Query mode (default: hybrid)")
    parser.add_argument("--max-source-nodes", type=int, default=40,
                       help="Max source nodes (default: 40)")
    parser.add_argument("--min-flow-threshold", type=float, default=0.01,
                       help="Min flow threshold (default: 0.01)")
    parser.add_argument("--alpha", type=float, default=10.0,
                       help="Alpha (default: 10.0)")
    parser.add_argument("--evaluator", type=str, default="mini",
                       choices=["mini", "4o"],
                       help="Evaluator model: mini (gpt-4o-mini) or 4o (gpt-4o) (default: mini)")
    parser.add_argument("--api-key", type=str,
                       help="OpenAI API key")
    parser.add_argument("--build", action="store_true",
                       help="Build KG only (no benchmark)")
    parser.add_argument("--force-build", action="store_true",
                       help="Force rebuild KG even if it exists")
    parser.add_argument("--embedding", type=str, default="openai-small",
                       choices=["openai-small", "openai-large", "jina-v3", "gritlm", "nvidia-nv-embed-v2"],
                       help="Embedding model (default: openai-small)")
    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"],
                       help="LLM model (default: gpt-4o-mini)")
    parser.add_argument("--dataset", type=str, default="squality",
                       help="Dataset name (default: squality)")

    args = parser.parse_args()

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

    benchmark = SQuALITYBenchmark(api_key, args.embedding, args.llm)

    # Build-only mode
    if args.build:
        await benchmark.build_kg(dataset_name=args.dataset, max_documents=args.max_documents)
        return

    result = await benchmark.run_benchmark(
        dataset_name=args.dataset,
        question_count=args.questions,
        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,
        evaluator=args.evaluator,
        force_build=args.force_build
    )
    benchmark.print_results(result)
    benchmark.save_results(result, dataset_name=args.dataset)

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