Merge pull request #1149 from ZiyiXia/master

Update C-MTEB

Author: ZiyiXia

.gitignore

@@ -136,3 +136,8 @@ pic2.py
 
 # Pyre type checker
 .pyre/
+
+# results
+results/
+en_results/
+zh_results/

C_MTEB/C_MTEB/__init__.py

@@ -1,10 +1,13 @@
+# from .tasks import *
 from .tasks import *
 
-ChineseTaskList = ['TNews', 'IFlyTek', 'MultilingualSentiment', 'JDReview', 'OnlineShopping', 'Waimai',
-                   'CLSClusteringS2S', 'CLSClusteringP2P', 'ThuNewsClusteringS2S', 'ThuNewsClusteringP2P',
-                   'Ocnli', 'Cmnli',
-                   'T2Reranking', 'MMarcoReranking', 'CMedQAv1', 'CMedQAv2',
-                   'T2Retrieval', 'MMarcoRetrieval', 'DuRetrieval', 'CovidRetrieval', 'CmedqaRetrieval', 'EcomRetrieval', 'MedicalRetrieval', 'VideoRetrieval',
-                   'ATEC', 'BQ', 'LCQMC', 'PAWSX', 'STSB', 'AFQMC', 'QBQTC']
+ChineseTaskList = [
+    'TNews', 'IFlyTek', 'MultilingualSentiment', 'JDReview', 'OnlineShopping', 'Waimai',
+    'CLSClusteringS2S.v2', 'CLSClusteringP2P.v2', 'ThuNewsClusteringS2S.v2', 'ThuNewsClusteringP2P.v2',
+    'Ocnli', 'Cmnli',
+    'T2Reranking', 'MMarcoReranking', 'CMedQAv1-reranking', 'CMedQAv2-reranking',
+    'T2Retrieval', 'MMarcoRetrieval', 'DuRetrieval', 'CovidRetrieval', 'CmedqaRetrieval', 'EcomRetrieval', 'MedicalRetrieval', 'VideoRetrieval',
+    'ATEC', 'BQ', 'LCQMC', 'PAWSX', 'STSB', 'AFQMC', 'QBQTC'
+]
 
 

C_MTEB/eval_C-MTEB.py

@@ -1,6 +1,6 @@
 import argparse
 
-from C_MTEB.tasks import *
+from C_MTEB import ChineseTaskList
 from flag_dres_model import FlagDRESModel
 from mteb import MTEB
 
@@ -24,24 +24,20 @@ def get_args():
     return parser.parse_args()
 
 
-
 if __name__ == '__main__':
     args = get_args()
 
     model = FlagDRESModel(model_name_or_path=args.model_name_or_path,
                           query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章：",
                           pooling_method=args.pooling_method)
+    
+    print(ChineseTaskList)
 
-    task_names = [t.description["name"] for t in MTEB(task_types=args.task_type,
-                                                      task_langs=['zh', 'zh-CN']).tasks]
-
-    for task in task_names:
-        # if task not in ChineseTaskList:
-        #     continue
+    for task in ChineseTaskList:
         if task in ['T2Retrieval', 'MMarcoRetrieval', 'DuRetrieval',
                     'CovidRetrieval', 'CmedqaRetrieval',
                     'EcomRetrieval', 'MedicalRetrieval', 'VideoRetrieval',
-                    'T2Reranking', 'MMarcoReranking', 'CMedQAv1', 'CMedQAv2']:
+                    'T2Reranking', 'MMarcoReranking', 'CMedQAv1-reranking', 'CMedQAv2-reranking']:
             if args.model_name_or_path not in query_instruction_for_retrieval_dict:
                 if args.add_instruction:
                     instruction = "为这个句子生成表示以用于检索相关文章："
@@ -55,8 +51,5 @@ def get_args():
 
         model.query_instruction_for_retrieval = instruction
 
-        evaluation = MTEB(tasks=[task], task_langs=['zh', 'zh-CN'])
+        evaluation = MTEB(tasks=[task])
         evaluation.run(model, output_folder=f"zh_results/{args.model_name_or_path.split('/')[-1]}")
-
-
-

C_MTEB/eval_MTEB.py

@@ -31,7 +31,10 @@ def get_args():
                           query_instruction_for_retrieval="Represent this sentence for searching relevant passages: ",
                           pooling_method=args.pooling_method)
 
-    task_names = [t.description["name"] for t in MTEB(task_types=args.task_type,
+    # task_names = [t.description["name"] for t in MTEB(task_types=args.task_type,
+    #                                                   task_langs=['en']).tasks]
+    
+    task_names = [t.metadata.name for t in MTEB(task_types=args.task_type,
                                                       task_langs=['en']).tasks]
 
     for task in task_names:

C_MTEB/summarize_results.py

@@ -4,58 +4,47 @@
 from collections import defaultdict
 
 from C_MTEB import *
+import mteb
 from mteb import MTEB
 
 
-def read_results(task_types, except_tasks, args):
+CMTEB_tasks = [
+    'TNews', 'IFlyTek', 'MultilingualSentiment', 'JDReview', 'OnlineShopping', 'Waimai',
+    'CLSClusteringS2S.v2', 'CLSClusteringP2P.v2', 'ThuNewsClusteringS2S.v2', 'ThuNewsClusteringP2P.v2',
+    'Ocnli', 'Cmnli',
+    'T2Reranking', 'MMarcoReranking', 'CMedQAv1-reranking', 'CMedQAv2-reranking',
+    'T2Retrieval', 'MMarcoRetrieval', 'DuRetrieval', 'CovidRetrieval', 'CmedqaRetrieval', 'EcomRetrieval', 'MedicalRetrieval', 'VideoRetrieval',
+    'ATEC', 'BQ', 'LCQMC', 'PAWSX', 'STSB', 'AFQMC', 'QBQTC'
+]
+
+
+def read_results(task_types, args):
     tasks_results = {}
-    model_dirs = {}
+    # model_dirs = {}
     for t_type in task_types:
         tasks_results[t_type] = {}
-        for t in MTEB(task_types=[t_type], task_langs=args.lang).tasks:
-            task_name = t.description["name"]
-            if task_name in except_tasks: continue
+        for t in mteb.get_tasks(task_types=[t_type]):
+            task_name = t.metadata.name
+            if task_name not in CMTEB_tasks:
+                continue
 
-            metric = t.description["main_score"]
+            metric = t.metadata.main_score
             tasks_results[t_type][task_name] = defaultdict(None)
 
-            for model_name in os.listdir(args.results_dir):
-                model_dir = os.path.join(args.results_dir, model_name)
-                if not os.path.isdir(model_dir): continue
-                model_dirs[model_name] = model_dir
-                if os.path.exists(os.path.join(model_dir, task_name + '.json')):
-                    data = json.load(open(os.path.join(model_dir, task_name + '.json')))
-                    for s in ['test', 'dev', 'validation']:
-                        if s in data:
-                            split = s
-                            break
-
-                    if 'en' in args.lang:
-                        if 'en-en' in data[split]:
-                            temp_data = data[split]['en-en']
-                        elif 'en' in data[split]:
-                            temp_data = data[split]['en']
-                        else:
-                            temp_data = data[split]
-                    elif 'zh' in args.lang:
-                        if 'zh' in data[split]:
-                            temp_data = data[split]['zh']
-                        elif 'zh-CN' in data[split]:
-                            temp_data = data[split]['zh-CN']
-                        else:
-                            temp_data = data[split]
-
-                    if metric == 'ap':
-                        tasks_results[t_type][task_name][model_name] = round(temp_data['cos_sim']['ap'] * 100, 2)
-                    elif metric == 'cosine_spearman':
-                        tasks_results[t_type][task_name][model_name] = round(temp_data['cos_sim']['spearman'] * 100, 2)
-                    else:
-                        tasks_results[t_type][task_name][model_name] = round(temp_data[metric] * 100, 2)
-
-    return tasks_results, model_dirs
-
-
-def output_markdown(tasks_results, model_names, save_file):
+            if os.path.exists(os.path.join(args.results_dir, task_name + '.json')):
+                data = json.load(open(os.path.join(args.results_dir, task_name + '.json')))
+                for s in ['test', 'dev', 'validation']:
+                    if s in data['scores']:
+                        split = s
+                        break
+
+                temp_data = data['scores'][split][0]
+                tasks_results[t_type][task_name] = round(temp_data[metric] * 100, 2)
+
+    return tasks_results
+
+
+def output_markdown(tasks_results, model, save_file):
     task_type_res = {}
     with open(save_file, 'w') as f:
         for t_type, type_results in tasks_results.items():
@@ -79,33 +68,29 @@ def output_markdown(tasks_results, model_names, save_file):
             f.write(first_line + ' Avg |  \n')
             f.write(second_line + ':--------:|  \n')
 
-            for model in model_names:
-                write_line = f"| {model} |"
-                all_res = []
-                cqa_res = []
-                for task_name, results in type_results.items():
-                    if "CQADupstack" in task_name:
-                        if model in results:
-                            cqa_res.append(results[model])
-                        continue
-
+            write_line = f"| {model} |"
+            all_res = []
+            cqa_res = []
+            for task_name, results in type_results.items():
+                if "CQADupstack" in task_name:
                     if model in results:
-                        write_line += f" {results[model]} |"
-                        all_res.append(results[model])
-                    else:
-                        write_line += f"  |"
-
-                if len(cqa_res) > 0:
-                    write_line += f" {round(sum(cqa_res) / len(cqa_res), 2)} |"
-                    all_res.append(round(sum(cqa_res) / len(cqa_res), 2))
-
-                # if len(all_res) == len(type_results.keys()):
-                if len(all_res) == task_cnt:
-                    write_line += f" {round(sum(all_res) / len(all_res), 2)} |"
-                    task_type_res[t_type][model] = all_res
-                else:
-                    write_line += f"  |"
-                f.write(write_line + '  \n')
+                        cqa_res.append(results[model])
+                    continue
+
+                write_line += f" {results} |"
+                all_res.append(results)
+
+            if len(cqa_res) > 0:
+                write_line += f" {round(sum(cqa_res) / len(cqa_res), 2)} |"
+                all_res.append(round(sum(cqa_res) / len(cqa_res), 2))
+
+            # if len(all_res) == len(type_results.keys()):
+            if len(all_res) == task_cnt:
+                write_line += f" {round(sum(all_res) / len(all_res), 2)} |"
+                task_type_res[t_type][model] = all_res
+            else:
+                write_line += f"  |"
+            f.write(write_line + '  \n\n')
 
         f.write(f'Overall  \n')
         first_line = "| Model |"
@@ -116,47 +101,45 @@ def output_markdown(tasks_results, model_names, save_file):
         f.write(first_line + ' Avg |  \n')
         f.write(second_line + ':--------:|  \n')
 
-        for model in model_names:
-            write_line = f"| {model} |"
-            all_res = []
-            for type_name, results in task_type_res.items():
-                if model in results:
-                    write_line += f" {round(sum(results[model]) / len(results[model]), 2)} |"
-                    all_res.extend(results[model])
-                else:
-                    write_line += f"  |"
-
-            if len(all_res) > 0:
-                write_line += f" {round(sum(all_res) / len(all_res), 2)} |"
+        write_line = f"| {model} |"
+        all_res = []
+        for type_name, results in task_type_res.items():
+            if model in results:
+                write_line += f" {round(sum(results[model]) / len(results[model]), 2)} |"
+                all_res.extend(results[model])
+            else:
+                write_line += f"  |"
+
+        if len(all_res) > 0:
+            write_line += f" {round(sum(all_res) / len(all_res), 2)} |"
 
-            f.write(write_line + '  \n')
+        f.write(write_line + '  \n')
 
 
 def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument('--results_dir', default="./zh_results", type=str)
     parser.add_argument('--lang', default="zh", type=str)
+    parser.add_argument('--model', default="model", type=str)
     return parser.parse_args()
 
 
 if __name__ == '__main__':
     args = get_args()
 
-    if args.lang == 'zh':
+    if args.lang == 'zho':
         task_types = ["Retrieval", "STS", "PairClassification", "Classification", "Reranking", "Clustering"]
-        except_tasks = []
-        args.lang = ['zh', 'zh-CN']
-    elif args.lang == 'en':
+        args.lang = ['zho']
+    elif args.lang == 'eng':
         task_types = ["Retrieval", "Clustering", "PairClassification", "Reranking", "STS", "Summarization",
                       "Classification"]
-        except_tasks = ['MSMARCOv2']
-        args.lang = ['en']
+        args.lang = ['eng']
     else:
         raise NotImplementedError(f"args.lang must be zh or en, but{args.lang}")
 
-    task_results, model_dirs = read_results(task_types, except_tasks, args=args)
+    task_results = read_results(task_types, args=args)
 
-    output_markdown(task_results, model_dirs.keys(),
+    output_markdown(task_results, args.model,
                     save_file=os.path.join(args.results_dir, f'{args.lang[0]}_results.md'))
 
 

Tutorials/4_Evaluation/4.2.2_MTEB_Leaderboard.ipynb

@@ -143,7 +143,7 @@
     "for task in MTEB_MAIN_EN.tasks:\n",
     "    # get the test set to evaluate on\n",
     "    eval_splits = [\"dev\"] if task == \"MSMARCO\" else [\"test\"]\n",
-    "    evaluation = MTEB(\n",
+    "    evaluation = mteb.MTEB(\n",
     "        tasks=[task], task_langs=[\"en\"]\n",
     "    )  # Remove \"en\" to run all available languages\n",
     "    evaluation.run(\n",