Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """ | |
| Benchmark: run pipeline on N copies of a PDF and collect timing + score metrics. | |
| Usage: | |
| python scripts/benchmark_agents.py path/to/paper.pdf --runs 5 | |
| """ | |
| import argparse | |
| import asyncio | |
| import json | |
| import sys | |
| import time | |
| from pathlib import Path | |
| sys.path.insert(0, str(Path(__file__).parent.parent)) | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| from orchestrators.custom_orchestrator import CitationEdgeOrchestrator | |
| from utils.helpers import generate_id | |
| async def run_once(orch, pdf_path, run_num): | |
| job_id = generate_id(f"bench{run_num}") | |
| t0 = time.monotonic() | |
| try: | |
| result = await orch.run(job_id=job_id, pdf_path=str(pdf_path)) | |
| duration = time.monotonic() - t0 | |
| scores = result.get("scores", {}) | |
| return { | |
| "run": run_num, | |
| "job_id": job_id, | |
| "status": result.get("status"), | |
| "duration_s": round(duration, 2), | |
| "overall_score": scores.get("overall_score", 0), | |
| "literary_score": scores.get("literary_score", 0), | |
| "argument_score": scores.get("argument_score", 0), | |
| "agents": {k: v.get("duration_s") for k, v in result.get("agents", {}).items()}, | |
| } | |
| except Exception as e: | |
| return {"run": run_num, "status": "error", "error": str(e)} | |
| async def main(): | |
| parser = argparse.ArgumentParser(description="Benchmark CitationEdge pipeline") | |
| parser.add_argument("pdf_path", help="Path to the research paper PDF") | |
| parser.add_argument("--runs", type=int, default=3, help="Number of runs (default: 3)") | |
| args = parser.parse_args() | |
| pdf_path = Path(args.pdf_path) | |
| if not pdf_path.exists(): | |
| print(f"Error: PDF not found: {pdf_path}") | |
| sys.exit(1) | |
| print(f"Benchmarking {args.runs} runs on: {pdf_path}\n") | |
| orch = CitationEdgeOrchestrator() | |
| results = [] | |
| for i in range(1, args.runs + 1): | |
| print(f"Run {i}/{args.runs}...", end=" ", flush=True) | |
| r = await run_once(orch, pdf_path, i) | |
| print(f"{r.get('duration_s', '?')}s | score={r.get('overall_score', '?')}") | |
| results.append(r) | |
| # Summary | |
| durations = [r["duration_s"] for r in results if "duration_s" in r] | |
| scores = [r["overall_score"] for r in results if "overall_score" in r] | |
| summary = { | |
| "runs": args.runs, | |
| "avg_duration_s": round(sum(durations) / len(durations), 2) if durations else 0, | |
| "avg_score": round(sum(scores) / len(scores), 2) if scores else 0, | |
| "results": results, | |
| } | |
| output_path = Path("data/benchmark_results.json") | |
| output_path.parent.mkdir(exist_ok=True) | |
| with open(output_path, "w") as f: | |
| json.dump(summary, f, indent=2) | |
| print(f"\nBenchmark complete. Avg duration: {summary['avg_duration_s']}s | Avg score: {summary['avg_score']}/10") | |
| print(f"Results saved to {output_path}") | |
| if __name__ == "__main__": | |
| asyncio.run(main()) | |