muthuk1/graphrag-inference-hackathon
1
1"""2Benchmark Runner — Runs all 3 pipelines on HotpotQA and evaluates3==================================================================4Pipeline 1: LLM-Only (no retrieval)5Pipeline 2: Basic RAG (vector search + LLM)6Pipeline 3: GraphRAG (TigerGraph + novelty engine)7 8Evaluates with: F1, EM, LLM-as-a-Judge, BERTScore, Context Hit Rate9"""10import json11import logging12from typing import Dict, List, Optional13from .layers.orchestration_layer import InferenceOrchestrator14from .layers.evaluation_layer import (15 EvaluationLayer, EvalSample, compute_bertscore16)17 18logger = logging.getLogger(__name__)19 20 21class BenchmarkRunner:22 """Runs benchmarks on HotpotQA with all 3 pipelines and generates comparison metrics."""23 24 def __init__(self, orchestrator, evaluator):25 self.orchestrator = orchestrator26 self.evaluator = evaluator27 self.benchmark_results = []28 self.eval_samples: List[EvalSample] = []29 30 def run_hotpotqa_benchmark(self, num_samples=100, split="validation",31 top_k=5, hops=2, progress_callback=None,32 run_judge=True, run_bertscore=True):33 """Run all 3 pipelines on HotpotQA and evaluate."""34 from datasets import load_dataset35 logger.info(f"Loading HotpotQA ({split}, n={num_samples})...")36 ds = load_dataset("hotpotqa/hotpot_qa", "distractor", split=split)37 38 results = []39 self.eval_samples = []40 41 for idx in range(min(num_samples, len(ds))):42 row = ds[idx]43 query, gold = row["question"], row["answer"]44 qtype = row.get("type", "unknown")45 level = row.get("level", "unknown")46 47 # Build passages from context48 passages = [f"{t}: {' '.join(s)}"49 for t, s in zip(row["context"]["title"], row["context"]["sentences"])]50 51 # Extract supporting facts for context hit rate52 sf = []53 for t, si in zip(row["supporting_facts"]["title"], row["supporting_facts"]["sent_id"]):54 for ct, cs in zip(row["context"]["title"], row["context"]["sentences"]):55 if ct == t and si < len(cs):56 sf.append(cs[si])57 58 try:59 # Run all 3 pipelines60 lo = self.orchestrator.run_llm_only(query)61 b = self.orchestrator.run_baseline_rag(query, passages, top_k)62 g = self.orchestrator.run_graphrag(query, passages, hops=hops)63 64 sample = EvalSample(65 query=query, reference_answer=gold,66 llm_only_answer=lo.answer,67 baseline_answer=b.answer,68 graphrag_answer=g.answer,69 baseline_contexts=b.contexts,70 graphrag_contexts=g.contexts,71 question_type=qtype, difficulty=str(level),72 supporting_facts=sf)73 self.eval_samples.append(sample)74 75 er = self.evaluator.evaluate_sample(76 sample,77 llm_only_tokens=lo.total_tokens,78 baseline_tokens=b.total_tokens,79 graphrag_tokens=g.total_tokens,80 llm_only_cost=lo.cost_usd,81 baseline_cost=b.cost_usd,82 graphrag_cost=g.cost_usd,83 llm_only_latency=lo.latency_ms,84 baseline_latency=b.latency_ms,85 graphrag_latency=g.latency_ms,86 run_judge=run_judge,87 )88 89 rd = {90 "idx": idx, "query": query, "gold_answer": gold,91 "question_type": qtype, "level": level,92 # Answers93 "llm_only_answer": lo.answer,94 "baseline_answer": b.answer,95 "graphrag_answer": g.answer,96 # F1 / EM97 "llm_only_f1": er.llm_only_f1,98 "baseline_f1": er.baseline_f1,99 "graphrag_f1": er.graphrag_f1,100 "llm_only_em": er.llm_only_em,101 "baseline_em": er.baseline_em,102 "graphrag_em": er.graphrag_em,103 # LLM-as-Judge104 "llm_only_judge": er.llm_only_judge,105 "baseline_judge": er.baseline_judge,106 "graphrag_judge": er.graphrag_judge,107 # Tokens / Cost / Latency108 "llm_only_tokens": lo.total_tokens,109 "baseline_tokens": b.total_tokens,110 "graphrag_tokens": g.total_tokens,111 "llm_only_cost": lo.cost_usd,112 "baseline_cost": b.cost_usd,113 "graphrag_cost": g.cost_usd,114 "llm_only_latency": lo.latency_ms,115 "baseline_latency": b.latency_ms,116 "graphrag_latency": g.latency_ms,117 # Context118 "baseline_context_hit": er.baseline_context_hit,119 "graphrag_context_hit": er.graphrag_context_hit,120 "entities_found": len(g.entities_found),121 "relations_traversed": len(g.relations_traversed),122 }123 results.append(rd)124 self.benchmark_results.append(rd)125 126 if progress_callback:127 progress_callback(idx + 1, num_samples, rd)128 if (idx + 1) % 10 == 0:129 logger.info(f"Processed {idx + 1}/{num_samples} queries...")130 131 except Exception as e:132 logger.error(f"Error on query {idx}: {e}")133 134 # Run BERTScore on full batch (more efficient than per-sample)135 bertscore_results = {}136 if run_bertscore and self.eval_samples:137 logger.info("Computing BERTScore for all pipelines...")138 for pipe in ["llm_only", "baseline", "graphrag"]:139 try:140 bs = self.evaluator.evaluate_bertscore_batch(self.eval_samples, pipeline=pipe)141 bertscore_results[pipe] = bs142 logger.info(f" {pipe}: mean_f1={bs.get('mean_f1', 0):.4f}, pass_rate={bs.get('pass_rate', 0):.1%}")143 except Exception as e:144 logger.warning(f" BERTScore for {pipe} failed: {e}")145 146 aggregate = self.evaluator.compute_aggregate_metrics()147 report = self.evaluator.generate_report()148 return {149 "results": results,150 "aggregate": aggregate,151 "bertscore": bertscore_results,152 "report": report,153 "num_completed": len(results),154 "num_requested": num_samples,155 }156 157 def get_results_dataframe(self):158 import pandas as pd159 return pd.DataFrame(self.benchmark_results) if self.benchmark_results else pd.DataFrame()160 161 def save_results(self, filepath):162 with open(filepath, 'w') as f:163 json.dump({164 "results": self.benchmark_results,165 "aggregate": self.evaluator.compute_aggregate_metrics(),166 }, f, indent=2, default=str)167