Add support for FinMTEB benchmark

mteb/abstasks/AbsTaskClassification.py

@@ -75,6 +75,7 @@ class AbsTaskClassification(AbsTask):
     n_experiments: int = 10
     k: int = 3
     train_split = "train"
+    sentence_column: str = "text"
 
     def evaluate(
         self,
@@ -137,7 +138,7 @@ def _evaluate_subset(
             )
             # Bootstrap `self.samples_per_label` samples per label for each split
             X_sampled, y_sampled, idxs = self._undersample_data(
-                train_split["text"],  # type: ignore
+                train_split[self.sentence_column],  # type: ignore
                 train_split["label"],  # type: ignore
                 self.samples_per_label,
                 idxs,
@@ -146,14 +147,15 @@ def _evaluate_subset(
             evaluator = self.evaluator(
                 X_sampled,
                 y_sampled,
-                eval_split["text"],  # type: ignore
+                eval_split[self.sentence_column],  # type: ignore
                 eval_split["label"],  # type: ignore
                 task_name=self.metadata.name,
                 **params,
             )
             scores_exp, test_cache = evaluator(
                 model, encode_kwargs=encode_kwargs, test_cache=test_cache
             )
+
             scores.append(scores_exp)
 
         avg_scores: dict[str, Any] = {
@@ -183,20 +185,20 @@ def _calculate_metrics_from_split(
     ) -> ClassificationDescriptiveStatistics:
         train_text = []
         if hf_subset:
-            text = self.dataset[hf_subset][split]["text"]
+            text = self.dataset[hf_subset][split][self.sentence_column]
             label = self.dataset[hf_subset][split]["label"]
             if split != "train":
                 train_text = self.dataset[hf_subset]["train"]["text"]
         elif compute_overall:
             text = []
             label = []
             for hf_subset in self.metadata.eval_langs:
-                text.extend(self.dataset[hf_subset][split]["text"])
+                text.extend(self.dataset[hf_subset][split][self.sentence_column])
                 label.extend(self.dataset[hf_subset][split]["label"])
                 if split != "train":
                     train_text.extend(self.dataset[hf_subset]["train"]["text"])
         else:
-            text = self.dataset[split]["text"]
+            text = self.dataset[split][self.sentence_column]
             label = self.dataset[split]["label"]
             if split != "train":
                 train_text = self.dataset["train"]["text"]

mteb/abstasks/AbsTaskPairClassification.py

@@ -66,6 +66,8 @@ class AbsTaskPairClassification(AbsTask):
     """
 
     abstask_prompt = "Retrieve text that are semantically similar to the given text."
+    sentence_1_column: str = "sentence1"
+    sentence_2_column: str = "sentence2"
 
     def _evaluate_subset(
         self,
@@ -80,8 +82,8 @@ def _evaluate_subset(
             "sentence_transformers.evaluation.PairClassificationEvaluator"
         ).setLevel(logging.WARN)
         evaluator = PairClassificationEvaluator(
-            data_split["sentence1"],
-            data_split["sentence2"],
+            data_split[self.sentence_1_column],
+            data_split[self.sentence_2_column],
             data_split["labels"],
             task_name=self.metadata.name,
             **kwargs,
@@ -111,14 +113,14 @@ def _calculate_metrics_from_split(
             dataset = dataset[0]
 
         sentence1 = (
-            dataset["sentence1"][0]
-            if len(dataset["sentence1"]) == 1
-            else dataset["sentence1"]
+            dataset[self.sentence_1_column][0]
+            if len(dataset[self.sentence_1_column]) == 1
+            else dataset[self.sentence_1_column]
         )
         sentence2 = (
-            dataset["sentence2"][0]
-            if len(dataset["sentence2"]) == 1
-            else dataset["sentence2"]
+            dataset[self.sentence_2_column][0]
+            if len(dataset[self.sentence_2_column]) == 1
+            else dataset[self.sentence_2_column]
         )
         labels = (
             dataset["labels"][0] if len(dataset["labels"]) == 1 else dataset["labels"]

mteb/abstasks/AbsTaskSTS.py

@@ -61,6 +61,12 @@ class AbsTaskSTS(AbsTask):
     """
 
     abstask_prompt = "Retrieve semantically similar text."
+    sentence_1_column: str = "sentence1"
+    sentence_2_column: str = "sentence2"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
     min_score: int
     max_score: int
 
@@ -72,8 +78,8 @@ def normalize(x):
 
         normalized_scores = list(map(normalize, data_split["score"]))
         evaluator = STSEvaluator(
-            data_split["sentence1"],
-            data_split["sentence2"],
+            data_split[self.sentence_1_column],
+            data_split[self.sentence_2_column],
             normalized_scores,
             task_name=self.metadata.name,
             **kwargs,
@@ -87,20 +93,24 @@ def _calculate_metrics_from_split(
         self, split: str, hf_subset: str | None = None, compute_overall: bool = False
     ) -> STSDescriptiveStatistics:
         if hf_subset:
-            sentence1 = self.dataset[hf_subset][split]["sentence1"]
-            sentence2 = self.dataset[hf_subset][split]["sentence2"]
+            sentence1 = self.dataset[hf_subset][split][self.sentence_1_column]
+            sentence2 = self.dataset[hf_subset][split][self.sentence_2_column]
             score = self.dataset[hf_subset][split]["score"]
         elif compute_overall:
             sentence1 = []
             sentence2 = []
             score = []
             for hf_subset in self.metadata.eval_langs:
-                sentence1.extend(self.dataset[hf_subset][split]["sentence1"])
-                sentence2.extend(self.dataset[hf_subset][split]["sentence2"])
+                sentence1.extend(
+                    self.dataset[hf_subset][split][self.sentence_1_column]
+                )
+                sentence2.extend(
+                    self.dataset[hf_subset][split][self.sentence_2_column]
+                )
                 score.extend(self.dataset[hf_subset][split]["score"])
         else:
-            sentence1 = self.dataset[split]["sentence1"]
-            sentence2 = self.dataset[split]["sentence2"]
+            sentence1 = self.dataset[split][self.sentence_1_column]
+            sentence2 = self.dataset[split][self.sentence_2_column]
             score = self.dataset[split]["score"]
 
         sentence1_len = [len(s) for s in sentence1]

mteb/abstasks/AbsTaskSummarization.py

@@ -82,6 +82,13 @@ class AbsTaskSummarization(AbsTask):
     abstask_prompt = (
         "Given a news summary, retrieve other semantically similar summaries."
     )
+
+    reference_summaries_column: str = "human_summaries"
+    generated_summaries_column: str = "machine_summaries"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
     # SummEval has DeprecatedSummarizationEvaluator
     evaluator = SummarizationEvaluator
 
@@ -97,9 +104,10 @@ def _evaluate_subset(
             (np.array(x) - self.min_score) / (self.max_score - self.min_score)
             for x in data_split["relevance"]
         ]
-        evaluator = self.evaluator(
-            machine_summaries=data_split["machine_summaries"],
-            human_summaries=data_split["human_summaries"],
+
+        evaluator = self.evalutor(
+            machine_summaries=data_split[self.generated_summaries_column],
+            human_summaries=data_split[self.reference_summaries_column],
             texts=data_split["text"],
             gold_scores=normalized_scores,
             task_name=self.metadata.name,
@@ -114,8 +122,12 @@ def _calculate_metrics_from_split(
     ) -> SummarizationDescriptiveStatistics:
         if hf_subset:
             text = self.dataset[hf_subset][split]["text"]
-            human_summaries = self.dataset[hf_subset][split]["human_summaries"]
-            machine_summaries = self.dataset[hf_subset][split]["machine_summaries"]
+            human_summaries = self.dataset[hf_subset][split][
+                self.reference_summaries_column
+            ]
+            machine_summaries = self.dataset[hf_subset][split][
+                self.generated_summaries_column
+            ]
             relevance = self.dataset[hf_subset][split]["relevance"]
         elif compute_overall:
             text = []
@@ -126,16 +138,16 @@ def _calculate_metrics_from_split(
             for hf_subset in self.metadata.eval_langs:
                 text.extend(self.dataset[hf_subset][split]["text"])
                 human_summaries.extend(
-                    self.dataset[hf_subset][split]["human_summaries"]
+                    self.dataset[hf_subset][split][self.reference_summaries_column]
                 )
                 machine_summaries.extend(
-                    self.dataset[hf_subset][split]["machine_summaries"]
+                    self.dataset[hf_subset][split][self.generated_summaries_column]
                 )
                 relevance.extend(self.dataset[hf_subset][split]["relevance"])
         else:
             text = self.dataset[split]["text"]
-            human_summaries = self.dataset[split]["human_summaries"]
-            machine_summaries = self.dataset[split]["machine_summaries"]
+            human_summaries = self.dataset[split][self.reference_summaries_column]
+            machine_summaries = self.dataset[split][self.generated_summaries_column]
             relevance = self.dataset[split]["relevance"]
 
         all_human_summaries = []

mteb/benchmarks/benchmarks.py

@@ -14,7 +14,8 @@
 
 http_url_adapter = TypeAdapter(AnyUrl)
 UrlString = Annotated[
-    str, BeforeValidator(lambda value: str(http_url_adapter.validate_python(value)))
+    str,
+    BeforeValidator(lambda value: str(http_url_adapter.validate_python(value))),
 ]  # Allows the type to be a string, but ensures that the string is a URL
 
 
@@ -1471,3 +1472,86 @@ def load_results(
           url={https://arxiv.org/abs/2305.19840},
 }""",
 )
+
+FinMTEB = Benchmark(
+    name="MTEB(Finance)",
+    tasks=get_tasks(
+        languages=["eng", "zho"],
+        tasks=[
+            "FinancialPhraseBankClassification",
+            "FinSentClassification",
+            "FiQAClassification",
+            "SemEva2017Classification",
+            "FLSClassification",
+            "ESGClassification",
+            "FOMCClassification",
+            "FinancialFraudClassification",
+            "FinNSPClassification",
+            "FinChinaSentimentClassification",
+            "FinFEClassification",
+            "OpenFinDataSentimentClassification",
+            "Weibo21Classification",
+            "MInDS14EnClustering",
+            "ComplaintsClustering",
+            "PiiClustering",
+            "FinanceArxivS2SClustering",
+            "FinanceArxivP2PClustering",
+            "WikiCompany2IndustryClustering",
+            "MInDS14ZhClustering",
+            "FinNLClustering",
+            "CCKS2022Clustering",
+            "CCKS2020Clustering",
+            "CCKS2019Clustering",
+            "HeadlineACPairClassification",
+            "HeadlinePDDPairClassification",
+            "HeadlinePDUPairClassification",
+            "AFQMCPairClassification",
+            "FinFactReranking",
+            "FiQA2018Reranking",
+            "HC3Reranking",
+            "FinEvaReranking",
+            "DISCFinLLMReranking",
+            "FiQA2018",
+            "FinanceBenchRetrieval",
+            "HC3Retrieval",
+            "Apple10KRetrieval",
+            "FinQARetrieval",
+            "TATQARetrieval",
+            "USNewsRetrieval",
+            "TradeTheEventEncyclopediaRetrieval",
+            "TradeTheEventNewsRetrieval",
+            "TheGoldmanEnRetrieval",
+            "FinTruthQARetrieval",
+            "FinEvaRetrieval",
+            "AlphaFinRetrieval",
+            "DISCFinLLMRetrieval",
+            "DISCFinLLMComputingRetrieval",
+            "DuEEFinRetrieval",
+            "SmoothNLPRetrieval",
+            "THUCNewsRetrieval",
+            "FinEvaEncyclopediaRetrieval",
+            "TheGoldmanZhRetrieval",
+            "FINAL",
+            "FinSTS",
+            "AFQMC",
+            "BQCorpus",
+            "Ectsum",
+            "FINDsum",
+            "FNS2022sum",
+            "FiNNAsum",
+            "FinEvaHeadlinesum",
+            "FinEvasum",
+        ],
+    ),
+    description="FinMTEB is an embedding benchmark consists of 64 financial domain-specific text datasets, across English and Chinese.",
+    reference="https://arxiv.org/abs/2409.18511v1",
+    citation="""@misc{tang2024needdomainspecificembeddingmodels,
+          title={Do We Need Domain-Specific Embedding Models? An Empirical Investigation},
+          author={Yixuan Tang and Yi Yang},
+          year={2024},
+          eprint={2409.18511},
+          archivePrefix={arXiv},
+          primaryClass={cs.CL},
+          url={https://arxiv.org/abs/2409.18511},
+    }""",
+)

mteb/tasks/Classification/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+
 from .ara import (
     AJGT,
     HotelReviewSentimentClassification,

mteb/tasks/Classification/eng/ESGClassification.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from mteb.abstasks.AbsTaskClassification import AbsTaskClassification
+from mteb.abstasks.TaskMetadata import TaskMetadata
+
+
+class ESGClassification(AbsTaskClassification):
+    metadata = TaskMetadata(
+        name="ESGClassification",
+        description="A finance dataset performs sentence classification under the environmental, social, and corporate governance (ESG) framework.",
+        reference="https://arxiv.org/abs/2309.13064",
+        dataset={
+            "path": "FinanceMTEB/ESG",
+            "revision": "521d56feabadda80b11d6adcc6b335d4c5ad8285",
+        },
+        type="Classification",
+        category="s2s",
+        eval_splits=["test"],
+        eval_langs=["eng-Latn"],
+        main_score="accuracy",
+    )

mteb/tasks/Classification/eng/FLSClassification.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from mteb.abstasks.AbsTaskClassification import AbsTaskClassification
+from mteb.abstasks.TaskMetadata import TaskMetadata
+
+
+class FLSClassification(AbsTaskClassification):
+    metadata = TaskMetadata(
+        name="FLSClassification",
+        description="A finance dataset detects whether the sentence is a forward-looking statement.",
+        reference="https://arxiv.org/abs/2309.13064",
+        dataset={
+            "path": "FinanceMTEB/FLS",
+            "revision": "39b6719f1d7197df4498fea9fce20d4ad782a083",
+        },
+        type="Classification",
+        category="s2s",
+        eval_splits=["test"],
+        eval_langs=["eng-Latn"],
+        main_score="accuracy",
+    )

mteb/tasks/Classification/eng/FOMCClassification.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from mteb.abstasks.AbsTaskClassification import AbsTaskClassification
+from mteb.abstasks.TaskMetadata import TaskMetadata
+
+
+class FOMCClassification(AbsTaskClassification):
+    metadata = TaskMetadata(
+        name="FOMCClassification",
+        description="A task of hawkish-dovish classification in finance domain.",
+        reference="https://github.com/gtfintechlab/fomc-hawkish-dovish",
+        dataset={
+            "path": "FinanceMTEB/FOMC",
+            "revision": "cdaf1306a24bc5e7441c7c871343efdf4c721bc2",
+        },
+        type="Classification",
+        category="s2s",
+        eval_splits=["test"],
+        eval_langs=["eng-Latn"],
+        main_score="accuracy",
+    )
+
+    def dataset_transform(self):
+        self.dataset = self.dataset.rename_column("sentence", "text")