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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flashrag_dev
-Version: 0.1.4.dev20250322
+Version: 0.1.4.dev20250326
 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.1.4.dev20250322 → flashrag_dev-0.1.4.dev20250326}/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.dev20250326/flashrag/version.py

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

{flashrag_dev-0.1.4.dev20250322 → flashrag_dev-0.1.4.dev20250326}/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.dev20250326
 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.1.4.dev20250322/flashrag/version.py

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