Add Korean Text Search Tasks to MTEB

mteb/abstasks/AbsTaskRetrieval.py

@@ -142,4 +142,4 @@ def encode_corpus(self, corpus: List[Dict[str, str]], batch_size: int, **kwargs)
                 (doc["title"] + self.sep + doc["text"]).strip() if "title" in doc else doc["text"].strip()
                 for doc in corpus
             ]
-        return self.model.encode(sentences, batch_size=batch_size, **kwargs)
+        return self.model.encode(sentences, batch_size=batch_size, **kwargs)

mteb/abstasks/BeIRKOTask.py

@@ -0,0 +1,42 @@
+import os
+from .AbsTask import AbsTask
+
+
+class BeIRKOTask(AbsTask):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def load_data(self, eval_splits=None, **kwargs):
+        """
+        Load dataset from BeIR-KO benchmark.
+        """
+        try:
+            from beir.datasets.data_loader_hf import HFDataLoader
+        except ImportError:
+            raise Exception("Retrieval tasks require beir package. Please install it with `pip install mteb[beir]`")
+
+
+        if self.data_loaded:
+            return
+        if eval_splits is None:
+            eval_splits = self.description["eval_splits"]
+        dataset = self.description["beir_name"]
+
+        hf_repo = self.description["hf_repo"]
+        hf_repo_qrels = self.description["hf_repo_qrels"]
+
+        # cqadupstack not on huggingface yet
+        # dataset, sub_dataset = dataset.split("/") if "cqadupstack" in dataset else (dataset, None)
+
+        self.corpus, self.queries, self.relevant_docs = {}, {}, {}
+
+        for split in eval_splits:
+
+            corpus, queries, qrels = HFDataLoader(hf_repo=hf_repo, hf_repo_qrels=hf_repo_qrels, streaming=False, keep_in_memory=False).load(split=split)
+            # Conversion from DataSet
+            queries = {query['id']: query['text'] for query in queries}
+            corpus = {doc['id']: {'title': doc['title'] , 'text': doc['text']} for doc in corpus}
+
+            self.corpus[split], self.queries[split], self.relevant_docs[split] = corpus, queries, qrels
+
+        self.data_loaded = True

mteb/abstasks/__init__.py

@@ -11,3 +11,4 @@
 from .CrosslingualTask import *
 from .MultilingualTask import *
 from .BeIRPLTask import *
+from .BeIRKOTask import *

mteb/tasks/Retrieval/KoMiracl.py

@@ -0,0 +1,20 @@
+from ...abstasks.AbsTaskRetrieval import AbsTaskRetrieval
+from ...abstasks.BeIRKOTask import BeIRKOTask
+
+
+class KoMiracl(AbsTaskRetrieval, BeIRKOTask):
+    @property
+    def description(self):
+        return {
+            "name": "Ko-miracl",
+            "hf_repo": "taeminlee/Ko-miracl",
+            "hf_repo_qrels": "taeminlee/Ko-miracl",
+            "beir_name": "Ko-miracl",
+            "description": "Ko-miracl",
+            "reference": "",
+            "type": "Retrieval",
+            "category": "s2p",
+            "eval_splits": ["dev"],
+            "eval_langs": ["ko"],
+            "main_score": "ndcg_at_10",
+        }

mteb/tasks/Retrieval/KoMrtydi.py

@@ -0,0 +1,20 @@
+from ...abstasks.AbsTaskRetrieval import AbsTaskRetrieval
+from ...abstasks.BeIRKOTask import BeIRKOTask
+
+
+class KoMrtydi(AbsTaskRetrieval, BeIRKOTask):
+    @property
+    def description(self):
+        return {
+            "name": "Ko-mrtydi",
+            "hf_repo": "taeminlee/Ko-mrtydi",
+            "hf_repo_qrels": "taeminlee/Ko-mrtydi",
+            "beir_name": "Ko-mrtydi",
+            "description": "Ko-mrtydi",
+            "reference": "",
+            "type": "Retrieval",
+            "category": "s2p",
+            "eval_splits": ["dev"],
+            "eval_langs": ["ko"],
+            "main_score": "ndcg_at_10",
+        }

mteb/tasks/Retrieval/KoStrategyQA.py

@@ -0,0 +1,20 @@
+from ...abstasks.AbsTaskRetrieval import AbsTaskRetrieval
+from ...abstasks.BeIRKOTask import BeIRKOTask
+
+
+class KoStrategyQA(AbsTaskRetrieval, BeIRKOTask):
+    @property
+    def description(self):
+        return {
+            "name": "Ko-StrategyQA",
+            "hf_repo": "taeminlee/Ko-StrategyQA",
+            "hf_repo_qrels": "taeminlee/Ko-StrategyQA",
+            "beir_name": "Ko-StrategyQA",
+            "description": "Ko-StrategyQA",
+            "reference": "",
+            "type": "Retrieval",
+            "category": "s2p",
+            "eval_splits": ["dev"],
+            "eval_langs": ["ko"],
+            "main_score": "ndcg_at_10",
+        }

mteb/tasks/Retrieval/__init__.py

@@ -43,3 +43,6 @@
 from .GermanDPRRetrieval import *
 from .GerDaLIRRetrieval import *
 from .XMarketRetrieval import *
+from .KoStrategyQA import *
+from .KoMrtydi import *
+from .KoMiracl import *

scripts/run_mteb_korean.py

@@ -0,0 +1,48 @@
+"""Example script for benchmarking all datasets constituting the MTEB English leaderboard & average scores"""
+
+import logging
+
+from mteb import MTEB
+from sentence_transformers import SentenceTransformer
+
+logging.basicConfig(level=logging.INFO)
+
+logger = logging.getLogger("main")
+
+TASK_LIST_CLASSIFICATION = [
+]
+
+TASK_LIST_CLUSTERING = [
+]
+
+TASK_LIST_PAIR_CLASSIFICATION = [
+]
+
+TASK_LIST_RERANKING = [
+]
+
+TASK_LIST_RETRIEVAL = [
+    'Ko-StrategyQA',
+    'Ko-mrtydi',
+    'Ko-miracl'
+]
+
+TASK_LIST_STS = [
+]
+
+TASK_LIST = (
+    TASK_LIST_CLASSIFICATION
+    + TASK_LIST_CLUSTERING
+    + TASK_LIST_PAIR_CLASSIFICATION
+    + TASK_LIST_RERANKING
+    + TASK_LIST_RETRIEVAL
+    + TASK_LIST_STS
+)
+
+model_name = "average_word_embeddings_komninos"
+model = SentenceTransformer(model_name)
+
+for task in TASK_LIST:
+    logger.info(f"Running task: {task}")
+    evaluation = MTEB(tasks=[task], task_langs=["ko"])  # Remove "ko" for running all languages
+    evaluation.run(model, output_folder=f"results/{model_name}")