From 98bf9fdd1243f7fbe2d6929ebc72b6f709e8ecda Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Thu, 22 May 2025 22:34:56 +0000 Subject: [PATCH 01/17] Add RuSciBench --- mteb/tasks/BitextMining/__init__.py | 1 + .../multilingual/RuSciBenchBitexMining.py | 71 ++++++++++ mteb/tasks/Classification/__init__.py | 1 + .../multilingual/RuSciBenchClassification.py | 91 +++++++++++++ mteb/tasks/Retrieval/__init__.py | 1 + .../multilingual/RuSciBenchRetrieval.py | 127 ++++++++++++++++++ 6 files changed, 292 insertions(+) create mode 100644 mteb/tasks/BitextMining/multilingual/RuSciBenchBitexMining.py create mode 100644 mteb/tasks/Classification/multilingual/RuSciBenchClassification.py create mode 100644 mteb/tasks/Retrieval/multilingual/RuSciBenchRetrieval.py diff --git a/mteb/tasks/BitextMining/__init__.py b/mteb/tasks/BitextMining/__init__.py index 3d3a4c3538..86954506d3 100644 --- a/mteb/tasks/BitextMining/__init__.py +++ b/mteb/tasks/BitextMining/__init__.py @@ -22,6 +22,7 @@ from .multilingual.NusaXBitextMining import * from .multilingual.PhincBitextMining import * from .multilingual.RomaTalesBitextMining import * +from .multilingual.RuSciBenchBitexMining import * from .multilingual.TatoebaBitextMining import * from .multilingual.WebFAQBitextMining import * from .srn.SRNCorpusBitextMining import * diff --git a/mteb/tasks/BitextMining/multilingual/RuSciBenchBitexMining.py b/mteb/tasks/BitextMining/multilingual/RuSciBenchBitexMining.py new file mode 100644 index 0000000000..0e6937b026 --- /dev/null +++ b/mteb/tasks/BitextMining/multilingual/RuSciBenchBitexMining.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +import datasets + +from mteb.abstasks.AbsTaskBitextMining import AbsTaskBitextMining +from mteb.abstasks.MultilingualTask import MultilingualTask +from mteb.abstasks.TaskMetadata import TaskMetadata + +_LANGUAGES = { + "ru-en": ["rus-Cyrl", "eng-Latn"], + "en-ru": ["rus-Cyrl", "rus-Cyrl"], +} + + +_SPLITS = ["test"] + + +class RuSciBenchBitexMining(AbsTaskBitextMining, MultilingualTask): + fast_loading = True + metadata = TaskMetadata( + name="RuSciBenchBitexMining", + dataset={ + "path": "mlsa-iai-msu-lab/ru_sci_bench_bitext_mining", + "revision": "927d95897a168b79568e96591276b995ed1c4da8", + }, + description="Find translation of a scientific article", + reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", + type="BitextMining", + category="p2p", + modalities=["text"], + eval_splits=_SPLITS, + eval_langs=_LANGUAGES, + main_score="f1", + date=("2007-01-01", "2023-01-01"), + domains=["Academic", "Non-fiction", "Written"], + task_subtypes=[], + license="not specified", + dialect=[], + sample_creation="found", + annotations_creators="derived", + bibtex_citation=""" +@article{vatolin2024ruscibench, + author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, + title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + journal = {Doklady Mathematics}, + year = {2024}, + volume = {110}, + number = {1}, + pages = {S251--S260}, + month = {12}, + doi = {10.1134/S1064562424602191}, + url = {https://doi.org/10.1134/S1064562424602191}, + issn = {1531-8362} +}""", + prompt="Given the following title and abstract of the scientific article, find its translation", + ) + + def load_data(self, **kwargs): + if self.data_loaded: + return + + self.dataset = {} + for lang in self.hf_subsets: + self.dataset.setdefault(lang, {})[_SPLITS[0]] = datasets.load_dataset( + split=_SPLITS[0], + name=lang, + **self.metadata_dict["dataset"], + ) + + self.dataset_transform() + self.data_loaded = True diff --git a/mteb/tasks/Classification/__init__.py b/mteb/tasks/Classification/__init__.py index eafd1a7109..ca1d3b1ed8 100644 --- a/mteb/tasks/Classification/__init__.py +++ b/mteb/tasks/Classification/__init__.py @@ -107,6 +107,7 @@ from .multilingual.NusaParagraphEmotionClassification import * from .multilingual.NusaParagraphTopicClassification import * from .multilingual.NusaXSenti import * +from .multilingual.RuSciBenchClassification import * from .multilingual.ScalaClassification import * from .multilingual.ScandiSentClassification import * from .multilingual.SIB200Classification import * diff --git a/mteb/tasks/Classification/multilingual/RuSciBenchClassification.py b/mteb/tasks/Classification/multilingual/RuSciBenchClassification.py new file mode 100644 index 0000000000..f7f778de47 --- /dev/null +++ b/mteb/tasks/Classification/multilingual/RuSciBenchClassification.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +from mteb.abstasks.AbsTaskClassification import AbsTaskClassification +from mteb.abstasks.MultilingualTask import MultilingualTask +from mteb.abstasks.TaskMetadata import TaskMetadata + + +class RuSciBenchCoreRiscClassification(MultilingualTask, AbsTaskClassification): + metadata = TaskMetadata( + name="RuSciBenchCoreRiscClassification", + dataset={ + "path": "mlsa-iai-msu-lab/ru_sci_bench_mteb", + "revision": "fbc0599a0b5f00b3c7d87ab4d13490f04fb77f8e", + }, + description="Classification of scientific papers (title+abstract) by publication Core RISC status", + reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", + type="Classification", + category="p2p", + modalities=["text"], + eval_splits=["test"], + eval_langs={ + "corerisc_ru": ["rus-Cyrl"], + "corerisc_en": ["eng-Latn"], + }, + main_score="accuracy", + date=("2007-01-01", "2023-01-01"), + domains=["Academic", "Non-fiction", "Written"], + task_subtypes=[], + license="mit", + annotations_creators="derived", + dialect=[], + sample_creation="found", + bibtex_citation=""" +@article{vatolin2024ruscibench, + author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, + title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + journal = {Doklady Mathematics}, + year = {2024}, + volume = {110}, + number = {1}, + pages = {S251--S260}, + month = {12}, + doi = {10.1134/S1064562424602191}, + url = {https://doi.org/10.1134/S1064562424602191}, + issn = {1531-8362} +}""", + prompt="Classify whether a scientific article is part of the core RISC or not based on the title and abstract", + ) + + +class RuSciBenchPubTypeClassification(MultilingualTask, AbsTaskClassification): + metadata = TaskMetadata( + name="RuSciBenchPubTypeClassification", + dataset={ + "path": "mlsa-iai-msu-lab/ru_sci_bench_mteb", + "revision": "fbc0599a0b5f00b3c7d87ab4d13490f04fb77f8e", + }, + description="Classification of scientific papers (title+abstract) by publication type", + reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", + type="Classification", + category="p2p", + modalities=["text"], + eval_splits=["test"], + eval_langs={ + "pub_type_ru": ["rus-Cyrl"], + "pub_type_en": ["eng-Latn"], + }, + main_score="accuracy", + date=("2007-01-01", "2023-01-01"), + domains=["Academic", "Non-fiction", "Written"], + task_subtypes=[], + license="mit", + annotations_creators="derived", + dialect=[], + sample_creation="found", + bibtex_citation=""" +@article{vatolin2024ruscibench, + author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, + title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + journal = {Doklady Mathematics}, + year = {2024}, + volume = {110}, + number = {1}, + pages = {S251--S260}, + month = {12}, + doi = {10.1134/S1064562424602191}, + url = {https://doi.org/10.1134/S1064562424602191}, + issn = {1531-8362} +}""", + prompt="Classify the type of scientific paper based on the title and abstract", + ) diff --git a/mteb/tasks/Retrieval/__init__.py b/mteb/tasks/Retrieval/__init__.py index b2abdd19e5..01f6deb2ef 100644 --- a/mteb/tasks/Retrieval/__init__.py +++ b/mteb/tasks/Retrieval/__init__.py @@ -134,6 +134,7 @@ from .multilingual.NeuCLIR2022Retrieval import * from .multilingual.NeuCLIR2023Retrieval import * from .multilingual.PublicHealthQARetrieval import * +from .multilingual.RuSciBenchRetrieval import * from .multilingual.StatcanDialogueDatasetRetrieval import * from .multilingual.WebFAQRetrieval import * from .multilingual.WikipediaRetrievalMultilingual import * diff --git a/mteb/tasks/Retrieval/multilingual/RuSciBenchRetrieval.py b/mteb/tasks/Retrieval/multilingual/RuSciBenchRetrieval.py new file mode 100644 index 0000000000..94b639d2b7 --- /dev/null +++ b/mteb/tasks/Retrieval/multilingual/RuSciBenchRetrieval.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +from mteb.abstasks.AbsTaskRetrieval import AbsTaskRetrieval +from mteb.abstasks.MultilingualTask import MultilingualTask +from mteb.abstasks.TaskMetadata import TaskMetadata +from mteb.tasks.Retrieval.multilingual.NeuCLIR2023Retrieval import load_neuclir_data + +_LANGUAGES = { + "ru": ["rus-Cyrl"], + "en": ["eng-Latn"], +} + + +class RuSciBenchCiteRetrieval(AbsTaskRetrieval, MultilingualTask): + ignore_identical_ids = True + + metadata = TaskMetadata( + name="RuSciBenchCiteRetrieval", + dataset={ + "path": "mlsa-iai-msu-lab/ru_sci_bench_cite_retrieval", + "revision": "6cb447d02f41b8b775d5d9df7faf472f44d2f1db", + }, + description="Retrieval of related scientific papers based on their title and abstract", + reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", + type="Retrieval", + category="p2p", + modalities=["text"], + eval_splits=["test"], + eval_langs=_LANGUAGES, + main_score="ndcg_at_10", + date=("2007-01-01", "2023-01-01"), + domains=["Academic", "Non-fiction", "Written"], + task_subtypes=["Article retrieval"], + license="mit", + dialect=[], + sample_creation="found", + annotations_creators="derived", + bibtex_citation=""" +@article{vatolin2024ruscibench, + author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, + title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + journal = {Doklady Mathematics}, + year = {2024}, + volume = {110}, + number = {1}, + pages = {S251--S260}, + month = {12}, + doi = {10.1134/S1064562424602191}, + url = {https://doi.org/10.1134/S1064562424602191}, + issn = {1531-8362} +}""", + prompt={ + "query": "Given a title and abstract of a scientific paper, retrieve the titles and abstracts of other relevant papers", + }, + ) + + def load_data(self, **kwargs): + if self.data_loaded: + return + + self.corpus, self.queries, self.relevant_docs = load_neuclir_data( + path=self.metadata_dict["dataset"]["path"], + langs=self.metadata.eval_langs, + eval_splits=self.metadata_dict["eval_splits"], + cache_dir=kwargs.get("cache_dir", None), + revision=self.metadata_dict["dataset"]["revision"], + ) + + self.data_loaded = True + + +class RuSciBenchCociteRetrieval(MultilingualTask, AbsTaskRetrieval): + ignore_identical_ids = True + + metadata = TaskMetadata( + name="RuSciBenchCociteRetrieval", + dataset={ + "path": "mlsa-iai-msu-lab/ru_sci_bench_cocite_retrieval", + "revision": "a5da47a245275669d2b6ddf8f96c5338dd2428b4", + }, + description="Retrieval of related scientific papers based on their title and abstract", + reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", + type="Retrieval", + category="p2p", + modalities=["text"], + eval_splits=["test"], + eval_langs=_LANGUAGES, + main_score="ndcg_at_10", + date=("2007-01-01", "2023-01-01"), + domains=["Academic", "Non-fiction", "Written"], + task_subtypes=["Article retrieval"], + license="mit", + dialect=[], + sample_creation="found", + annotations_creators="derived", + bibtex_citation=""" +@article{vatolin2024ruscibench, + author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, + title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + journal = {Doklady Mathematics}, + year = {2024}, + volume = {110}, + number = {1}, + pages = {S251--S260}, + month = {12}, + doi = {10.1134/S1064562424602191}, + url = {https://doi.org/10.1134/S1064562424602191}, + issn = {1531-8362} +}""", + prompt={ + "query": "Given a title and abstract of a scientific paper, retrieve the titles and abstracts of other relevant papers", + }, + ) + + def load_data(self, **kwargs): + if self.data_loaded: + return + + self.corpus, self.queries, self.relevant_docs = load_neuclir_data( + path=self.metadata_dict["dataset"]["path"], + langs=self.metadata.eval_langs, + eval_splits=self.metadata_dict["eval_splits"], + cache_dir=kwargs.get("cache_dir", None), + revision=self.metadata_dict["dataset"]["revision"], + ) + + self.data_loaded = True From a95f42dbbab45b5eb785d2e5e00ba0cac15316e5 Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Thu, 22 May 2025 22:42:54 +0000 Subject: [PATCH 02/17] fix bitext mining lang --- mteb/tasks/BitextMining/multilingual/RuSciBenchBitexMining.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mteb/tasks/BitextMining/multilingual/RuSciBenchBitexMining.py b/mteb/tasks/BitextMining/multilingual/RuSciBenchBitexMining.py index 0e6937b026..844ff6b514 100644 --- a/mteb/tasks/BitextMining/multilingual/RuSciBenchBitexMining.py +++ b/mteb/tasks/BitextMining/multilingual/RuSciBenchBitexMining.py @@ -8,7 +8,7 @@ _LANGUAGES = { "ru-en": ["rus-Cyrl", "eng-Latn"], - "en-ru": ["rus-Cyrl", "rus-Cyrl"], + "en-ru": ["eng-Latn", "rus-Cyrl"], } From 7ec15ddd9c65f107ee28c0dd8eb44dca622f34fd Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Sun, 13 Jul 2025 17:50:53 +0000 Subject: [PATCH 03/17] Add regression task --- mteb/abstasks/AbsTaskRegression.py | 171 ++++++++++++++++++ mteb/abstasks/TaskMetadata.py | 7 +- mteb/abstasks/__init__.py | 1 + mteb/benchmarks/benchmarks.py | 39 ++++ .../evaluators/RegressionEvaluator.py | 75 ++++++++ mteb/evaluation/evaluators/__init__.py | 1 + mteb/tasks/BitextMining/__init__.py | 2 +- ...texMining.py => RuSciBenchBitextMining.py} | 50 ++--- mteb/tasks/Classification/__init__.py | 1 + .../multilingual/RuSciBenchClassification.py | 145 ++++++++++++--- 10 files changed, 429 insertions(+), 63 deletions(-) create mode 100644 mteb/abstasks/AbsTaskRegression.py create mode 100644 mteb/evaluation/evaluators/RegressionEvaluator.py rename mteb/tasks/BitextMining/multilingual/{RuSciBenchBitexMining.py => RuSciBenchBitextMining.py} (50%) diff --git a/mteb/abstasks/AbsTaskRegression.py b/mteb/abstasks/AbsTaskRegression.py new file mode 100644 index 0000000000..097fd66ab4 --- /dev/null +++ b/mteb/abstasks/AbsTaskRegression.py @@ -0,0 +1,171 @@ +from __future__ import annotations + +import logging +from typing import Any + +from mteb.abstasks.TaskMetadata import DescriptiveStatistics, HFSubset +from mteb.encoder_interface import Encoder +from mteb.evaluation.evaluators.RegressionEvaluator import LinearRegressionEvaluator + +from ..load_results.task_results import ScoresDict +from .AbsTask import AbsTask + +logger = logging.getLogger(__name__) + + +class RegressionDescriptiveStatistics(DescriptiveStatistics): + """Descriptive statistics for Regression + + Attributes: + num_samples: number of samples in the dataset. + number_of_characters: Total number of symbols in the dataset. + num_texts_in_train: Number of texts in the train split + + min_text_length: Minimum length of text + average_text_length: Average length of text + max_text_length: Maximum length of text + unique_text: Number of unique texts + + min_value: Minimum of the target variable + average_value: Average of the target variable + max_value: Maximum of the target variable + """ + + num_samples: int + number_of_characters: int + num_texts_in_train: int | None + + min_text_length: int + average_text_length: float + max_text_length: int + unique_text: int + + min_value: float + average_value: float + max_value: float + + +class AbsTaskRegression(AbsTask): + """Abstract class for regression tasks + + self.load_data() must generate a huggingface dataset with a split matching self.metadata_dict["eval_splits"], and assign it to self.dataset. It + must contain the following columns: + text: str + value: float + """ + + def __init__(self, seed: int = 42, **kwargs: Any): + super().__init__(seed, **kwargs) + if hasattr(self, "metadata"): + self.metadata + + def _evaluate_subset( + self, + model: Encoder, + dataset, + eval_split: str = "test", + train_split: str = "train", + encode_kwargs: dict[str, Any] = {}, + **kwargs, + ) -> ScoresDict: + train_split = dataset[train_split] + eval_split = dataset[eval_split] + + evaluator = LinearRegressionEvaluator( + train_split["text"], + train_split["value"], + eval_split["text"], + eval_split["value"], + task_name=self.metadata.name, + encode_kwargs=encode_kwargs, + **kwargs, + ) + scores = evaluator(model) + return scores + + def _add_main_score(self, scores): + scores["main_score"] = scores[self.metadata.main_score] + + def evaluate( + self, + model: Encoder, + eval_split: str = "test", + train_split: str = "train", + *, + encode_kwargs: dict[str, Any] = {}, + **kwargs: Any, + ) -> dict[HFSubset, ScoresDict]: + if not self.data_loaded: + self.load_data() + + scores = {} + hf_subsets = list(self.dataset) if self.is_multilingual else ["default"] + + for hf_subset in hf_subsets: + logger.info( + f"\nTask: {self.metadata.name}, split: {eval_split}, subset: {hf_subset}. Running..." + ) + + if hf_subset not in self.dataset and hf_subset == "default": + ds = self.dataset + else: + ds = self.dataset[hf_subset] + scores[hf_subset] = self._evaluate_subset( + model, + ds, + eval_split, + train_split, + encode_kwargs=encode_kwargs, + **kwargs, + ) + self._add_main_score(scores[hf_subset]) + + return scores + + def __hash__(self) -> int: + return hash(self.metadata) + + def _calculate_metrics_from_split( + self, split: str, hf_subset: str | None = None, compute_overall: bool = False + ) -> RegressionDescriptiveStatistics: + train_text = [] + if hf_subset: + texts = self.dataset[hf_subset][split]["text"] + values = self.dataset[hf_subset][split]["value"] + if split != "train": + train_text = self.dataset[hf_subset]["train"]["text"] + elif compute_overall: + texts = [] + values = [] + for lang_subset in self.metadata.eval_langs: + texts.extend(self.dataset[lang_subset][split]["text"]) + values.extend(self.dataset[lang_subset][split]["value"]) + if split != "train": + train_text.extend(self.dataset[lang_subset]["train"]["text"]) + else: + texts = self.dataset[split]["text"] + values = self.dataset[split]["value"] + if split != "train": + train_text = self.dataset["train"]["text"] + + text_lengths = [len(t) for t in texts] + total_text_length = sum(text_lengths) + + num_texts_in_train_val = ( + len(set(texts) & set(train_text)) if split != "train" else None + ) + + return RegressionDescriptiveStatistics( + num_samples=len(texts), + number_of_characters=total_text_length, + num_texts_in_train=num_texts_in_train_val, + min_text_length=min(text_lengths) if text_lengths else 0, + average_text_length=(total_text_length / len(texts)) + if len(texts) > 0 + else 0, + max_text_length=max(text_lengths) if text_lengths else 0, + unique_text=len(set(texts)), + min_value=min(values) if values else 0.0, + average_value=(sum(values) / len(values)) if len(values) > 0 else 0.0, + max_value=max(values) if values else 0.0, + ) diff --git a/mteb/abstasks/TaskMetadata.py b/mteb/abstasks/TaskMetadata.py index ac8b52890d..f2c631b3b2 100644 --- a/mteb/abstasks/TaskMetadata.py +++ b/mteb/abstasks/TaskMetadata.py @@ -114,15 +114,16 @@ _TASK_TYPE = ( "BitextMining", "Classification", - "MultilabelClassification", "Clustering", + "InstructionRetrieval", + "MultilabelClassification", "PairClassification", + "Regression", "Reranking", "Retrieval", + "Speed", "STS", "Summarization", - "InstructionRetrieval", - "Speed", ) + MIEB_TASK_TYPE TASK_TYPE = Literal[_TASK_TYPE] diff --git a/mteb/abstasks/__init__.py b/mteb/abstasks/__init__.py index 720f8747e8..d5401fb00e 100644 --- a/mteb/abstasks/__init__.py +++ b/mteb/abstasks/__init__.py @@ -8,6 +8,7 @@ from .AbsTaskInstructionRetrieval import * from .AbsTaskMultilabelClassification import * from .AbsTaskPairClassification import * +from .AbsTaskRegression import * from .AbsTaskReranking import * from .AbsTaskRetrieval import * from .AbsTaskSpeedTask import * diff --git a/mteb/benchmarks/benchmarks.py b/mteb/benchmarks/benchmarks.py index 7f47cfe689..bc16218212 100644 --- a/mteb/benchmarks/benchmarks.py +++ b/mteb/benchmarks/benchmarks.py @@ -250,6 +250,45 @@ """, ) + +RU_SCI_BENCH = Benchmark( + name="RuSciBench", + tasks=get_tasks( + tasks=[ + # BitextMining + "RuSciBenchBitextMining", + # Classification + "RuSciBenchCoreRiscClassification", + "RuSciBenchGRNTIOrigClassification", + "RuSciBenchOECDOrigClassification", + "RuSciBenchPubTypeClassification", + # Retrieval + "RuSciBenchCiteRetrieval", + "RuSciBenchCociteRetrieval", + # Regression + "RuSciBenchCitedCountRegression", + "RuSciBenchYearPublRegression", + ], + ), + description="RuSciBench is a benchmark designed for evaluating sentence encoders and language models on scientific texts in both Russian and English. The data is sourced from eLibrary (www.elibrary.ru), Russia's largest electronic library of scientific publications. This benchmark facilitates the evaluation and comparison of models on various research-related tasks.", + reference="https://link.springer.com/article/10.1134/S1064562424602191", + citation=r""" +@article{vatolin2024ruscibench, + author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, + doi = {10.1134/S1064562424602191}, + issn = {1531-8362}, + journal = {Doklady Mathematics}, + month = {12}, + number = {1}, + pages = {S251--S260}, + title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + url = {https://doi.org/10.1134/S1064562424602191}, + volume = {110}, + year = {2024}, +} +""", +) + MTEB_RETRIEVAL_WITH_INSTRUCTIONS = Benchmark( name="FollowIR", display_name="Instruction Following", diff --git a/mteb/evaluation/evaluators/RegressionEvaluator.py b/mteb/evaluation/evaluators/RegressionEvaluator.py new file mode 100644 index 0000000000..3c5cae128d --- /dev/null +++ b/mteb/evaluation/evaluators/RegressionEvaluator.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import logging +from typing import Any + +import numpy as np +from scipy.stats import kendalltau +from sklearn.linear_model import LinearRegression +from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score + +from mteb.encoder_interface import Encoder + +from .Evaluator import Evaluator + +logger = logging.getLogger(__name__) + + +class LinearRegressionEvaluator(Evaluator): + def __init__( + self, + sentences_train, + y_train, + sentences_test, + y_test, + task_name: str | None = None, + fit_intercept: bool = True, + encode_kwargs: dict[str, Any] = {}, + limit: int | None = None, + **kwargs, + ): + super().__init__(**kwargs) + if limit is not None: + sentences_train = sentences_train[:limit] + y_train = y_train[:limit] + sentences_test = sentences_test[:limit] + y_test = y_test[:limit] + self.sentences_train = sentences_train + self.y_train = y_train + self.sentences_test = sentences_test + self.y_test = y_test + + self.task_name = task_name + self.encode_kwargs = encode_kwargs + + if "batch_size" not in self.encode_kwargs: + self.encode_kwargs["batch_size"] = 32 + + self.fit_intercept = fit_intercept + + def __call__(self, model: Encoder) -> dict[str, float]: + scores = {} + X_train = model.encode( + self.sentences_train, + model=model, + task_name=self.task_name, + **self.encode_kwargs, + ) + X_test = model.encode( + self.sentences_test, + model=model, + task_name=self.task_name, + **self.encode_kwargs, + ) + + linear_regression = LinearRegression(fit_intercept=self.fit_intercept) + linear_regression.fit(X_train, self.y_train) + y_pred = linear_regression.predict(X_test) + + scores["mae"] = mean_absolute_error(self.y_test, y_pred) + scores["mse"] = mean_squared_error(self.y_test, y_pred) + scores["rmse"] = np.sqrt(scores["mse"]) + scores["r2"] = r2_score(self.y_test, y_pred) + scores["kendalltau"] = kendalltau(self.y_test, y_pred).statistic + + return scores diff --git a/mteb/evaluation/evaluators/__init__.py b/mteb/evaluation/evaluators/__init__.py index ebc6aa6f94..69301b41e5 100644 --- a/mteb/evaluation/evaluators/__init__.py +++ b/mteb/evaluation/evaluators/__init__.py @@ -11,6 +11,7 @@ from .Image.VisualSTSEvaluator import * from .Image.ZeroShotClassificationEvaluator import * from .PairClassificationEvaluator import * +from .RegressionEvaluator import * from .RerankingEvaluator import * from .RetrievalEvaluator import * from .STSEvaluator import * diff --git a/mteb/tasks/BitextMining/__init__.py b/mteb/tasks/BitextMining/__init__.py index 86954506d3..93e8d670b1 100644 --- a/mteb/tasks/BitextMining/__init__.py +++ b/mteb/tasks/BitextMining/__init__.py @@ -22,7 +22,7 @@ from .multilingual.NusaXBitextMining import * from .multilingual.PhincBitextMining import * from .multilingual.RomaTalesBitextMining import * -from .multilingual.RuSciBenchBitexMining import * +from .multilingual.RuSciBenchBitextMining import * from .multilingual.TatoebaBitextMining import * from .multilingual.WebFAQBitextMining import * from .srn.SRNCorpusBitextMining import * diff --git a/mteb/tasks/BitextMining/multilingual/RuSciBenchBitexMining.py b/mteb/tasks/BitextMining/multilingual/RuSciBenchBitextMining.py similarity index 50% rename from mteb/tasks/BitextMining/multilingual/RuSciBenchBitexMining.py rename to mteb/tasks/BitextMining/multilingual/RuSciBenchBitextMining.py index 844ff6b514..2e91e935b3 100644 --- a/mteb/tasks/BitextMining/multilingual/RuSciBenchBitexMining.py +++ b/mteb/tasks/BitextMining/multilingual/RuSciBenchBitextMining.py @@ -1,7 +1,5 @@ from __future__ import annotations -import datasets - from mteb.abstasks.AbsTaskBitextMining import AbsTaskBitextMining from mteb.abstasks.MultilingualTask import MultilingualTask from mteb.abstasks.TaskMetadata import TaskMetadata @@ -15,15 +13,15 @@ _SPLITS = ["test"] -class RuSciBenchBitexMining(AbsTaskBitextMining, MultilingualTask): +class RuSciBenchBitextMining(AbsTaskBitextMining, MultilingualTask): fast_loading = True metadata = TaskMetadata( - name="RuSciBenchBitexMining", + name="RuSciBenchBitextMining", dataset={ "path": "mlsa-iai-msu-lab/ru_sci_bench_bitext_mining", - "revision": "927d95897a168b79568e96591276b995ed1c4da8", + "revision": "e5840033c5cf2573932db027ac8001fe0a7eb6fa", }, - description="Find translation of a scientific article", + description="This task focuses on finding translations of scientific articles. The dataset is sourced from eLibrary, Russia's largest electronic library of scientific publications. Russian authors often provide English translations for their abstracts and titles, and the data consists of these paired titles and abstracts. The task evaluates a model's ability to match an article's Russian title and abstract to its English counterpart, or vice versa.", reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", type="BitextMining", category="p2p", @@ -38,34 +36,20 @@ class RuSciBenchBitexMining(AbsTaskBitextMining, MultilingualTask): dialect=[], sample_creation="found", annotations_creators="derived", - bibtex_citation=""" + bibtex_citation=r""" @article{vatolin2024ruscibench, - author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, - title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, + doi = {10.1134/S1064562424602191}, + issn = {1531-8362}, journal = {Doklady Mathematics}, - year = {2024}, - volume = {110}, - number = {1}, - pages = {S251--S260}, - month = {12}, - doi = {10.1134/S1064562424602191}, - url = {https://doi.org/10.1134/S1064562424602191}, - issn = {1531-8362} -}""", + month = {12}, + number = {1}, + pages = {S251--S260}, + title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + url = {https://doi.org/10.1134/S1064562424602191}, + volume = {110}, + year = {2024}, +} +""", prompt="Given the following title and abstract of the scientific article, find its translation", ) - - def load_data(self, **kwargs): - if self.data_loaded: - return - - self.dataset = {} - for lang in self.hf_subsets: - self.dataset.setdefault(lang, {})[_SPLITS[0]] = datasets.load_dataset( - split=_SPLITS[0], - name=lang, - **self.metadata_dict["dataset"], - ) - - self.dataset_transform() - self.data_loaded = True diff --git a/mteb/tasks/Classification/__init__.py b/mteb/tasks/Classification/__init__.py index ca1d3b1ed8..49981be5c0 100644 --- a/mteb/tasks/Classification/__init__.py +++ b/mteb/tasks/Classification/__init__.py @@ -33,6 +33,7 @@ from .eng.NewsClassification import * from .eng.PatentClassification import * from .eng.PoemSentimentClassification import * +from .eng.RuSciBenchClassification import * from .eng.SDSEyeProtectionClassification import * from .eng.SDSGlovesClassification import * from .eng.ToxicChatClassification import * diff --git a/mteb/tasks/Classification/multilingual/RuSciBenchClassification.py b/mteb/tasks/Classification/multilingual/RuSciBenchClassification.py index f7f778de47..5855620dc2 100644 --- a/mteb/tasks/Classification/multilingual/RuSciBenchClassification.py +++ b/mteb/tasks/Classification/multilingual/RuSciBenchClassification.py @@ -12,7 +12,7 @@ class RuSciBenchCoreRiscClassification(MultilingualTask, AbsTaskClassification): "path": "mlsa-iai-msu-lab/ru_sci_bench_mteb", "revision": "fbc0599a0b5f00b3c7d87ab4d13490f04fb77f8e", }, - description="Classification of scientific papers (title+abstract) by publication Core RISC status", + description="This binary classification task aims to determine whether a scientific paper (based on its title and abstract) belongs to the Core of the Russian Science Citation Index (RISC). The RISC includes a wide range of publications, but the Core RISC comprises the most cited and prestigious journals, dissertations, theses, monographs, and studies. The task is provided for both Russian and English versions of the paper's title and abstract.", reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", type="Classification", category="p2p", @@ -30,20 +30,21 @@ class RuSciBenchCoreRiscClassification(MultilingualTask, AbsTaskClassification): annotations_creators="derived", dialect=[], sample_creation="found", - bibtex_citation=""" + bibtex_citation=r""" @article{vatolin2024ruscibench, - author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, - title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, + doi = {10.1134/S1064562424602191}, + issn = {1531-8362}, journal = {Doklady Mathematics}, - year = {2024}, - volume = {110}, - number = {1}, - pages = {S251--S260}, - month = {12}, - doi = {10.1134/S1064562424602191}, - url = {https://doi.org/10.1134/S1064562424602191}, - issn = {1531-8362} -}""", + month = {12}, + number = {1}, + pages = {S251--S260}, + title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + url = {https://doi.org/10.1134/S1064562424602191}, + volume = {110}, + year = {2024}, +} +""", prompt="Classify whether a scientific article is part of the core RISC or not based on the title and abstract", ) @@ -55,7 +56,7 @@ class RuSciBenchPubTypeClassification(MultilingualTask, AbsTaskClassification): "path": "mlsa-iai-msu-lab/ru_sci_bench_mteb", "revision": "fbc0599a0b5f00b3c7d87ab4d13490f04fb77f8e", }, - description="Classification of scientific papers (title+abstract) by publication type", + description="This task involves classifying scientific papers (based on their title and abstract) into different publication types. The dataset identifies the following types: 'Article', 'Conference proceedings', 'Survey', 'Miscellanea', 'Short message', 'Review', and 'Personalia'. This task is available for both Russian and English versions of the paper's title and abstract.", reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", type="Classification", category="p2p", @@ -73,19 +74,111 @@ class RuSciBenchPubTypeClassification(MultilingualTask, AbsTaskClassification): annotations_creators="derived", dialect=[], sample_creation="found", - bibtex_citation=""" + bibtex_citation=r""" @article{vatolin2024ruscibench, - author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, - title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, + doi = {10.1134/S1064562424602191}, + issn = {1531-8362}, journal = {Doklady Mathematics}, - year = {2024}, - volume = {110}, - number = {1}, - pages = {S251--S260}, - month = {12}, - doi = {10.1134/S1064562424602191}, - url = {https://doi.org/10.1134/S1064562424602191}, - issn = {1531-8362} -}""", + month = {12}, + number = {1}, + pages = {S251--S260}, + title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + url = {https://doi.org/10.1134/S1064562424602191}, + volume = {110}, + year = {2024}, +} +""", prompt="Classify the type of scientific paper based on the title and abstract", ) + + +class RuSciBenchGRNTIOrigClassification(MultilingualTask, AbsTaskClassification): + metadata = TaskMetadata( + name="RuSciBenchGRNTIOrigClassification", + dataset={ + "path": "mlsa-iai-msu-lab/ru_sci_bench_mteb", + "revision": "fbc0599a0b5f00b3c7d87ab4d13490f04fb77f8e", + }, + description="""Classification of scientific papers based on the GRNTI (State Rubricator of Scientific and + Technical Information) rubricator. GRNTI is a universal hierarchical classification of knowledge domains + adopted in Russia and CIS countries to systematize the entire flow of scientific and technical information. + This task uses the first level of the GRNTI hierarchy and top 28 classes by frequency.""", + reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", + type="Classification", + category="p2p", + modalities=["text"], + eval_splits=["test"], + eval_langs={ + "grnti_ru": ["rus-Cyrl"], + "grnti_en": ["eng-Latn"], + }, + main_score="accuracy", + date=("2007-01-01", "2023-01-01"), + domains=["Academic", "Non-fiction", "Written"], + task_subtypes=[], + license="mit", + annotations_creators="derived", + dialect=[], + sample_creation="found", + bibtex_citation=r""" +@article{vatolin2024ruscibench, + author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, + doi = {10.1134/S1064562424602191}, + issn = {1531-8362}, + journal = {Doklady Mathematics}, + month = {12}, + number = {1}, + pages = {S251--S260}, + title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + url = {https://doi.org/10.1134/S1064562424602191}, + volume = {110}, + year = {2024}, +} +""", + prompt="Classify the category of scientific papers based on the titles and abstracts", + ) + + +class RuSciBenchOECDOrigClassification(MultilingualTask, AbsTaskClassification): + metadata = TaskMetadata( + name="RuSciBenchOECDOrigClassification", + dataset={ + "path": "mlsa-iai-msu-lab/ru_sci_bench_mteb", + "revision": "fbc0599a0b5f00b3c7d87ab4d13490f04fb77f8e", + }, + description="Classification of scientific papers based on the OECD (Organization for Economic Co-operation and Development) rubricator. OECD provides a hierarchical 3-level system of classes for labeling scientific articles. This task uses the first two levels of the OECD hierarchy, top 29 classes.", + reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", + type="Classification", + category="p2p", + modalities=["text"], + eval_splits=["test"], + eval_langs={ + "oecd_ru": ["rus-Cyrl"], + "oecd_en": ["eng-Latn"], + }, + main_score="accuracy", + date=("2007-01-01", "2023-01-01"), + domains=["Academic", "Non-fiction", "Written"], + task_subtypes=[], + license="mit", + annotations_creators="derived", + dialect=[], + sample_creation="found", + bibtex_citation=r""" +@article{vatolin2024ruscibench, + author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, + doi = {10.1134/S1064562424602191}, + issn = {1531-8362}, + journal = {Doklady Mathematics}, + month = {12}, + number = {1}, + pages = {S251--S260}, + title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + url = {https://doi.org/10.1134/S1064562424602191}, + volume = {110}, + year = {2024}, +} +""", + prompt="Classify the category of scientific papers based on the titles and abstracts", + ) From 611e7ff4e97763c4faf991adfde7884a53b49995 Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Sun, 13 Jul 2025 17:57:50 +0000 Subject: [PATCH 04/17] fix init --- mteb/tasks/Classification/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/mteb/tasks/Classification/__init__.py b/mteb/tasks/Classification/__init__.py index 49981be5c0..ca1d3b1ed8 100644 --- a/mteb/tasks/Classification/__init__.py +++ b/mteb/tasks/Classification/__init__.py @@ -33,7 +33,6 @@ from .eng.NewsClassification import * from .eng.PatentClassification import * from .eng.PoemSentimentClassification import * -from .eng.RuSciBenchClassification import * from .eng.SDSEyeProtectionClassification import * from .eng.SDSGlovesClassification import * from .eng.ToxicChatClassification import * From 67136efe80c0d2ff92e2941e361028000cfad3ca Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Sun, 13 Jul 2025 18:14:01 +0000 Subject: [PATCH 05/17] add missing files --- mteb/tasks/Regression/__init__.py | 3 + .../multilingual/RuSciBenchRegression.py | 105 ++++++++++++++++ .../tasks/Regression/multilingual/__init__.py | 0 .../multilingual/RuSciBenchRetrieval.py | 118 +++++++++++++----- mteb/tasks/__init__.py | 1 + 5 files changed, 198 insertions(+), 29 deletions(-) create mode 100644 mteb/tasks/Regression/__init__.py create mode 100644 mteb/tasks/Regression/multilingual/RuSciBenchRegression.py create mode 100644 mteb/tasks/Regression/multilingual/__init__.py diff --git a/mteb/tasks/Regression/__init__.py b/mteb/tasks/Regression/__init__.py new file mode 100644 index 0000000000..ec32ed84c1 --- /dev/null +++ b/mteb/tasks/Regression/__init__.py @@ -0,0 +1,3 @@ +from __future__ import annotations + +from .multilingual.RuSciBenchRegression import * diff --git a/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py b/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py new file mode 100644 index 0000000000..46b9b11528 --- /dev/null +++ b/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +from mteb.abstasks.AbsTaskRegression import AbsTaskRegression +from mteb.abstasks.MultilingualTask import MultilingualTask +from mteb.abstasks.TaskMetadata import TaskMetadata + + +class RuSciBenchCitedCountRegression(MultilingualTask, AbsTaskRegression): + metadata = TaskMetadata( + name="RuSciBenchCitedCountRegression", + description="Predicts the number of times a scientific article has been cited by other papers. The prediction is based on the article's title and abstract. The data is sourced from the Russian electronic library of scientific publications (eLibrary.ru) and includes papers with both Russian and English abstracts.", + reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", + dataset={ + "path": "mlsa-iai-msu-lab/ru_sci_bench_mteb", + "revision": "fbc0599a0b5f00b3c7d87ab4d13490f04fb77f8e", + }, + type="Regression", + category="p2p", + modalities=["text"], + eval_splits=["test"], + eval_langs={ + "cited_count_ru": ["rus-Cyrl"], + "cited_count_en": ["eng-Latn"], + }, + main_score="kendalltau", + date=("2007-01-01", "2023-01-01"), + domains=["Academic", "Non-fiction", "Written"], + task_subtypes=[], + license="mit", + sample_creation="found", + annotations_creators="derived", + dialect=None, + bibtex_citation=r""" +@article{vatolin2024ruscibench, + author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, + doi = {10.1134/S1064562424602191}, + issn = {1531-8362}, + journal = {Doklady Mathematics}, + month = {12}, + number = {1}, + pages = {S251--S260}, + title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + url = {https://doi.org/10.1134/S1064562424602191}, + volume = {110}, + year = {2024}, +} +""", + prompt="Predict the number of citations for a scientific article based on the title and abstract", + ) + + def dataset_transform(self): + for subset in self.dataset: + self.dataset[subset]["train"] = self.dataset[subset][ + "train" + ].train_test_split(test_size=2048, seed=self.seed)["test"] + + +class RuSciBenchYearPublRegression(MultilingualTask, AbsTaskRegression): + metadata = TaskMetadata( + name="RuSciBenchYearPublRegression", + description="Predicts the publication year of a scientific article. The prediction is based on the article's title and abstract. The data is sourced from the Russian electronic library of scientific publications (eLibrary.ru) and includes papers with both Russian and English abstracts.", + reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", + dataset={ + "path": "mlsa-iai-msu-lab/ru_sci_bench_mteb", + "revision": "fbc0599a0b5f00b3c7d87ab4d13490f04fb77f8e", + }, + type="Regression", + category="p2p", + modalities=["text"], + eval_splits=["test"], + eval_langs={ + "yearpubl_ru": ["rus-Cyrl"], + "yearpubl_en": ["eng-Latn"], + }, + main_score="kendalltau", + date=("2007-01-01", "2023-01-01"), + domains=["Academic", "Non-fiction", "Written"], + task_subtypes=[], + license="mit", + sample_creation="found", + annotations_creators="derived", + dialect=None, + bibtex_citation=r""" +@article{vatolin2024ruscibench, + author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, + doi = {10.1134/S1064562424602191}, + issn = {1531-8362}, + journal = {Doklady Mathematics}, + month = {12}, + number = {1}, + pages = {S251--S260}, + title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + url = {https://doi.org/10.1134/S1064562424602191}, + volume = {110}, + year = {2024}, +} +""", + prompt="Predict paper publitaction year based on the title and abstract", + ) + + def dataset_transform(self): + for subset in self.dataset: + self.dataset[subset]["train"] = self.dataset[subset][ + "train" + ].train_test_split(test_size=2048, seed=self.seed)["test"] diff --git a/mteb/tasks/Regression/multilingual/__init__.py b/mteb/tasks/Regression/multilingual/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/mteb/tasks/Retrieval/multilingual/RuSciBenchRetrieval.py b/mteb/tasks/Retrieval/multilingual/RuSciBenchRetrieval.py index 94b639d2b7..12fa7b2ec6 100644 --- a/mteb/tasks/Retrieval/multilingual/RuSciBenchRetrieval.py +++ b/mteb/tasks/Retrieval/multilingual/RuSciBenchRetrieval.py @@ -1,9 +1,13 @@ from __future__ import annotations +from collections import defaultdict +from typing import cast + +import datasets + from mteb.abstasks.AbsTaskRetrieval import AbsTaskRetrieval from mteb.abstasks.MultilingualTask import MultilingualTask from mteb.abstasks.TaskMetadata import TaskMetadata -from mteb.tasks.Retrieval.multilingual.NeuCLIR2023Retrieval import load_neuclir_data _LANGUAGES = { "ru": ["rus-Cyrl"], @@ -11,6 +15,60 @@ } +def load_ruscibench_data( + path: str, + langs: list, + eval_splits: list, + cache_dir: str | None = None, + revision: str | None = None, +): + corpus: dict[str, dict[str, dict[str, dict[str, str]] | None]] = { + lang: dict.fromkeys(eval_splits) for lang in langs + } + queries: dict[str, dict[str, dict[str, str] | None]] = { + lang: dict.fromkeys(eval_splits) for lang in langs + } + relevant_docs: dict[str, dict[str, dict[str, dict[str, int]] | None]] = { + lang: dict.fromkeys(eval_splits) for lang in langs + } + + for lang in langs: + lang_corpus = cast( + datasets.Dataset, + datasets.load_dataset( + path, f"corpus-{lang}", cache_dir=cache_dir, revision=revision + ), + )["corpus"] + lang_queries = cast( + datasets.Dataset, + datasets.load_dataset( + path, f"queries-{lang}", cache_dir=cache_dir, revision=revision + ), + )["queries"] + lang_qrels = cast( + datasets.Dataset, + datasets.load_dataset( + path, f"{lang}", cache_dir=cache_dir, revision=revision + ), + )["test"] + corpus[lang] = { + "test": { + str(e["_id"]): {"text": e["text"], "title": e["title"]} + for e in lang_corpus + } + } + queries[lang] = {"test": {str(e["_id"]): e["text"] for e in lang_queries}} + relevant_docs[lang]["test"] = defaultdict(dict) + for item in lang_qrels: + relevant_docs[lang]["test"][str(item["query-id"])].update( + {str(item["corpus-id"]): item["score"]} + ) + corpus = datasets.DatasetDict(corpus) + queries = datasets.DatasetDict(queries) + relevant_docs = datasets.DatasetDict(relevant_docs) + return corpus, queries, relevant_docs + + class RuSciBenchCiteRetrieval(AbsTaskRetrieval, MultilingualTask): ignore_identical_ids = True @@ -20,7 +78,7 @@ class RuSciBenchCiteRetrieval(AbsTaskRetrieval, MultilingualTask): "path": "mlsa-iai-msu-lab/ru_sci_bench_cite_retrieval", "revision": "6cb447d02f41b8b775d5d9df7faf472f44d2f1db", }, - description="Retrieval of related scientific papers based on their title and abstract", + description="This task is focused on Direct Citation Prediction for scientific papers. Given a query paper (title and abstract), the goal is to retrieve papers that are directly cited by it from a larger corpus of papers. The dataset for this task consists of 3,000 query papers, 15,000 relevant (cited) papers, and 75,000 irrelevant papers. The task is available for both Russian and English scientific texts.", reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", type="Retrieval", category="p2p", @@ -35,20 +93,21 @@ class RuSciBenchCiteRetrieval(AbsTaskRetrieval, MultilingualTask): dialect=[], sample_creation="found", annotations_creators="derived", - bibtex_citation=""" + bibtex_citation=r""" @article{vatolin2024ruscibench, - author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, - title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, + doi = {10.1134/S1064562424602191}, + issn = {1531-8362}, journal = {Doklady Mathematics}, - year = {2024}, - volume = {110}, - number = {1}, - pages = {S251--S260}, - month = {12}, - doi = {10.1134/S1064562424602191}, - url = {https://doi.org/10.1134/S1064562424602191}, - issn = {1531-8362} -}""", + month = {12}, + number = {1}, + pages = {S251--S260}, + title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + url = {https://doi.org/10.1134/S1064562424602191}, + volume = {110}, + year = {2024}, +} +""", prompt={ "query": "Given a title and abstract of a scientific paper, retrieve the titles and abstracts of other relevant papers", }, @@ -58,7 +117,7 @@ def load_data(self, **kwargs): if self.data_loaded: return - self.corpus, self.queries, self.relevant_docs = load_neuclir_data( + self.corpus, self.queries, self.relevant_docs = load_ruscibench_data( path=self.metadata_dict["dataset"]["path"], langs=self.metadata.eval_langs, eval_splits=self.metadata_dict["eval_splits"], @@ -78,7 +137,7 @@ class RuSciBenchCociteRetrieval(MultilingualTask, AbsTaskRetrieval): "path": "mlsa-iai-msu-lab/ru_sci_bench_cocite_retrieval", "revision": "a5da47a245275669d2b6ddf8f96c5338dd2428b4", }, - description="Retrieval of related scientific papers based on their title and abstract", + description="This task focuses on Co-citation Prediction for scientific papers. Given a query paper (title and abstract), the goal is to retrieve other papers that are co-cited with it. Two papers are considered co-cited if they are both cited by at least 5 of the same other papers. Similar to the Direct Citation task, this task employs a retrieval setup: for a given query paper, all other papers in the corpus that are not co-cited with it are considered negative examples. The task is available for both Russian and English scientific texts.", reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", type="Retrieval", category="p2p", @@ -93,20 +152,21 @@ class RuSciBenchCociteRetrieval(MultilingualTask, AbsTaskRetrieval): dialect=[], sample_creation="found", annotations_creators="derived", - bibtex_citation=""" + bibtex_citation=r""" @article{vatolin2024ruscibench, - author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, - title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, + doi = {10.1134/S1064562424602191}, + issn = {1531-8362}, journal = {Doklady Mathematics}, - year = {2024}, - volume = {110}, - number = {1}, - pages = {S251--S260}, - month = {12}, - doi = {10.1134/S1064562424602191}, - url = {https://doi.org/10.1134/S1064562424602191}, - issn = {1531-8362} -}""", + month = {12}, + number = {1}, + pages = {S251--S260}, + title = {RuSciBench: Open Benchmark for Russian and English Scientific Document Representations}, + url = {https://doi.org/10.1134/S1064562424602191}, + volume = {110}, + year = {2024}, +} +""", prompt={ "query": "Given a title and abstract of a scientific paper, retrieve the titles and abstracts of other relevant papers", }, @@ -116,7 +176,7 @@ def load_data(self, **kwargs): if self.data_loaded: return - self.corpus, self.queries, self.relevant_docs = load_neuclir_data( + self.corpus, self.queries, self.relevant_docs = load_ruscibench_data( path=self.metadata_dict["dataset"]["path"], langs=self.metadata.eval_langs, eval_splits=self.metadata_dict["eval_splits"], diff --git a/mteb/tasks/__init__.py b/mteb/tasks/__init__.py index 8abdf1f811..5304bd961d 100644 --- a/mteb/tasks/__init__.py +++ b/mteb/tasks/__init__.py @@ -15,6 +15,7 @@ from .InstructionRetrieval import * from .MultiLabelClassification import * from .PairClassification import * +from .Regression import * from .Reranking import * from .Retrieval import * from .SpeedTask import * From 42c68cc10905d8a39b9fdde0a86c7fec7c4f8d2b Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Sun, 13 Jul 2025 18:46:10 +0000 Subject: [PATCH 06/17] Improve description --- .../multilingual/RuSciBenchBitextMining.py | 21 +++++++++---------- .../multilingual/RuSciBenchClassification.py | 16 +++++++++++--- .../multilingual/RuSciBenchRegression.py | 12 +++++++---- .../multilingual/RuSciBenchRetrieval.py | 14 +++++++++++-- 4 files changed, 43 insertions(+), 20 deletions(-) diff --git a/mteb/tasks/BitextMining/multilingual/RuSciBenchBitextMining.py b/mteb/tasks/BitextMining/multilingual/RuSciBenchBitextMining.py index 2e91e935b3..f696b4abf0 100644 --- a/mteb/tasks/BitextMining/multilingual/RuSciBenchBitextMining.py +++ b/mteb/tasks/BitextMining/multilingual/RuSciBenchBitextMining.py @@ -4,14 +4,6 @@ from mteb.abstasks.MultilingualTask import MultilingualTask from mteb.abstasks.TaskMetadata import TaskMetadata -_LANGUAGES = { - "ru-en": ["rus-Cyrl", "eng-Latn"], - "en-ru": ["eng-Latn", "rus-Cyrl"], -} - - -_SPLITS = ["test"] - class RuSciBenchBitextMining(AbsTaskBitextMining, MultilingualTask): fast_loading = True @@ -21,13 +13,20 @@ class RuSciBenchBitextMining(AbsTaskBitextMining, MultilingualTask): "path": "mlsa-iai-msu-lab/ru_sci_bench_bitext_mining", "revision": "e5840033c5cf2573932db027ac8001fe0a7eb6fa", }, - description="This task focuses on finding translations of scientific articles. The dataset is sourced from eLibrary, Russia's largest electronic library of scientific publications. Russian authors often provide English translations for their abstracts and titles, and the data consists of these paired titles and abstracts. The task evaluates a model's ability to match an article's Russian title and abstract to its English counterpart, or vice versa.", + description="""This task focuses on finding translations of scientific articles. + The dataset is sourced from eLibrary, Russia's largest electronic library of scientific publications. + Russian authors often provide English translations for their abstracts and titles, + and the data consists of these paired titles and abstracts. The task evaluates a model's ability + to match an article's Russian title and abstract to its English counterpart, or vice versa.""", reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", type="BitextMining", category="p2p", modalities=["text"], - eval_splits=_SPLITS, - eval_langs=_LANGUAGES, + eval_splits=["test"], + eval_langs={ + "ru-en": ["rus-Cyrl", "eng-Latn"], + "en-ru": ["eng-Latn", "rus-Cyrl"], + }, main_score="f1", date=("2007-01-01", "2023-01-01"), domains=["Academic", "Non-fiction", "Written"], diff --git a/mteb/tasks/Classification/multilingual/RuSciBenchClassification.py b/mteb/tasks/Classification/multilingual/RuSciBenchClassification.py index 5855620dc2..43ee567170 100644 --- a/mteb/tasks/Classification/multilingual/RuSciBenchClassification.py +++ b/mteb/tasks/Classification/multilingual/RuSciBenchClassification.py @@ -12,7 +12,11 @@ class RuSciBenchCoreRiscClassification(MultilingualTask, AbsTaskClassification): "path": "mlsa-iai-msu-lab/ru_sci_bench_mteb", "revision": "fbc0599a0b5f00b3c7d87ab4d13490f04fb77f8e", }, - description="This binary classification task aims to determine whether a scientific paper (based on its title and abstract) belongs to the Core of the Russian Science Citation Index (RISC). The RISC includes a wide range of publications, but the Core RISC comprises the most cited and prestigious journals, dissertations, theses, monographs, and studies. The task is provided for both Russian and English versions of the paper's title and abstract.", + description="""This binary classification task aims to determine whether a scientific paper + (based on its title and abstract) belongs to the Core of the Russian Science Citation Index (RISC). + The RISC includes a wide range of publications, but the Core RISC comprises the most cited and prestigious + journals, dissertations, theses, monographs, and studies. The task is provided for both Russian and English + versions of the paper's title and abstract.""", reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", type="Classification", category="p2p", @@ -56,7 +60,10 @@ class RuSciBenchPubTypeClassification(MultilingualTask, AbsTaskClassification): "path": "mlsa-iai-msu-lab/ru_sci_bench_mteb", "revision": "fbc0599a0b5f00b3c7d87ab4d13490f04fb77f8e", }, - description="This task involves classifying scientific papers (based on their title and abstract) into different publication types. The dataset identifies the following types: 'Article', 'Conference proceedings', 'Survey', 'Miscellanea', 'Short message', 'Review', and 'Personalia'. This task is available for both Russian and English versions of the paper's title and abstract.", + description="""This task involves classifying scientific papers (based on their title and abstract) + into different publication types. The dataset identifies the following types: + 'Article', 'Conference proceedings', 'Survey', 'Miscellanea', 'Short message', 'Review', and 'Personalia'. + This task is available for both Russian and English versions of the paper's title and abstract.""", reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", type="Classification", category="p2p", @@ -147,7 +154,10 @@ class RuSciBenchOECDOrigClassification(MultilingualTask, AbsTaskClassification): "path": "mlsa-iai-msu-lab/ru_sci_bench_mteb", "revision": "fbc0599a0b5f00b3c7d87ab4d13490f04fb77f8e", }, - description="Classification of scientific papers based on the OECD (Organization for Economic Co-operation and Development) rubricator. OECD provides a hierarchical 3-level system of classes for labeling scientific articles. This task uses the first two levels of the OECD hierarchy, top 29 classes.", + description="""Classification of scientific papers based on the OECD + (Organization for Economic Co-operation and Development) rubricator. OECD provides + a hierarchical 3-level system of classes for labeling scientific articles. + This task uses the first two levels of the OECD hierarchy, top 29 classes.""", reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", type="Classification", category="p2p", diff --git a/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py b/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py index 46b9b11528..476c0afe5c 100644 --- a/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py +++ b/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py @@ -8,7 +8,9 @@ class RuSciBenchCitedCountRegression(MultilingualTask, AbsTaskRegression): metadata = TaskMetadata( name="RuSciBenchCitedCountRegression", - description="Predicts the number of times a scientific article has been cited by other papers. The prediction is based on the article's title and abstract. The data is sourced from the Russian electronic library of scientific publications (eLibrary.ru) and includes papers with both Russian and English abstracts.", + description="""Predicts the number of times a scientific article has been cited by other papers. + The prediction is based on the article's title and abstract. The data is sourced from the Russian electronic + library of scientific publications (eLibrary.ru) and includes papers with both Russian and English abstracts.""", reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", dataset={ "path": "mlsa-iai-msu-lab/ru_sci_bench_mteb", @@ -29,7 +31,7 @@ class RuSciBenchCitedCountRegression(MultilingualTask, AbsTaskRegression): license="mit", sample_creation="found", annotations_creators="derived", - dialect=None, + dialect=[], bibtex_citation=r""" @article{vatolin2024ruscibench, author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, @@ -58,7 +60,9 @@ def dataset_transform(self): class RuSciBenchYearPublRegression(MultilingualTask, AbsTaskRegression): metadata = TaskMetadata( name="RuSciBenchYearPublRegression", - description="Predicts the publication year of a scientific article. The prediction is based on the article's title and abstract. The data is sourced from the Russian electronic library of scientific publications (eLibrary.ru) and includes papers with both Russian and English abstracts.", + description="""Predicts the publication year of a scientific article. The prediction is based on the + article's title and abstract. The data is sourced from the Russian electronic library of scientific + publications (eLibrary.ru) and includes papers with both Russian and English abstracts.""", reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", dataset={ "path": "mlsa-iai-msu-lab/ru_sci_bench_mteb", @@ -79,7 +83,7 @@ class RuSciBenchYearPublRegression(MultilingualTask, AbsTaskRegression): license="mit", sample_creation="found", annotations_creators="derived", - dialect=None, + dialect=[], bibtex_citation=r""" @article{vatolin2024ruscibench, author = {Vatolin, A. and Gerasimenko, N. and Ianina, A. and Vorontsov, K.}, diff --git a/mteb/tasks/Retrieval/multilingual/RuSciBenchRetrieval.py b/mteb/tasks/Retrieval/multilingual/RuSciBenchRetrieval.py index 12fa7b2ec6..d70f7527a5 100644 --- a/mteb/tasks/Retrieval/multilingual/RuSciBenchRetrieval.py +++ b/mteb/tasks/Retrieval/multilingual/RuSciBenchRetrieval.py @@ -78,7 +78,11 @@ class RuSciBenchCiteRetrieval(AbsTaskRetrieval, MultilingualTask): "path": "mlsa-iai-msu-lab/ru_sci_bench_cite_retrieval", "revision": "6cb447d02f41b8b775d5d9df7faf472f44d2f1db", }, - description="This task is focused on Direct Citation Prediction for scientific papers. Given a query paper (title and abstract), the goal is to retrieve papers that are directly cited by it from a larger corpus of papers. The dataset for this task consists of 3,000 query papers, 15,000 relevant (cited) papers, and 75,000 irrelevant papers. The task is available for both Russian and English scientific texts.", + description="""This task is focused on Direct Citation Prediction for scientific papers from eLibrary, + Russia's largest electronic library of scientific publications. Given a query paper (title and abstract), + the goal is to retrieve papers that are directly cited by it from a larger corpus of papers. + The dataset for this task consists of 3,000 query papers, 15,000 relevant (cited) papers, + and 75,000 irrelevant papers. The task is available for both Russian and English scientific texts.""", reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", type="Retrieval", category="p2p", @@ -137,7 +141,13 @@ class RuSciBenchCociteRetrieval(MultilingualTask, AbsTaskRetrieval): "path": "mlsa-iai-msu-lab/ru_sci_bench_cocite_retrieval", "revision": "a5da47a245275669d2b6ddf8f96c5338dd2428b4", }, - description="This task focuses on Co-citation Prediction for scientific papers. Given a query paper (title and abstract), the goal is to retrieve other papers that are co-cited with it. Two papers are considered co-cited if they are both cited by at least 5 of the same other papers. Similar to the Direct Citation task, this task employs a retrieval setup: for a given query paper, all other papers in the corpus that are not co-cited with it are considered negative examples. The task is available for both Russian and English scientific texts.", + description="""This task focuses on Co-citation Prediction for scientific papers from eLibrary, + Russia's largest electronic library of scientific publications. Given a query paper (title and abstract), + the goal is to retrieve other papers that are co-cited with it. Two papers are considered co-cited + if they are both cited by at least 5 of the same other papers. Similar to the Direct Citation task, + this task employs a retrieval setup: for a given query paper, all other papers in the corpus that + are not co-cited with it are considered negative examples. The task is available for both Russian + and English scientific texts.""", reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", type="Retrieval", category="p2p", From 100f64f9d05e22ec167a28584083ab2272a9edf8 Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Mon, 14 Jul 2025 11:48:30 +0000 Subject: [PATCH 07/17] Add superseded_by --- mteb/tasks/Classification/rus/RuSciBenchGRNTIClassification.py | 2 ++ mteb/tasks/Classification/rus/RuSciBenchOECDClassification.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/mteb/tasks/Classification/rus/RuSciBenchGRNTIClassification.py b/mteb/tasks/Classification/rus/RuSciBenchGRNTIClassification.py index c2c737eea5..7c87a98aaf 100644 --- a/mteb/tasks/Classification/rus/RuSciBenchGRNTIClassification.py +++ b/mteb/tasks/Classification/rus/RuSciBenchGRNTIClassification.py @@ -5,6 +5,8 @@ class RuSciBenchGRNTIClassification(AbsTaskClassification): + superseded_by="RuSciBenchGRNTIOrigClassification" + metadata = TaskMetadata( name="RuSciBenchGRNTIClassification", dataset={ diff --git a/mteb/tasks/Classification/rus/RuSciBenchOECDClassification.py b/mteb/tasks/Classification/rus/RuSciBenchOECDClassification.py index b32f7c2b6e..3b7062ae6e 100644 --- a/mteb/tasks/Classification/rus/RuSciBenchOECDClassification.py +++ b/mteb/tasks/Classification/rus/RuSciBenchOECDClassification.py @@ -5,6 +5,8 @@ class RuSciBenchOECDClassification(AbsTaskClassification): + superseded_by="RuSciBenchOECDOrigClassification" + metadata = TaskMetadata( name="RuSciBenchOECDClassification", dataset={ From 43a8ceb93288f31f046d2f88a07b5d1655c68c3f Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Mon, 14 Jul 2025 11:59:01 +0000 Subject: [PATCH 08/17] fix lint --- mteb/tasks/Classification/rus/RuSciBenchGRNTIClassification.py | 2 +- mteb/tasks/Classification/rus/RuSciBenchOECDClassification.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/mteb/tasks/Classification/rus/RuSciBenchGRNTIClassification.py b/mteb/tasks/Classification/rus/RuSciBenchGRNTIClassification.py index 7c87a98aaf..69dcbc06bd 100644 --- a/mteb/tasks/Classification/rus/RuSciBenchGRNTIClassification.py +++ b/mteb/tasks/Classification/rus/RuSciBenchGRNTIClassification.py @@ -5,7 +5,7 @@ class RuSciBenchGRNTIClassification(AbsTaskClassification): - superseded_by="RuSciBenchGRNTIOrigClassification" + superseded_by = "RuSciBenchGRNTIOrigClassification" metadata = TaskMetadata( name="RuSciBenchGRNTIClassification", diff --git a/mteb/tasks/Classification/rus/RuSciBenchOECDClassification.py b/mteb/tasks/Classification/rus/RuSciBenchOECDClassification.py index 3b7062ae6e..ec1b7ba5ef 100644 --- a/mteb/tasks/Classification/rus/RuSciBenchOECDClassification.py +++ b/mteb/tasks/Classification/rus/RuSciBenchOECDClassification.py @@ -5,7 +5,7 @@ class RuSciBenchOECDClassification(AbsTaskClassification): - superseded_by="RuSciBenchOECDOrigClassification" + superseded_by = "RuSciBenchOECDOrigClassification" metadata = TaskMetadata( name="RuSciBenchOECDClassification", From fff7575b01b783dfa66331be14d3f39555d7eed9 Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Mon, 14 Jul 2025 14:26:36 +0000 Subject: [PATCH 09/17] Update regression task to match with v2 --- mteb/abstasks/AbsTaskRegression.py | 63 +++++++++++-------- .../evaluators/RegressionEvaluator.py | 24 ++++--- 2 files changed, 50 insertions(+), 37 deletions(-) diff --git a/mteb/abstasks/AbsTaskRegression.py b/mteb/abstasks/AbsTaskRegression.py index 097fd66ab4..71f37f9dd0 100644 --- a/mteb/abstasks/AbsTaskRegression.py +++ b/mteb/abstasks/AbsTaskRegression.py @@ -3,6 +3,10 @@ import logging from typing import Any +from datasets import DatasetDict +from sklearn.base import BaseEstimator +from sklearn.linear_model import LinearRegression + from mteb.abstasks.TaskMetadata import DescriptiveStatistics, HFSubset from mteb.encoder_interface import Encoder from mteb.evaluation.evaluators.RegressionEvaluator import LinearRegressionEvaluator @@ -54,33 +58,36 @@ class AbsTaskRegression(AbsTask): value: float """ - def __init__(self, seed: int = 42, **kwargs: Any): - super().__init__(seed, **kwargs) - if hasattr(self, "metadata"): - self.metadata + evaluator: type[LinearRegressionEvaluator] = LinearRegressionEvaluator + model: BaseEstimator = LinearRegression(n_jobs=-1) + + train_split: str = "train" + label_column_name: str = "value" + input_column_name: str = "text" + abstask_prompt = "Predict the value of the user passage." def _evaluate_subset( self, model: Encoder, - dataset, - eval_split: str = "test", - train_split: str = "train", - encode_kwargs: dict[str, Any] = {}, + dataset: DatasetDict, + hf_split: str, + hf_subset: str, + encode_kwargs: dict[str, Any], **kwargs, ) -> ScoresDict: - train_split = dataset[train_split] - eval_split = dataset[eval_split] - - evaluator = LinearRegressionEvaluator( - train_split["text"], - train_split["value"], - eval_split["text"], - eval_split["value"], + train_split = dataset[self.train_split] + eval_split = dataset[hf_split] + + evaluator = self.evaluator( + train_split[self.input_column_name], + train_split[self.label_column_name], + eval_split[self.input_column_name], + eval_split[self.label_column_name], task_name=self.metadata.name, - encode_kwargs=encode_kwargs, + hf_subset=hf_subset, **kwargs, ) - scores = evaluator(model) + scores = evaluator(model, encode_kwargs=encode_kwargs) return scores def _add_main_score(self, scores): @@ -89,8 +96,8 @@ def _add_main_score(self, scores): def evaluate( self, model: Encoder, - eval_split: str = "test", - train_split: str = "train", + split: str = "test", + subsets_to_run: list[HFSubset] | None = None, *, encode_kwargs: dict[str, Any] = {}, **kwargs: Any, @@ -98,12 +105,17 @@ def evaluate( if not self.data_loaded: self.load_data() + if "random_state" in self.classifier.get_params(): + self.classifier = self.classifier.set_params(random_state=self.seed) + scores = {} - hf_subsets = list(self.dataset) if self.is_multilingual else ["default"] + hf_subsets = self.hf_subsets + if subsets_to_run is not None: + hf_subsets = [s for s in hf_subsets if s in subsets_to_run] for hf_subset in hf_subsets: logger.info( - f"\nTask: {self.metadata.name}, split: {eval_split}, subset: {hf_subset}. Running..." + f"\nTask: {self.metadata.name}, split: {split}, subset: {hf_subset}. Running..." ) if hf_subset not in self.dataset and hf_subset == "default": @@ -113,8 +125,8 @@ def evaluate( scores[hf_subset] = self._evaluate_subset( model, ds, - eval_split, - train_split, + hf_split=split, + hf_subset=hf_subset, encode_kwargs=encode_kwargs, **kwargs, ) @@ -122,9 +134,6 @@ def evaluate( return scores - def __hash__(self) -> int: - return hash(self.metadata) - def _calculate_metrics_from_split( self, split: str, hf_subset: str | None = None, compute_overall: bool = False ) -> RegressionDescriptiveStatistics: diff --git a/mteb/evaluation/evaluators/RegressionEvaluator.py b/mteb/evaluation/evaluators/RegressionEvaluator.py index 3c5cae128d..418250782a 100644 --- a/mteb/evaluation/evaluators/RegressionEvaluator.py +++ b/mteb/evaluation/evaluators/RegressionEvaluator.py @@ -22,9 +22,10 @@ def __init__( y_train, sentences_test, y_test, - task_name: str | None = None, + task_name: str, + hf_split: str, + hf_subset: str, fit_intercept: bool = True, - encode_kwargs: dict[str, Any] = {}, limit: int | None = None, **kwargs, ): @@ -39,27 +40,30 @@ def __init__( self.sentences_test = sentences_test self.y_test = y_test + self.hf_split = hf_split + self.hf_subset = hf_subset self.task_name = task_name - self.encode_kwargs = encode_kwargs - - if "batch_size" not in self.encode_kwargs: - self.encode_kwargs["batch_size"] = 32 - self.fit_intercept = fit_intercept - def __call__(self, model: Encoder) -> dict[str, float]: + def __call__( + self, model: Encoder, *, encode_kwargs: dict[str, Any] = {} + ) -> dict[str, float]: scores = {} X_train = model.encode( self.sentences_train, model=model, task_name=self.task_name, - **self.encode_kwargs, + hf_split="train", + hf_subset=self.hf_subset, + **encode_kwargs, ) X_test = model.encode( self.sentences_test, model=model, task_name=self.task_name, - **self.encode_kwargs, + hf_split=self.hf_split, + hf_subset=self.hf_subset, + **encode_kwargs, ) linear_regression = LinearRegression(fit_intercept=self.fit_intercept) From d2a11af659e5fdc4375a96deb6f3483b5eb6fe3d Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Tue, 15 Jul 2025 14:55:41 +0000 Subject: [PATCH 10/17] Add stratified_subsampling for regression task --- mteb/abstasks/AbsTaskRegression.py | 61 +++++++++++++++++-- .../multilingual/RuSciBenchRegression.py | 22 ++++--- 2 files changed, 69 insertions(+), 14 deletions(-) diff --git a/mteb/abstasks/AbsTaskRegression.py b/mteb/abstasks/AbsTaskRegression.py index 71f37f9dd0..01e577c1ce 100644 --- a/mteb/abstasks/AbsTaskRegression.py +++ b/mteb/abstasks/AbsTaskRegression.py @@ -3,16 +3,16 @@ import logging from typing import Any -from datasets import DatasetDict +import datasets +import pandas as pd from sklearn.base import BaseEstimator from sklearn.linear_model import LinearRegression +from mteb.abstasks.AbsTask import AbsTask from mteb.abstasks.TaskMetadata import DescriptiveStatistics, HFSubset from mteb.encoder_interface import Encoder from mteb.evaluation.evaluators.RegressionEvaluator import LinearRegressionEvaluator - -from ..load_results.task_results import ScoresDict -from .AbsTask import AbsTask +from mteb.load_results.task_results import ScoresDict logger = logging.getLogger(__name__) @@ -69,9 +69,10 @@ class AbsTaskRegression(AbsTask): def _evaluate_subset( self, model: Encoder, - dataset: DatasetDict, + dataset: datasets.DatasetDict, hf_split: str, - hf_subset: str, + *, + hf_subset: str | None = None, encode_kwargs: dict[str, Any], **kwargs, ) -> ScoresDict: @@ -134,6 +135,54 @@ def evaluate( return scores + @staticmethod + def stratified_subsampling( + dataset_dict: datasets.DatasetDict, + seed: int, + splits: list[str] = ["test"], + label: str = "value", + n_samples: int = 2048, + n_bins: int = 10, + ) -> datasets.DatasetDict: + """Subsamples the dataset with stratification by the supplied label, which is assumed to be a continuous value. + The continuous values are bucketized into `n_bins` bins based on quantiles. + Returns a DatasetDict object. + + Args: + dataset_dict: the DatasetDict object. + seed: the random seed. + splits: the splits of the dataset. + label: the label with which the stratified sampling is based on. + n_samples: Optional, number of samples to subsample. + n_bins: Optional, number of bins to bucketize the continuous label. + """ + stratify_col_name = f"{label}_binned_for_stratification" + + for split in splits: + if n_samples >= len(dataset_dict[split]): + logger.debug( + "Subsampling not needed for split %s, as n_samples is equal or greater than the number of samples.", + split, + ) + continue + + dataset = dataset_dict[split] + labels = dataset[label] + + binned_labels = pd.qcut(labels, q=n_bins, labels=False, duplicates="drop") + dataset_with_bins: datasets.Dataset = dataset.add_column( + name=stratify_col_name, column=binned_labels.tolist() + ) + + subsampled_dataset = dataset_with_bins.train_test_split( + test_size=n_samples, seed=seed, stratify_by_column=stratify_col_name + )["test"] + + subsampled_dataset = subsampled_dataset.remove_columns([stratify_col_name]) + dataset_dict[split] = subsampled_dataset + + return dataset_dict + def _calculate_metrics_from_split( self, split: str, hf_subset: str | None = None, compute_overall: bool = False ) -> RegressionDescriptiveStatistics: diff --git a/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py b/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py index 476c0afe5c..f3cb0fa9fd 100644 --- a/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py +++ b/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py @@ -51,10 +51,13 @@ class RuSciBenchCitedCountRegression(MultilingualTask, AbsTaskRegression): ) def dataset_transform(self): - for subset in self.dataset: - self.dataset[subset]["train"] = self.dataset[subset][ - "train" - ].train_test_split(test_size=2048, seed=self.seed)["test"] + self.dataset = self.stratified_subsampling( + self.dataset, + seed=self.seed, + splits=["train"], + label="value", + n_samples=2048, + ) class RuSciBenchYearPublRegression(MultilingualTask, AbsTaskRegression): @@ -103,7 +106,10 @@ class RuSciBenchYearPublRegression(MultilingualTask, AbsTaskRegression): ) def dataset_transform(self): - for subset in self.dataset: - self.dataset[subset]["train"] = self.dataset[subset][ - "train" - ].train_test_split(test_size=2048, seed=self.seed)["test"] + self.dataset = self.stratified_subsampling( + self.dataset, + seed=self.seed, + splits=["train"], + label="value", + n_samples=2048, + ) From f7a7907acfaa6da024f5f0b2f9490d0c3c2b911d Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Thu, 17 Jul 2025 18:06:31 +0000 Subject: [PATCH 11/17] Add boostrap for regression task --- mteb/abstasks/AbsTaskRegression.py | 68 +++++++++++++++---- .../evaluators/RegressionEvaluator.py | 30 ++++---- .../multilingual/RuSciBenchRegression.py | 18 ----- 3 files changed, 72 insertions(+), 44 deletions(-) diff --git a/mteb/abstasks/AbsTaskRegression.py b/mteb/abstasks/AbsTaskRegression.py index 01e577c1ce..7c0563fd90 100644 --- a/mteb/abstasks/AbsTaskRegression.py +++ b/mteb/abstasks/AbsTaskRegression.py @@ -4,6 +4,7 @@ from typing import Any import datasets +import numpy as np import pandas as pd from sklearn.base import BaseEstimator from sklearn.linear_model import LinearRegression @@ -66,6 +67,14 @@ class AbsTaskRegression(AbsTask): input_column_name: str = "text" abstask_prompt = "Predict the value of the user passage." + n_experiments: int = 10 + n_samples: int = 2048 + + def __init__(self, **kwargs: Any): + super().__init__(**kwargs) + self.n_experiments = self.metadata_dict.get("n_experiments", self.n_experiments) + self.n_samples = self.metadata_dict.get("n_samples", self.n_samples) + def _evaluate_subset( self, model: Encoder, @@ -79,17 +88,43 @@ def _evaluate_subset( train_split = dataset[self.train_split] eval_split = dataset[hf_split] - evaluator = self.evaluator( - train_split[self.input_column_name], - train_split[self.label_column_name], - eval_split[self.input_column_name], - eval_split[self.label_column_name], - task_name=self.metadata.name, - hf_subset=hf_subset, - **kwargs, - ) - scores = evaluator(model, encode_kwargs=encode_kwargs) - return scores + scores_list, test_cache = [], None + for i in range(self.n_experiments): + logger.info( + "=" * 10 + f" Experiment {i + 1}/{self.n_experiments} " + "=" * 10 + ) + + if self.n_samples >= len(train_split): + train_split_sampled = train_split + else: + train_split_sampled = self.stratified_subsampling( + datasets.DatasetDict({"train": train_split}), + seed=self.seed + i, + splits=["train"], + label=self.label_column_name, + n_samples=self.n_samples, + )["train"] + + evaluator = self.evaluator( + train_split_sampled[self.input_column_name], + train_split_sampled[self.label_column_name], + eval_split[self.input_column_name], + eval_split[self.label_column_name], + task_name=self.metadata.name, + hf_split=hf_split, + hf_subset=hf_subset, + **kwargs, + ) + scores, test_cache = evaluator( + model, encode_kwargs=encode_kwargs, test_cache=test_cache + ) + scores_list.append(scores) + + avg_scores: dict[str, Any] = { + k: np.mean([s[k] for s in scores_list]) for k in scores_list[0] + } + avg_scores["scores_per_experiment"] = scores_list + return avg_scores def _add_main_score(self, scores): scores["main_score"] = scores[self.metadata.main_score] @@ -106,8 +141,8 @@ def evaluate( if not self.data_loaded: self.load_data() - if "random_state" in self.classifier.get_params(): - self.classifier = self.classifier.set_params(random_state=self.seed) + if "random_state" in self.model.get_params(): + self.model = self.model.set_params(random_state=self.seed) scores = {} hf_subsets = self.hf_subsets @@ -171,7 +206,12 @@ def stratified_subsampling( binned_labels = pd.qcut(labels, q=n_bins, labels=False, duplicates="drop") dataset_with_bins: datasets.Dataset = dataset.add_column( - name=stratify_col_name, column=binned_labels.tolist() + name=stratify_col_name, + column=binned_labels.tolist(), + ) + dataset_with_bins = dataset_with_bins.cast_column( + stratify_col_name, + datasets.ClassLabel(names=np.unique(binned_labels).tolist()), ) subsampled_dataset = dataset_with_bins.train_test_split( diff --git a/mteb/evaluation/evaluators/RegressionEvaluator.py b/mteb/evaluation/evaluators/RegressionEvaluator.py index 418250782a..f83a6bafde 100644 --- a/mteb/evaluation/evaluators/RegressionEvaluator.py +++ b/mteb/evaluation/evaluators/RegressionEvaluator.py @@ -46,8 +46,12 @@ def __init__( self.fit_intercept = fit_intercept def __call__( - self, model: Encoder, *, encode_kwargs: dict[str, Any] = {} - ) -> dict[str, float]: + self, + model: Encoder, + *, + encode_kwargs: dict[str, Any] = {}, + test_cache: np.ndarray | None = None, + ) -> tuple[dict[str, float], np.ndarray]: scores = {} X_train = model.encode( self.sentences_train, @@ -57,18 +61,20 @@ def __call__( hf_subset=self.hf_subset, **encode_kwargs, ) - X_test = model.encode( - self.sentences_test, - model=model, - task_name=self.task_name, - hf_split=self.hf_split, - hf_subset=self.hf_subset, - **encode_kwargs, - ) + if test_cache is None: + X_test = model.encode( + self.sentences_test, + model=model, + task_name=self.task_name, + hf_split=self.hf_split, + hf_subset=self.hf_subset, + **encode_kwargs, + ) + test_cache = X_test linear_regression = LinearRegression(fit_intercept=self.fit_intercept) linear_regression.fit(X_train, self.y_train) - y_pred = linear_regression.predict(X_test) + y_pred = linear_regression.predict(test_cache) scores["mae"] = mean_absolute_error(self.y_test, y_pred) scores["mse"] = mean_squared_error(self.y_test, y_pred) @@ -76,4 +82,4 @@ def __call__( scores["r2"] = r2_score(self.y_test, y_pred) scores["kendalltau"] = kendalltau(self.y_test, y_pred).statistic - return scores + return scores, test_cache diff --git a/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py b/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py index f3cb0fa9fd..9d952e4d9b 100644 --- a/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py +++ b/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py @@ -50,15 +50,6 @@ class RuSciBenchCitedCountRegression(MultilingualTask, AbsTaskRegression): prompt="Predict the number of citations for a scientific article based on the title and abstract", ) - def dataset_transform(self): - self.dataset = self.stratified_subsampling( - self.dataset, - seed=self.seed, - splits=["train"], - label="value", - n_samples=2048, - ) - class RuSciBenchYearPublRegression(MultilingualTask, AbsTaskRegression): metadata = TaskMetadata( @@ -104,12 +95,3 @@ class RuSciBenchYearPublRegression(MultilingualTask, AbsTaskRegression): """, prompt="Predict paper publitaction year based on the title and abstract", ) - - def dataset_transform(self): - self.dataset = self.stratified_subsampling( - self.dataset, - seed=self.seed, - splits=["train"], - label="value", - n_samples=2048, - ) From f4736ba368a75b823345fee5bf53f8c28e1dd351 Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Tue, 22 Jul 2025 15:25:21 +0000 Subject: [PATCH 12/17] Rename task class, add model as evaluator argument --- mteb/abstasks/AbsTaskRegression.py | 13 +++++++++---- mteb/evaluation/evaluators/RegressionEvaluator.py | 9 +++++---- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/mteb/abstasks/AbsTaskRegression.py b/mteb/abstasks/AbsTaskRegression.py index 7c0563fd90..b6c6b4de82 100644 --- a/mteb/abstasks/AbsTaskRegression.py +++ b/mteb/abstasks/AbsTaskRegression.py @@ -1,12 +1,11 @@ from __future__ import annotations import logging -from typing import Any +from typing import Any, Protocol import datasets import numpy as np import pandas as pd -from sklearn.base import BaseEstimator from sklearn.linear_model import LinearRegression from mteb.abstasks.AbsTask import AbsTask @@ -50,7 +49,12 @@ class RegressionDescriptiveStatistics(DescriptiveStatistics): max_value: float -class AbsTaskRegression(AbsTask): +class RegressorModel(Protocol): + def fit(self, X, y, sample_weight=None): ... + def predict(self, X): ... + + +class AbsTaskTextRegression(AbsTask): """Abstract class for regression tasks self.load_data() must generate a huggingface dataset with a split matching self.metadata_dict["eval_splits"], and assign it to self.dataset. It @@ -60,7 +64,7 @@ class AbsTaskRegression(AbsTask): """ evaluator: type[LinearRegressionEvaluator] = LinearRegressionEvaluator - model: BaseEstimator = LinearRegression(n_jobs=-1) + regressor: RegressorModel = LinearRegression(n_jobs=-1) train_split: str = "train" label_column_name: str = "value" @@ -113,6 +117,7 @@ def _evaluate_subset( task_name=self.metadata.name, hf_split=hf_split, hf_subset=hf_subset, + regressor=self.regressor, **kwargs, ) scores, test_cache = evaluator( diff --git a/mteb/evaluation/evaluators/RegressionEvaluator.py b/mteb/evaluation/evaluators/RegressionEvaluator.py index f83a6bafde..5a747996c1 100644 --- a/mteb/evaluation/evaluators/RegressionEvaluator.py +++ b/mteb/evaluation/evaluators/RegressionEvaluator.py @@ -5,9 +5,9 @@ import numpy as np from scipy.stats import kendalltau -from sklearn.linear_model import LinearRegression from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score +from mteb.abstasks.AbsTaskRegression import RegressorModel from mteb.encoder_interface import Encoder from .Evaluator import Evaluator @@ -25,6 +25,7 @@ def __init__( task_name: str, hf_split: str, hf_subset: str, + regressor: RegressorModel, fit_intercept: bool = True, limit: int | None = None, **kwargs, @@ -44,6 +45,7 @@ def __init__( self.hf_subset = hf_subset self.task_name = task_name self.fit_intercept = fit_intercept + self.regressor = regressor def __call__( self, @@ -72,9 +74,8 @@ def __call__( ) test_cache = X_test - linear_regression = LinearRegression(fit_intercept=self.fit_intercept) - linear_regression.fit(X_train, self.y_train) - y_pred = linear_regression.predict(test_cache) + self.regressor.fit(X_train, self.y_train) + y_pred = self.regressor.predict(test_cache) scores["mae"] = mean_absolute_error(self.y_test, y_pred) scores["mse"] = mean_squared_error(self.y_test, y_pred) From d8aa9e9cef9309b2a5f4ab6ec54612f15d96dfdb Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Tue, 22 Jul 2025 15:59:12 +0000 Subject: [PATCH 13/17] fix import --- mteb/abstasks/AbsTaskRegression.py | 12 +++++------- mteb/evaluation/evaluators/RegressionEvaluator.py | 8 ++++++-- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/mteb/abstasks/AbsTaskRegression.py b/mteb/abstasks/AbsTaskRegression.py index b6c6b4de82..8bf602cf1f 100644 --- a/mteb/abstasks/AbsTaskRegression.py +++ b/mteb/abstasks/AbsTaskRegression.py @@ -1,7 +1,7 @@ from __future__ import annotations import logging -from typing import Any, Protocol +from typing import Any import datasets import numpy as np @@ -11,7 +11,10 @@ from mteb.abstasks.AbsTask import AbsTask from mteb.abstasks.TaskMetadata import DescriptiveStatistics, HFSubset from mteb.encoder_interface import Encoder -from mteb.evaluation.evaluators.RegressionEvaluator import LinearRegressionEvaluator +from mteb.evaluation.evaluators.RegressionEvaluator import ( + LinearRegressionEvaluator, + RegressorModel, +) from mteb.load_results.task_results import ScoresDict logger = logging.getLogger(__name__) @@ -49,11 +52,6 @@ class RegressionDescriptiveStatistics(DescriptiveStatistics): max_value: float -class RegressorModel(Protocol): - def fit(self, X, y, sample_weight=None): ... - def predict(self, X): ... - - class AbsTaskTextRegression(AbsTask): """Abstract class for regression tasks diff --git a/mteb/evaluation/evaluators/RegressionEvaluator.py b/mteb/evaluation/evaluators/RegressionEvaluator.py index 5a747996c1..752a573767 100644 --- a/mteb/evaluation/evaluators/RegressionEvaluator.py +++ b/mteb/evaluation/evaluators/RegressionEvaluator.py @@ -1,13 +1,12 @@ from __future__ import annotations import logging -from typing import Any +from typing import Any, Protocol import numpy as np from scipy.stats import kendalltau from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score -from mteb.abstasks.AbsTaskRegression import RegressorModel from mteb.encoder_interface import Encoder from .Evaluator import Evaluator @@ -15,6 +14,11 @@ logger = logging.getLogger(__name__) +class RegressorModel(Protocol): + def fit(self, X, y, sample_weight=None): ... + def predict(self, X): ... + + class LinearRegressionEvaluator(Evaluator): def __init__( self, From 3f9ba9bc1a00beca206ee53eb236b8fed054dec2 Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Tue, 22 Jul 2025 21:00:10 +0000 Subject: [PATCH 14/17] fix import 2 --- .../{AbsTaskRegression.py => AbsTaskTextRegression.py} | 0 mteb/abstasks/__init__.py | 2 +- mteb/tasks/Regression/multilingual/RuSciBenchRegression.py | 6 +++--- 3 files changed, 4 insertions(+), 4 deletions(-) rename mteb/abstasks/{AbsTaskRegression.py => AbsTaskTextRegression.py} (100%) diff --git a/mteb/abstasks/AbsTaskRegression.py b/mteb/abstasks/AbsTaskTextRegression.py similarity index 100% rename from mteb/abstasks/AbsTaskRegression.py rename to mteb/abstasks/AbsTaskTextRegression.py diff --git a/mteb/abstasks/__init__.py b/mteb/abstasks/__init__.py index d5401fb00e..aa6a049440 100644 --- a/mteb/abstasks/__init__.py +++ b/mteb/abstasks/__init__.py @@ -8,12 +8,12 @@ from .AbsTaskInstructionRetrieval import * from .AbsTaskMultilabelClassification import * from .AbsTaskPairClassification import * -from .AbsTaskRegression import * from .AbsTaskReranking import * from .AbsTaskRetrieval import * from .AbsTaskSpeedTask import * from .AbsTaskSTS import * from .AbsTaskSummarization import * +from .AbsTaskTextRegression import * from .Image.AbsTaskAny2AnyMultiChoice import * from .Image.AbsTaskAny2AnyRetrieval import * from .Image.AbsTaskImageClassification import * diff --git a/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py b/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py index 9d952e4d9b..f868e818e0 100644 --- a/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py +++ b/mteb/tasks/Regression/multilingual/RuSciBenchRegression.py @@ -1,11 +1,11 @@ from __future__ import annotations -from mteb.abstasks.AbsTaskRegression import AbsTaskRegression +from mteb.abstasks.AbsTaskTextRegression import AbsTaskTextRegression from mteb.abstasks.MultilingualTask import MultilingualTask from mteb.abstasks.TaskMetadata import TaskMetadata -class RuSciBenchCitedCountRegression(MultilingualTask, AbsTaskRegression): +class RuSciBenchCitedCountRegression(MultilingualTask, AbsTaskTextRegression): metadata = TaskMetadata( name="RuSciBenchCitedCountRegression", description="""Predicts the number of times a scientific article has been cited by other papers. @@ -51,7 +51,7 @@ class RuSciBenchCitedCountRegression(MultilingualTask, AbsTaskRegression): ) -class RuSciBenchYearPublRegression(MultilingualTask, AbsTaskRegression): +class RuSciBenchYearPublRegression(MultilingualTask, AbsTaskTextRegression): metadata = TaskMetadata( name="RuSciBenchYearPublRegression", description="""Predicts the publication year of a scientific article. The prediction is based on the From 38a004893af1bd34a86e83f4c46e498d0994e98b Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Fri, 25 Jul 2025 12:53:33 +0000 Subject: [PATCH 15/17] fixes --- mteb/abstasks/AbsTaskTextRegression.py | 4 ++-- mteb/benchmarks/benchmarks.py | 6 +++--- .../evaluators/RegressionEvaluator.py | 18 +++--------------- .../multilingual/RuSciBenchClassification.py | 18 ++++++++++++------ 4 files changed, 20 insertions(+), 26 deletions(-) diff --git a/mteb/abstasks/AbsTaskTextRegression.py b/mteb/abstasks/AbsTaskTextRegression.py index 8bf602cf1f..be15e43e43 100644 --- a/mteb/abstasks/AbsTaskTextRegression.py +++ b/mteb/abstasks/AbsTaskTextRegression.py @@ -6,6 +6,7 @@ import datasets import numpy as np import pandas as pd +from sklearn.base import RegressorMixin from sklearn.linear_model import LinearRegression from mteb.abstasks.AbsTask import AbsTask @@ -13,7 +14,6 @@ from mteb.encoder_interface import Encoder from mteb.evaluation.evaluators.RegressionEvaluator import ( LinearRegressionEvaluator, - RegressorModel, ) from mteb.load_results.task_results import ScoresDict @@ -62,7 +62,7 @@ class AbsTaskTextRegression(AbsTask): """ evaluator: type[LinearRegressionEvaluator] = LinearRegressionEvaluator - regressor: RegressorModel = LinearRegression(n_jobs=-1) + regressor: RegressorMixin = LinearRegression(n_jobs=-1) train_split: str = "train" label_column_name: str = "value" diff --git a/mteb/benchmarks/benchmarks.py b/mteb/benchmarks/benchmarks.py index d625cadeba..39c5255dd2 100644 --- a/mteb/benchmarks/benchmarks.py +++ b/mteb/benchmarks/benchmarks.py @@ -259,8 +259,8 @@ "RuSciBenchBitextMining", # Classification "RuSciBenchCoreRiscClassification", - "RuSciBenchGRNTIOrigClassification", - "RuSciBenchOECDOrigClassification", + "RuSciBenchGRNTIClassification.v2", + "RuSciBenchOECDClassification.v2", "RuSciBenchPubTypeClassification", # Retrieval "RuSciBenchCiteRetrieval", @@ -2055,7 +2055,7 @@ ] ), description="""R2MED: First Reasoning-Driven Medical Retrieval Benchmark. - R2MED is a high-quality, high-resolution information retrieval (IR) dataset designed for medical scenarios. + R2MED is a high-quality, high-resolution information retrieval (IR) dataset designed for medical scenarios. It contains 876 queries with three retrieval tasks, five medical scenarios, and twelve body systems. """, reference="https://r2med.github.io/", diff --git a/mteb/evaluation/evaluators/RegressionEvaluator.py b/mteb/evaluation/evaluators/RegressionEvaluator.py index 752a573767..47551aea34 100644 --- a/mteb/evaluation/evaluators/RegressionEvaluator.py +++ b/mteb/evaluation/evaluators/RegressionEvaluator.py @@ -1,10 +1,11 @@ from __future__ import annotations import logging -from typing import Any, Protocol +from typing import Any import numpy as np from scipy.stats import kendalltau +from sklearn.base import RegressorMixin from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score from mteb.encoder_interface import Encoder @@ -14,11 +15,6 @@ logger = logging.getLogger(__name__) -class RegressorModel(Protocol): - def fit(self, X, y, sample_weight=None): ... - def predict(self, X): ... - - class LinearRegressionEvaluator(Evaluator): def __init__( self, @@ -29,17 +25,10 @@ def __init__( task_name: str, hf_split: str, hf_subset: str, - regressor: RegressorModel, - fit_intercept: bool = True, - limit: int | None = None, + regressor: RegressorMixin, **kwargs, ): super().__init__(**kwargs) - if limit is not None: - sentences_train = sentences_train[:limit] - y_train = y_train[:limit] - sentences_test = sentences_test[:limit] - y_test = y_test[:limit] self.sentences_train = sentences_train self.y_train = y_train self.sentences_test = sentences_test @@ -48,7 +37,6 @@ def __init__( self.hf_split = hf_split self.hf_subset = hf_subset self.task_name = task_name - self.fit_intercept = fit_intercept self.regressor = regressor def __call__( diff --git a/mteb/tasks/Classification/multilingual/RuSciBenchClassification.py b/mteb/tasks/Classification/multilingual/RuSciBenchClassification.py index 43ee567170..97d0f11b97 100644 --- a/mteb/tasks/Classification/multilingual/RuSciBenchClassification.py +++ b/mteb/tasks/Classification/multilingual/RuSciBenchClassification.py @@ -100,9 +100,9 @@ class RuSciBenchPubTypeClassification(MultilingualTask, AbsTaskClassification): ) -class RuSciBenchGRNTIOrigClassification(MultilingualTask, AbsTaskClassification): +class RuSciBenchGRNTIClassificationV2(MultilingualTask, AbsTaskClassification): metadata = TaskMetadata( - name="RuSciBenchGRNTIOrigClassification", + name="RuSciBenchGRNTIClassification.v2", dataset={ "path": "mlsa-iai-msu-lab/ru_sci_bench_mteb", "revision": "fbc0599a0b5f00b3c7d87ab4d13490f04fb77f8e", @@ -110,7 +110,10 @@ class RuSciBenchGRNTIOrigClassification(MultilingualTask, AbsTaskClassification) description="""Classification of scientific papers based on the GRNTI (State Rubricator of Scientific and Technical Information) rubricator. GRNTI is a universal hierarchical classification of knowledge domains adopted in Russia and CIS countries to systematize the entire flow of scientific and technical information. - This task uses the first level of the GRNTI hierarchy and top 28 classes by frequency.""", + This task uses the first level of the GRNTI hierarchy and top 28 classes by frequency. + + In this version, English language support has been added and data partitioning has been slightly modified. + """, reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", type="Classification", category="p2p", @@ -147,9 +150,9 @@ class RuSciBenchGRNTIOrigClassification(MultilingualTask, AbsTaskClassification) ) -class RuSciBenchOECDOrigClassification(MultilingualTask, AbsTaskClassification): +class RuSciBenchOECDClassificationV2(MultilingualTask, AbsTaskClassification): metadata = TaskMetadata( - name="RuSciBenchOECDOrigClassification", + name="RuSciBenchOECDClassification.v2", dataset={ "path": "mlsa-iai-msu-lab/ru_sci_bench_mteb", "revision": "fbc0599a0b5f00b3c7d87ab4d13490f04fb77f8e", @@ -157,7 +160,10 @@ class RuSciBenchOECDOrigClassification(MultilingualTask, AbsTaskClassification): description="""Classification of scientific papers based on the OECD (Organization for Economic Co-operation and Development) rubricator. OECD provides a hierarchical 3-level system of classes for labeling scientific articles. - This task uses the first two levels of the OECD hierarchy, top 29 classes.""", + This task uses the first two levels of the OECD hierarchy, top 29 classes. + + In this version, English language support has been added and data partitioning has been slightly modified. + """, reference="https://github.com/mlsa-iai-msu-lab/ru_sci_bench_mteb", type="Classification", category="p2p", From 41bfb21f6d618e7b5801763b595d1b7e2aa6e12c Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Fri, 25 Jul 2025 21:50:23 +0000 Subject: [PATCH 16/17] fix --- mteb/tasks/Classification/rus/RuSciBenchGRNTIClassification.py | 2 +- mteb/tasks/Classification/rus/RuSciBenchOECDClassification.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/mteb/tasks/Classification/rus/RuSciBenchGRNTIClassification.py b/mteb/tasks/Classification/rus/RuSciBenchGRNTIClassification.py index 69dcbc06bd..76bd475741 100644 --- a/mteb/tasks/Classification/rus/RuSciBenchGRNTIClassification.py +++ b/mteb/tasks/Classification/rus/RuSciBenchGRNTIClassification.py @@ -5,7 +5,7 @@ class RuSciBenchGRNTIClassification(AbsTaskClassification): - superseded_by = "RuSciBenchGRNTIOrigClassification" + superseded_by = "RuSciBenchGRNTIClassification.v2" metadata = TaskMetadata( name="RuSciBenchGRNTIClassification", diff --git a/mteb/tasks/Classification/rus/RuSciBenchOECDClassification.py b/mteb/tasks/Classification/rus/RuSciBenchOECDClassification.py index ec1b7ba5ef..33affb702e 100644 --- a/mteb/tasks/Classification/rus/RuSciBenchOECDClassification.py +++ b/mteb/tasks/Classification/rus/RuSciBenchOECDClassification.py @@ -5,7 +5,7 @@ class RuSciBenchOECDClassification(AbsTaskClassification): - superseded_by = "RuSciBenchOECDOrigClassification" + superseded_by = "RuSciBenchOECDClassification.v2" metadata = TaskMetadata( name="RuSciBenchOECDClassification", From effab73ea8e835640a7be013f62982132b22fcf0 Mon Sep 17 00:00:00 2001 From: Alexey Vatolin Date: Fri, 25 Jul 2025 21:57:30 +0000 Subject: [PATCH 17/17] Rename regression model protocol --- mteb/abstasks/AbsTaskTextRegression.py | 4 ++-- mteb/evaluation/evaluators/RegressionEvaluator.py | 10 +++++++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/mteb/abstasks/AbsTaskTextRegression.py b/mteb/abstasks/AbsTaskTextRegression.py index be15e43e43..417bed14f0 100644 --- a/mteb/abstasks/AbsTaskTextRegression.py +++ b/mteb/abstasks/AbsTaskTextRegression.py @@ -6,7 +6,6 @@ import datasets import numpy as np import pandas as pd -from sklearn.base import RegressorMixin from sklearn.linear_model import LinearRegression from mteb.abstasks.AbsTask import AbsTask @@ -14,6 +13,7 @@ from mteb.encoder_interface import Encoder from mteb.evaluation.evaluators.RegressionEvaluator import ( LinearRegressionEvaluator, + SklearnRegressorModel, ) from mteb.load_results.task_results import ScoresDict @@ -62,7 +62,7 @@ class AbsTaskTextRegression(AbsTask): """ evaluator: type[LinearRegressionEvaluator] = LinearRegressionEvaluator - regressor: RegressorMixin = LinearRegression(n_jobs=-1) + regressor: SklearnRegressorModel = LinearRegression(n_jobs=-1) train_split: str = "train" label_column_name: str = "value" diff --git a/mteb/evaluation/evaluators/RegressionEvaluator.py b/mteb/evaluation/evaluators/RegressionEvaluator.py index 47551aea34..7242d1559e 100644 --- a/mteb/evaluation/evaluators/RegressionEvaluator.py +++ b/mteb/evaluation/evaluators/RegressionEvaluator.py @@ -1,11 +1,10 @@ from __future__ import annotations import logging -from typing import Any +from typing import Any, Protocol import numpy as np from scipy.stats import kendalltau -from sklearn.base import RegressorMixin from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score from mteb.encoder_interface import Encoder @@ -15,6 +14,11 @@ logger = logging.getLogger(__name__) +class SklearnRegressorModel(Protocol): + def fit(self, X, y, sample_weight=None): ... + def predict(self, X): ... + + class LinearRegressionEvaluator(Evaluator): def __init__( self, @@ -25,7 +29,7 @@ def __init__( task_name: str, hf_split: str, hf_subset: str, - regressor: RegressorMixin, + regressor: SklearnRegressorModel, **kwargs, ): super().__init__(**kwargs)