flashrag-dev 0.1.4.dev20250322__tar.gz → 0.1.4.dev20250403__tar.gz

{flashrag_dev-0.1.4.dev20250322 → flashrag_dev-0.1.4.dev20250403}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag_dev
-Version: 0.1.4.dev20250322
+Version: 0.1.4.dev20250403
 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
@@ -108,8 +108,6 @@ FlashRAG is still under development and there are many issues and room for impro
 
 [25/01/07] We have currently support the aggregation of multiple retrievers, see it in [**multi retriever usage**](https://github.com/RUC-NLPIR/FlashRAG/blob/main/docs/original_docs/multi_retriever_usage.md).
 
-[25/01/07] We have integrated a very flexible and lightweight corpus chunking library [**Chunkie**](https://github.com/chonkie-ai/chonkie?tab=readme-ov-file#usage), which supports various custom chunking methods (tokens, sentences, semantic, etc.). Use it in [<u>chunking doc corpus</u>](docs/original_docs/chunk-doc-corpus.md).
-
 [24/10/21] We have released a version based on the Paddle framework that supports Chinese hardware platforms. Please refer to [FlashRAG Paddle](https://github.com/RUC-NLPIR/FlashRAG-Paddle) for details.
 
 [24/10/13] A new in-domain dataset and corpus - [DomainRAG](https://arxiv.org/pdf/2406.05654) have been added to the dataset. The dataset is based on the internal enrollment data of Renmin University of China, covering seven types of tasks, which can be used for conducting domain-specific RAG testing.

{flashrag_dev-0.1.4.dev20250322 → flashrag_dev-0.1.4.dev20250403}/README.md

@@ -91,8 +91,6 @@ FlashRAG is still under development and there are many issues and room for impro
 
 [25/01/07] We have currently support the aggregation of multiple retrievers, see it in [**multi retriever usage**](https://github.com/RUC-NLPIR/FlashRAG/blob/main/docs/original_docs/multi_retriever_usage.md).
 
-[25/01/07] We have integrated a very flexible and lightweight corpus chunking library [**Chunkie**](https://github.com/chonkie-ai/chonkie?tab=readme-ov-file#usage), which supports various custom chunking methods (tokens, sentences, semantic, etc.). Use it in [<u>chunking doc corpus</u>](docs/original_docs/chunk-doc-corpus.md).
-
 [24/10/21] We have released a version based on the Paddle framework that supports Chinese hardware platforms. Please refer to [FlashRAG Paddle](https://github.com/RUC-NLPIR/FlashRAG-Paddle) for details.
 
 [24/10/13] A new in-domain dataset and corpus - [DomainRAG](https://arxiv.org/pdf/2406.05654) have been added to the dataset. The dataset is based on the internal enrollment data of Renmin University of China, covering seven types of tasks, which can be used for conducting domain-specific RAG testing.

{flashrag_dev-0.1.4.dev20250322 → flashrag_dev-0.1.4.dev20250403}/flashrag/retriever/retriever.py

@@ -156,6 +156,7 @@ class BaseRetriever:
             with open(self.cache_path, "r") as f:
                 self.cache = json.load(f)
         self.silent = self._config["silent_retrieval"] if "silent_retrieval" in self._config else False
+
     def update_additional_setting(self):
         pass
 
@@ -370,31 +371,54 @@ class DenseRetriever(BaseTextRetriever):
         self.batch_size = self._config["retrieval_batch_size"]
         self.instruction = self._config["instruction"]
 
-        self.retreival_model_path = self._config["retrieval_model_path"]
+        self.retrieval_model_path = self._config["retrieval_model_path"]
         self.use_st = self._config["use_sentence_transformer"]
         self.use_faiss_gpu = self._config["faiss_gpu"]
 
     def load_model(self):
         if self.use_st:
             self.encoder = STEncoder(
-                model_name = self.retrieval_method,
-                model_path = self._config["retrieval_model_path"],
-                max_length = self.query_max_length,
-                use_fp16 = self.use_fp16,
-                instruction = self.instruction,
-                silent = self.silent
+                model_name=self.retrieval_method,
+                model_path=self._config["retrieval_model_path"],
+                max_length=self.query_max_length,
+                use_fp16=self.use_fp16,
+                instruction=self.instruction,
+                silent=self.silent,
             )
         else:
             # check pooling method
             self._check_pooling_method(self.retrieval_model_path, self.pooling_method)
             self.encoder = Encoder(
-                model_name = self.retrieval_method,
-                model_path = self.retreival_model_path,
-                pooling_method = self.pooling_method,
-                max_length = self.query_max_length,
-                use_fp16 = self.use_fp16,
-                instruction = self.instruction,
-                silent = self.silent
+                model_name=self.retrieval_method,
+                model_path=self.retrieval_model_path,
+                pooling_method=self.pooling_method,
+                max_length=self.query_max_length,
+                use_fp16=self.use_fp16,
+                instruction=self.instruction,
+            )
+
+    def _check_pooling_method(self, model_path, pooling_method):
+        try:
+            # read pooling method from 1_Pooling/config.json
+            pooling_config = json.load(open(os.path.join(model_path, "1_Pooling/config.json")))
+            for k, v in pooling_config.items():
+                if k.startswith("pooling_mode") and v == True:
+                    detect_pooling_method = k.split("pooling_mode_")[-1]
+                    if detect_pooling_method == "mean_tokens":
+                        detect_pooling_method = "mean"
+                    elif detect_pooling_method == "cls_token":
+                        detect_pooling_method = "cls"
+                    else:
+                        # raise warning: not implemented pooling method
+                        warnings.warn(f"Pooling method {detect_pooling_method} is not implemented.", UserWarning)
+                        detect_pooling_method = "mean"
+                    break
+        except:
+            detect_pooling_method = None
+
+        if detect_pooling_method is not None and detect_pooling_method != pooling_method:
+            warnings.warn(
+                f"Pooling method in model config file is {detect_pooling_method}, but the input is {pooling_method}. Please check carefully."
             )
 
     def _search(self, query: str, num: int = None, return_score=False):
@@ -462,9 +486,7 @@ class MultiModalRetriever(BaseRetriever):
         self.batch_size = config["retrieval_batch_size"]
 
         self.encoder = ClipEncoder(
-            model_name=self.retrieval_method,
-            model_path=config["retrieval_model_path"],
-            silent=self.silent
+            model_name=self.retrieval_method, model_path=config["retrieval_model_path"], silent=self.silent
         )
 
     def _judge_input_modal(self, query):

flashrag_dev-0.1.4.dev20250403/flashrag/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.4dev20250403"

{flashrag_dev-0.1.4.dev20250322 → flashrag_dev-0.1.4.dev20250403}/flashrag_dev.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag-dev
-Version: 0.1.4.dev20250322
+Version: 0.1.4.dev20250403
 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
@@ -108,8 +108,6 @@ FlashRAG is still under development and there are many issues and room for impro
 
 [25/01/07] We have currently support the aggregation of multiple retrievers, see it in [**multi retriever usage**](https://github.com/RUC-NLPIR/FlashRAG/blob/main/docs/original_docs/multi_retriever_usage.md).
 
-[25/01/07] We have integrated a very flexible and lightweight corpus chunking library [**Chunkie**](https://github.com/chonkie-ai/chonkie?tab=readme-ov-file#usage), which supports various custom chunking methods (tokens, sentences, semantic, etc.). Use it in [<u>chunking doc corpus</u>](docs/original_docs/chunk-doc-corpus.md).
-
 [24/10/21] We have released a version based on the Paddle framework that supports Chinese hardware platforms. Please refer to [FlashRAG Paddle](https://github.com/RUC-NLPIR/FlashRAG-Paddle) for details.
 
 [24/10/13] A new in-domain dataset and corpus - [DomainRAG](https://arxiv.org/pdf/2406.05654) have been added to the dataset. The dataset is based on the internal enrollment data of Renmin University of China, covering seven types of tasks, which can be used for conducting domain-specific RAG testing.

flashrag_dev-0.1.4.dev20250322/flashrag/version.py

@@ -1 +0,0 @@
-__version__ = "0.1.4dev20250322"