From 7a4ad471c9026c7882504b1c8b730045b4bb74af Mon Sep 17 00:00:00 2001 From: CHEN SHENGYUAN Date: Thu, 18 Dec 2025 15:35:33 +0800 Subject: enable vectorized retrieval with sparse matrix operations --- run.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) (limited to 'run.py') diff --git a/run.py b/run.py index 0d5d017..4ea5b3b 100644 --- a/run.py +++ b/run.py @@ -25,6 +25,7 @@ def parse_arguments(): parser.add_argument("--iteration_threshold", type=float, default=0.4, help="The threshold for iteration") parser.add_argument("--passage_ratio", type=float, default=2, help="The ratio for passage") parser.add_argument("--top_k_sentence", type=int, default=3, help="The top k sentence to use") + parser.add_argument("--use_vectorized_retrieval", action="store_true", help="Use vectorized matrix-based retrieval instead of BFS iteration") return parser.parse_args() @@ -59,7 +60,8 @@ def main(): max_iterations=args.max_iterations, iteration_threshold=args.iteration_threshold, passage_ratio=args.passage_ratio, - top_k_sentence=args.top_k_sentence + top_k_sentence=args.top_k_sentence, + use_vectorized_retrieval=args.use_vectorized_retrieval ) rag_model = LinearRAG(global_config=config) rag_model.index(passages) -- cgit v1.2.3