-
Notifications
You must be signed in to change notification settings - Fork 238
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
d6bd6f4
commit 8140f5f
Showing
3 changed files
with
123 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,4 @@ | ||
from __future__ import annotations | ||
|
||
import mteb | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,117 @@ | ||
"""This script contains functions that are used to get an overview of the MTEB benchmark.""" | ||
|
||
from __future__ import annotations | ||
|
||
import logging | ||
|
||
from mteb.abstasks import AbsTask | ||
from mteb.abstasks.languages import ( | ||
ISO_TO_LANGUAGE, | ||
ISO_TO_SCRIPT, | ||
path_to_lang_codes, | ||
path_to_lang_scripts, | ||
) | ||
from mteb.abstasks.TaskMetadata import TASK_DOMAIN, TASK_TYPE | ||
from mteb.tasks import * # import all tasks | ||
|
||
logger = logging.getLogger(__name__) | ||
|
||
|
||
def check_is_valid_script(script: str) -> None: | ||
if script not in ISO_TO_SCRIPT: | ||
raise ValueError( | ||
f"Invalid script code: {script}, you can find valid ISO 15924 codes in {path_to_lang_scripts}" | ||
) | ||
|
||
|
||
def check_is_valid_language(lang: str) -> None: | ||
if lang not in ISO_TO_LANGUAGE: | ||
raise ValueError( | ||
f"Invalid language code: {lang}, you can find valid ISO 639-3 codes in {path_to_lang_codes}" | ||
) | ||
|
||
|
||
def filter_superseeded_datasets(tasks: list[AbsTask]) -> list[AbsTask]: | ||
return [t for t in tasks if t.is_superseeded is None] | ||
|
||
|
||
def filter_tasks_by_languages( | ||
tasks: list[AbsTask], languages: list[str] | ||
) -> list[AbsTask]: | ||
[check_is_valid_language(lang) for lang in languages] | ||
langs_to_keep = set(languages) | ||
return [t for t in tasks if langs_to_keep.intersection(t.metadata.languages)] | ||
|
||
|
||
def filter_tasks_by_script(tasks: list[AbsTask], script: list[str]) -> list[AbsTask]: | ||
[check_is_valid_script(s) for s in script] | ||
script_to_keep = set(script) | ||
return [t for t in tasks if script_to_keep.intersection(t.metadata.scripts)] | ||
|
||
|
||
def filter_tasks_by_domains( | ||
tasks: list[AbsTask], domains: list[TASK_DOMAIN] | ||
) -> list[AbsTask]: | ||
domains_to_keep = set(domains) | ||
|
||
def _convert_to_set(domain: list[TASK_DOMAIN] | None) -> set: | ||
return set(domain) if domain is not None else set() | ||
|
||
return [ | ||
t | ||
for t in tasks | ||
if domains_to_keep.intersection(_convert_to_set(t.metadata.domains)) | ||
] | ||
|
||
|
||
def filter_tasks_by_task_type( | ||
tasks: list[AbsTask], task_type: TASK_TYPE | ||
) -> list[AbsTask]: | ||
return [t for t in tasks if t.metadata.type == task_type] | ||
|
||
|
||
def get_tasks( | ||
languages: list[str] | None = None, | ||
script: list[str] | None = None, | ||
domains: list[TASK_DOMAIN] | None = None, | ||
task_type: TASK_TYPE | None = None, | ||
exclude_superseeded_datasets: bool = True, | ||
) -> list[AbsTask]: | ||
"""Get a list of tasks based on the specified filters. | ||
Args: | ||
languages: A list of languages either specified as 3 letter languages codes (ISO 639-3, e.g. "eng") or as script languages codes e.g. | ||
"eng-Latn". | ||
script: A list of script codes (ISO 15924 codes). If None, all scripts are included. | ||
domains: A list of task domains. | ||
task_type: A string specifying the type of task. If None, all tasks are included. | ||
exclude_superseeded_datasets: A boolean flag to exclude datasets which are superseeded by another. | ||
Returns: | ||
A list of all initialized tasks objects which pass all of the filters (AND operation). | ||
Examples: | ||
>>> get_tasks(languages=["eng", "deu"], script=["Latn"], domains=["Legal"]) | ||
>>> get_tasks(languages=["eng"], script=["Latn"], task_type="Classification") | ||
>>> get_tasks(languages=["eng"], script=["Latn"], task_type="Clustering", exclude_superseeded_datasets=False) | ||
""" | ||
tasks_categories_cls = [cls for cls in AbsTask.__subclasses__()] | ||
tasks = [ | ||
cls() | ||
for cat_cls in tasks_categories_cls | ||
for cls in cat_cls.__subclasses__() | ||
if cat_cls.__name__.startswith("AbsTask") | ||
] | ||
|
||
if languages: | ||
tasks = filter_tasks_by_languages(tasks, languages) | ||
if script: | ||
tasks = filter_tasks_by_script(tasks, script) | ||
if domains: | ||
tasks = filter_tasks_by_domains(tasks, domains) | ||
if exclude_superseeded_datasets: | ||
tasks = filter_superseeded_datasets(tasks) | ||
if task_type: | ||
tasks = filter_tasks_by_task_type(tasks, task_type) | ||
|
||
return tasks |