Merge branch 'new-flagembedding-v1' of https://github.com/hanhainebula/FlagEmbedding into new-flagembedding-v1

Author: ZiyiXia

FlagEmbedding/abc/evaluation/data_loader.py

@@ -192,7 +192,7 @@ def _download_file(self, download_url: str, save_dir: str):
         try:
             subprocess.run(cmd, check=True)
         except subprocess.CalledProcessError as e:
-            logger.error(f"Error code: {e.returncode}. Error message: {e.stderr}")
+            logger.warning(e.output)
 
         if not os.path.exists(save_path) or os.path.getsize(save_path) == 0:
             raise FileNotFoundError(f"Failed to download file from {download_url} to {save_path}")
@@ -217,7 +217,7 @@ def _download_gz_file(self, download_url: str, save_dir: str):
         try:
             subprocess.run(cmd, check=True)
         except subprocess.CalledProcessError as e:
-            logger.error(f"Error code: {e.returncode}. Error message: {e.output}")
+            logger.warning(e.output)
 
         file_path = gz_file_path.replace(".gz", "")
         if not os.path.exists(file_path) or self._get_fpath_size(file_path) == 0:
@@ -236,7 +236,7 @@ def _download_zip_file(self, download_url: str, save_dir: str):
         try:
             subprocess.run(cmd, check=True)
         except subprocess.CalledProcessError as e:
-            logger.error(f"Error code: {e.returncode}. Error message: {e.output}")
+            logger.warning(e.output)
 
         if not os.path.exists(file_path) or self._get_fpath_size(file_path) == 0:
             raise FileNotFoundError(f"Failed to unzip file {zip_file_path}")

FlagEmbedding/abc/evaluation/evaluator.py

@@ -173,6 +173,7 @@ def __call__(
         retriever_eval_results = self.evaluate_results(no_reranker_search_results_save_dir, k_values=k_values)
         self.output_eval_results_to_json(retriever_eval_results, eval_results_save_path)
 
+        retriever.stop_multi_process_pool()
         # Reranking Stage
         if reranker is not None:
             reranker_search_results_save_dir = os.path.join(
@@ -215,6 +216,7 @@ def __call__(
             eval_results_save_path = os.path.join(reranker_search_results_save_dir, 'EVAL', 'eval_results.json')
             reranker_eval_results = self.evaluate_results(reranker_search_results_save_dir, k_values=k_values)
             self.output_eval_results_to_json(reranker_eval_results, eval_results_save_path)
+            reranker.stop_multi_process_pool()
 
     @staticmethod
     def save_search_results(

FlagEmbedding/abc/evaluation/searcher.py

@@ -25,6 +25,10 @@ def __str__(self) -> str:
         """
         return os.path.basename(self.embedder.model.config._name_or_path)
 
+    def stop_multi_process_pool(self):
+        if self.embedder.pool is not None:
+            self.embedder.stop_multi_process_pool(self.embedder.pool)
+
     @abstractmethod
     def __call__(
         self,
@@ -144,6 +148,10 @@ def __str__(self) -> str:
         """
         return os.path.basename(self.reranker.model.config._name_or_path)
 
+    def stop_multi_process_pool(self):
+        if self.reranker.pool is not None:
+            self.reranker.stop_multi_process_pool(self.reranker.pool)
+
     def __call__(
         self,
         corpus: Dict[str, Dict[str, Any]],

FlagEmbedding/evaluation/air_bench/arguments.py

@@ -9,9 +9,15 @@ class AIRBenchEvalModelArgs:
     embedder_name_or_path: str = field(
         metadata={"help": "The embedder name or path.", "required": True}
     )
+    embedder_model_class: Optional[str] = field(
+        default=None, metadata={"help": "The embedder model class. Available classes: ['encoder-only-base', 'encoder-only-m3', 'decoder-only-base', 'decoder-only-icl']. Default: None. For the custom model, you need to specifiy the model class.", "choices": ["encoder-only-base", "encoder-only-m3", "decoder-only-base", "decoder-only-icl"]}
+    )
     normalize_embeddings: bool = field(
         default=True, metadata={"help": "whether to normalize the embeddings"}
     )
+    pooling_method: str = field(
+        default="cls", metadata={"help": "The pooling method fot the embedder."}
+    )
     use_fp16: bool = field(
         default=True, metadata={"help": "whether to use fp16 for inference"}
     )
@@ -36,6 +42,9 @@ class AIRBenchEvalModelArgs:
     reranker_name_or_path: Optional[str] = field(
         default=None, metadata={"help": "The reranker name or path."}
     )
+    reranker_model_class: Optional[str] = field(
+        default=None, metadata={"help": "The reranker model class. Available classes: ['encoder-only-base', 'decoder-only-base', 'decoder-only-layerwise', 'decoder-only-lightweight']. Default: None. For the custom model, you need to specify the model class.", "choices": ["encoder-only-base", "decoder-only-base", "decoder-only-layerwise", "decoder-only-lightweight"]}
+    )
     reranker_peft_path: Optional[str] = field(
         default=None, metadata={"help": "The reranker peft path."}
     )

FlagEmbedding/evaluation/beir/evaluator.py

@@ -145,6 +145,8 @@ def __call__(
             retriever_eval_results = self.evaluate_results(no_reranker_search_results_save_dir, k_values=k_values)
             self.output_eval_results_to_json(retriever_eval_results, eval_results_save_path)
 
+            retriever.stop_multi_process_pool()
+
             # Reranking Stage
             if reranker is not None:
                 reranker_search_results_save_dir = os.path.join(
@@ -314,7 +316,8 @@ def __call__(
                     eval_results_save_path = os.path.join(reranker_search_results_save_dir, 'EVAL', 'eval_results.json')
                     reranker_eval_results = self.evaluate_results(reranker_search_results_save_dir, k_values=k_values)
                     self.output_eval_results_to_json(reranker_eval_results, eval_results_save_path)
-
+            if reranker is not None:
+                reranker.stop_multi_process_pool()
     def evaluate_results(
         self,
         search_results_save_dir: str,

FlagEmbedding/evaluation/msmarco/data_loader.py

@@ -45,19 +45,33 @@ def _load_remote_corpus(
             corpus_dict = {}
             with open(save_path, "w", encoding="utf-8") as f:
                 for data in tqdm(corpus, desc="Loading and Saving corpus"):
-                    _data = {
-                        "id": data["docid"],
-                        "title": data["title"],
-                        "text": data.get("text", data.get("body", ""))
-                    }
-                    corpus_dict[data["docid"]] = {
-                        "title": data["title"],
-                        "text": data.get("text", data.get("body", ""))
-                    }
+                    if dataset_name == 'passage':
+                        _data = {
+                            "id": data["docid"],
+                            "title": data["title"],
+                            "text": data["text"]
+                        }
+                        corpus_dict[data["docid"]] = {
+                            "title": data["title"],
+                            "text": data["text"]
+                        }
+                    else:
+                        _data = {
+                            "id": data["doc_id"],
+                            "title": data["title"],
+                            "text": data["body"]
+                        }
+                        corpus_dict[data["doc_id"]] = {
+                            "title": data["title"],
+                            "text": data["body"]
+                        }
                     f.write(json.dumps(_data, ensure_ascii=False) + "\n")
             logging.info(f"{self.eval_name} {dataset_name} corpus saved to {save_path}")
         else:
-            corpus_dict = {data["docid"]: {"title": data["title"], "text": data.get("text", data.get("body", ""))} for data in tqdm(corpus, desc="Loading corpus")}
+            if dataset_name == 'passage':
+                corpus_dict = {data["docid"]: {"title": data["title"], "text": data["text"]} for data in tqdm(corpus, desc="Loading corpus")}
+            else:
+                corpus_dict = {data["doc_id"]: {"title": data["title"], "text": data["body"]} for data in tqdm(corpus, desc="Loading corpus")}
         return datasets.DatasetDict(corpus_dict)
 
     def _load_remote_qrels(

FlagEmbedding/finetune/embedder/encoder_only/m3/runner.py

@@ -38,8 +38,8 @@ def get_model(
         colbert_dim: int = -1,
         cache_dir: str = None
     ):
+        cache_folder = os.getenv('HF_HUB_CACHE', None) if cache_dir is None else cache_dir
         if not os.path.exists(model_name_or_path):
-            cache_folder = os.getenv('HF_HUB_CACHE', None) if cache_dir is None else cache_dir
             model_name_or_path = snapshot_download(
                 repo_id=model_name_or_path,
                 cache_dir=cache_folder,
@@ -48,6 +48,7 @@ def get_model(
 
         model = AutoModel.from_pretrained(
             model_name_or_path,
+            cache_dir=cache_folder,
             trust_remote_code=trust_remote_code
         )
         colbert_linear = torch.nn.Linear(