update eval

Author: 545999961

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/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/__main__.py

@@ -5,22 +5,19 @@
     MSMARCOEvalRunner
 )
 
-def main():
-    parser = HfArgumentParser((
-        MSMARCOEvalArgs,
-        MSMARCOEvalModelArgs
-    ))
 
-    eval_args, model_args = parser.parse_args_into_dataclasses()
-    eval_args: MSMARCOEvalArgs
-    model_args: MSMARCOEvalModelArgs
+parser = HfArgumentParser((
+    MSMARCOEvalArgs,
+    MSMARCOEvalModelArgs
+))
 
-    runner = MSMARCOEvalRunner(
-        eval_args=eval_args,
-        model_args=model_args
-    )
+eval_args, model_args = parser.parse_args_into_dataclasses()
+eval_args: MSMARCOEvalArgs
+model_args: MSMARCOEvalModelArgs
 
-    runner.run()
+runner = MSMARCOEvalRunner(
+    eval_args=eval_args,
+    model_args=model_args
+)
 
-if __name__ == "__main__":
-    main()
+runner.run()