Spaces:
Running
Running
| import time | |
| from benchmark_utils import ( | |
| BENCHMARK_RESULTS_PATH, | |
| EVAL_QUESTIONS_PATH, | |
| PIPELINES, | |
| estimate_cost, | |
| read_json, | |
| write_json, | |
| ) | |
| def main() -> None: | |
| questions = read_json(EVAL_QUESTIONS_PATH, []) | |
| if not questions: | |
| raise FileNotFoundError(f"No evaluation questions found in {EVAL_QUESTIONS_PATH}") | |
| from pipelines.basic_rag.rag_pipeline import run_basic_rag | |
| from pipelines.graphrag.graphrag_pipeline import run_graphrag | |
| from pipelines.llm_only.llm_pipeline import run_llm_only | |
| runners = { | |
| "llm_only": run_llm_only, | |
| "basic_rag": run_basic_rag, | |
| "graphrag": run_graphrag, | |
| } | |
| rows = [] | |
| for index, item in enumerate(questions, start=1): | |
| question = item["question"] | |
| print(f"[{index}/{len(questions)}] {question}") | |
| pipeline_results = {} | |
| for name in PIPELINES: | |
| pipeline_results[name] = normalize_result(runners[name], question) | |
| rows.append( | |
| { | |
| "question": question, | |
| "correct_answer": item.get("correct_answer", ""), | |
| "category": item.get("category", ""), | |
| "difficulty": item.get("difficulty", ""), | |
| "source_doc_ids": item.get("source_doc_ids", []), | |
| "pipelines": pipeline_results, | |
| } | |
| ) | |
| write_json(BENCHMARK_RESULTS_PATH, rows) | |
| print(f"Saved benchmark results to {BENCHMARK_RESULTS_PATH}") | |
| def normalize_result(runner, question): | |
| started = time.time() | |
| try: | |
| raw = runner(question) | |
| status = raw.get("status", "success") | |
| except Exception as exc: | |
| elapsed = time.time() - started | |
| return { | |
| "status": "error", | |
| "answer": "", | |
| "total_tokens": 0, | |
| "input_tokens": 0, | |
| "output_tokens": 0, | |
| "latency_seconds": elapsed, | |
| "estimated_cost": 0, | |
| "retrieved_context_count": 0, | |
| "graph_nodes_used": 0, | |
| "graph_edges_used": 0, | |
| "error": str(exc), | |
| } | |
| details = raw.get("details") or {} | |
| total_tokens = raw.get("total_tokens", raw.get("tokens", 0)) or 0 | |
| input_tokens = details.get("prompt_tokens", raw.get("prompt_tokens", 0)) or 0 | |
| output_tokens = details.get("completion_tokens", raw.get("completion_tokens", 0)) or 0 | |
| context_count = len(details.get("retrieved_chunks") or details.get("chunks") or []) | |
| reasoning_paths = details.get("reasoning_paths") or [] | |
| matched_entities = details.get("matched_entities") or [] | |
| graph_chunks = details.get("chunks") or [] | |
| retrieval_trace = details.get("retrieval_trace") or raw.get("retrieval_trace") or {} | |
| return { | |
| "status": status, | |
| "answer": raw.get("answer", ""), | |
| "total_tokens": total_tokens, | |
| "input_tokens": input_tokens, | |
| "output_tokens": output_tokens, | |
| "latency_seconds": raw.get("latency", time.time() - started) or 0, | |
| "estimated_cost": estimate_cost(total_tokens), | |
| "retrieved_context_count": context_count, | |
| "graph_nodes_used": retrieval_trace.get("expanded_node_count", len(matched_entities) + len(graph_chunks)), | |
| "graph_edges_used": retrieval_trace.get( | |
| "expanded_edge_count", | |
| sum(max(0, len(path) - 1) for path in reasoning_paths), | |
| ), | |
| "seed_chunks_used": retrieval_trace.get("seed_count", details.get("seed_chunks_used", 0)), | |
| "fallback_used": retrieval_trace.get("fallback_used", details.get("fallback_used", False)), | |
| "reranked_candidate_count": retrieval_trace.get( | |
| "reranked_candidate_count", | |
| details.get("reranked_candidate_count", 0), | |
| ), | |
| } | |
| if __name__ == "__main__": | |
| main() | |