Spaces:
Building
Building
File size: 2,928 Bytes
0e38162 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | #!/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())
|