From d14faaa77141981e62700724ca74d152bd7ad63c Mon Sep 17 00:00:00 2001 From: yichuan-w Date: Mon, 18 Aug 2025 17:01:43 -0700 Subject: [PATCH] feat(benchmarks): add --batch-size option and plumb through to HNSW search (default 0) --- benchmarks/run_evaluation.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/benchmarks/run_evaluation.py b/benchmarks/run_evaluation.py index a018d98..ab4e169 100644 --- a/benchmarks/run_evaluation.py +++ b/benchmarks/run_evaluation.py @@ -197,6 +197,12 @@ def main(): parser.add_argument( "--ef-search", type=int, default=120, help="The 'efSearch' parameter for HNSW." ) + parser.add_argument( + "--batch-size", + type=int, + default=0, + help="Batch size for HNSW batched search (0 disables batching)", + ) parser.add_argument( "--llm-type", type=str, @@ -331,13 +337,23 @@ def main(): for i in range(num_eval_queries): start_time = time.time() - new_results = searcher.search(queries[i], top_k=args.top_k, complexity=args.ef_search) + new_results = searcher.search( + queries[i], + top_k=args.top_k, + complexity=args.ef_search, + batch_size=args.batch_size, + ) search_times.append(time.time() - start_time) # Optional: also call the LLM with configurable backend/model (does not affect recall) llm_config = {"type": args.llm_type, "model": args.llm_model} chat = LeannChat(args.index_path, llm_config=llm_config, searcher=searcher) - answer = chat.ask(queries[i], top_k=args.top_k, complexity=args.ef_search) + answer = chat.ask( + queries[i], + top_k=args.top_k, + complexity=args.ef_search, + batch_size=args.batch_size, + ) print(f"Answer: {answer}") # Correct Recall Calculation: Based on TEXT content new_texts = {result.text for result in new_results}