Fix bug in cache_manager

RUC-NLPIR · Jul 28, 2024 · 8277011 · 8277011
1 parent cbd4c8d
commit 8277011
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/flashrag/retriever/retriever.py b/flashrag/retriever/retriever.py
@@ -62,12 +62,9 @@ def wrapper(self, query_list, num = None, return_score = False):
             if isinstance(query_list, str):
                 query_list = [query_list]
                 if 'batch' not in func.__name__:
-                    new_results = [results]
-                    new_scores = [scores]
-                else:
-                    new_results = results
-                    new_scores = scores
-            for query, doc_items, doc_scores in zip(query_list, new_results, new_scores):
+                    results = [results]
+                    scores = [scores]
+            for query, doc_items, doc_scores in zip(query_list, results, scores):
                 for item, score in zip(doc_items, doc_scores):
                     item['score'] = score
                 self.cache[query] = doc_items