flashrag-dev 0.3.0.dev20260821__tar.gz → 0.3.0.dev20260919__tar.gz

{flashrag_dev-0.3.0.dev20260821/flashrag_dev.egg-info → flashrag_dev-0.3.0.dev20260919}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashrag_dev
-Version: 0.3.0.dev20260821
+Version: 0.3.0.dev20260919
 Summary: A library for efficient Retrieval-Augmented Generation research
 Home-page: https://github.com/RUC-NLPIR/FlashRAG
 Author: Jiajie Jin, Yutao Zhu, Chenghao Zhang, Xinyu Yang, Zhicheng Dou

{flashrag_dev-0.3.0.dev20260821 → flashrag_dev-0.3.0.dev20260919}/flashrag/generator/generator.py

@@ -269,6 +269,7 @@ class VLLMGenerator(BaseGenerator):
         if return_scores:
             scores = []
             for output in outputs:
+                output_scores = []
                 for single_output in output.outputs:
                     if single_output.logprobs:
                         token_probs = [np.exp(list(score_dict.values())[0].logprob) 

{flashrag_dev-0.3.0.dev20260821 → flashrag_dev-0.3.0.dev20260919}/flashrag/refiner/refiner.py

@@ -168,7 +168,7 @@ class ExtractiveRefiner(BaseRefiner):
 
         # split into sentences: [[sent1, sent2,...], [...]]
         sent_lists = [
-            [i.strip() for i in re.split(r"(?<![A-Za-z]\.)(?<=[.!?])\s+", " ".join(res)) if len(i.strip()) > 5]
+            [i.strip() for i in re.split(r"(?<=[.!?])\s+(?![a-z])", " ".join(res)) if len(i.strip()) > 5]
             for res in retrieval_results
         ]
         score_lists = []  # matching scores, size == sent_lists

{flashrag_dev-0.3.0.dev20260821 → flashrag_dev-0.3.0.dev20260919}/flashrag/retriever/index_builder.py

@@ -150,7 +150,6 @@ class Index_Builder:
             else:
                 assert False, "Invalid bm25 backend!"
         elif self.retrieval_method == "splade":
-            from seismic import SeismicIndex
             self.build_seismic_index()
         else:
             self.build_dense_index()
@@ -176,6 +175,7 @@ class Index_Builder:
         --batched_indexing 10000 # seismic batch
         --nknn 32
         """
+        from seismic import SeismicIndex
 
         if self.pooling_method != 'max':
             print(

{flashrag_dev-0.3.0.dev20260821 → flashrag_dev-0.3.0.dev20260919}/flashrag/retriever/retriever.py

@@ -68,6 +68,8 @@ def cache_manager(func):
                 [t[0] for t in cache_results],
                 [t[1] for t in cache_results],
             )
+            if "batch" not in func.__name__:
+                results, scores = results[0], scores[0]
 
         else:
             results, scores = func(self, query=query, num=num, return_score=True)

flashrag_dev-0.3.0.dev20260919/flashrag/version.py

@@ -0,0 +1 @@
+__version__ = "0.3.0dev20260919"

{flashrag_dev-0.3.0.dev20260821 → flashrag_dev-0.3.0.dev20260919/flashrag_dev.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashrag_dev
-Version: 0.3.0.dev20260821
+Version: 0.3.0.dev20260919
 Summary: A library for efficient Retrieval-Augmented Generation research
 Home-page: https://github.com/RUC-NLPIR/FlashRAG
 Author: Jiajie Jin, Yutao Zhu, Chenghao Zhang, Xinyu Yang, Zhicheng Dou

{flashrag_dev-0.3.0.dev20260821 → flashrag_dev-0.3.0.dev20260919}/flashrag_dev.egg-info/SOURCES.txt

@@ -58,4 +58,5 @@ flashrag_dev.egg-info/PKG-INFO
 flashrag_dev.egg-info/SOURCES.txt
 flashrag_dev.egg-info/dependency_links.txt
 flashrag_dev.egg-info/requires.txt
-flashrag_dev.egg-info/top_level.txt
+flashrag_dev.egg-info/top_level.txt
+tests/test_retriever_cache.py

flashrag_dev-0.3.0.dev20260919/tests/test_retriever_cache.py

@@ -0,0 +1,86 @@
+import copy
+import json
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+from flashrag.retriever import BaseTextRetriever
+
+
+class TestRetrieverCache(unittest.TestCase):
+    def setUp(self):
+        self.directory = tempfile.TemporaryDirectory()
+        self.addCleanup(self.directory.cleanup)
+        self.cache_path = Path(self.directory.name) / "retrieval_cache.json"
+        self.docs = [
+            {"id": "1", "contents": "First document", "score": 0.8},
+            {"id": "2", "contents": "Second document", "score": 0.5},
+        ]
+        self.cache_path.write_text(json.dumps({"cached": self.docs, "empty": []}))
+
+    def retriever(self, save_cache=False):
+        return BaseTextRetriever(
+            {
+                "retrieval_method": "bm25",
+                "retrieval_topk": 2,
+                "index_path": None,
+                "corpus_path": None,
+                "save_retrieval_cache": save_cache,
+                "use_retrieval_cache": True,
+                "retrieval_cache_path": str(self.cache_path),
+                "use_reranker": False,
+                "save_dir": self.directory.name,
+            }
+        )
+
+    def test_single_search_returns_flat_results_and_scores(self):
+        for query in ["cached", "missing"]:
+            for num in [1, 2]:
+                for return_score in [False, True]:
+                    with self.subTest(query=query, num=num, return_score=return_score):
+                        retriever = self.retriever()
+                        docs = copy.deepcopy(self.docs[:num])
+                        scores = [doc["score"] for doc in docs]
+                        with patch.object(retriever, "_batch_search", return_value=([docs], [scores])) as search:
+                            result = retriever.search(query, num=num, return_score=return_score)
+                        self.assertEqual(result, (docs, scores) if return_score else docs)
+                        if query == "cached":
+                            search.assert_not_called()
+                        else:
+                            search.assert_called_once_with(query=[query], num=num, return_score=True)
+
+    def test_single_search_can_save_cache(self):
+        for query in ["cached", "missing"]:
+            with self.subTest(query=query):
+                retriever = self.retriever(save_cache=True)
+                docs = copy.deepcopy(self.docs)
+                scores = [doc["score"] for doc in docs]
+                with patch.object(retriever, "_batch_search", return_value=([docs], [scores])):
+                    self.assertEqual(retriever.search(query), docs)
+                retriever._save_cache()
+                saved = json.loads(self.cache_path.read_text())
+                self.assertEqual(saved[query], docs)
+
+    def test_single_search_with_no_cached_results(self):
+        for return_score in [False, True]:
+            with self.subTest(return_score=return_score):
+                with self.assertWarns(UserWarning):
+                    result = self.retriever().search("empty", return_score=return_score)
+                self.assertEqual(result, ([], []) if return_score else [])
+
+    def test_batch_search_keeps_query_dimension_and_order(self):
+        for return_score in [False, True]:
+            with self.subTest(return_score=return_score):
+                retriever = self.retriever()
+                missing_docs = [{"id": "3", "contents": "Another document"}]
+                with patch.object(retriever, "_batch_search", return_value=([missing_docs], [[0.7]])) as search:
+                    result = retriever.batch_search(["missing", "cached"], return_score=return_score)
+                expected_docs = [missing_docs, self.docs]
+                expected_scores = [[0.7], [0.8, 0.5]]
+                self.assertEqual(result, (expected_docs, expected_scores) if return_score else expected_docs)
+                search.assert_called_once_with(query=["missing"], num=2, return_score=True)
+
+
+if __name__ == "__main__":
+    unittest.main()

flashrag_dev-0.3.0.dev20260821/flashrag/version.py

@@ -1 +0,0 @@
-__version__ = "0.3.0dev20260821"