[WIP] Add RTEB retrieval code

mteb/model_meta.py

@@ -38,6 +38,7 @@
     "ColBERT",
 ]
 DISTANCE_METRICS = Literal["cosine", "max_sim", "dot"]
+EMBEDDING_DTYPES = Literal["float32", "int8", "binary"]
 
 
 def sentence_transformers_loader(
@@ -58,6 +59,34 @@ def get_loader_name(
     return loader.__name__
 
 
+def model_id(
+    model_name: str,
+    embd_dtype: str
+    | None,  # Keep None here as input can still be None before default assignment
+    embd_dim: int | None,
+) -> str:
+    # Handle potential None values passed to the function, even if the class attribute has a default
+    if model_name is None:
+        # Or handle appropriately, maybe raise error if name is critical for ID
+        model_name_part = "unknown_model"
+    else:
+        model_name_part = model_name.replace("/", "__")
+
+    dtype_str = embd_dtype if embd_dtype else "unknown_dtype"
+    dim_str = f"{embd_dim}d" if embd_dim else "unknown_dim"
+
+    # Check if default was used implicitly for dtype
+    if embd_dtype is None:
+        # If the class attribute defaults to 'float32', reflect that possibility if None is passed
+        # However, the class instance will have 'float32' if not specified.
+        # Let's assume the function should reflect the actual value passed or derived.
+        # If the intention is to always use the default if None is passed, adjust logic here.
+        # For now, stick to representing the input or lack thereof.
+        pass  # dtype_str is already "unknown_dtype"
+
+    return f"{model_name_part}_{dtype_str}_{dim_str}"
+
+
 class ModelMeta(BaseModel):
     """The model metadata object.
 
@@ -70,6 +99,7 @@ class ModelMeta(BaseModel):
         max_tokens: The maximum number of tokens the model can handle. Can be None if the maximum number of tokens is not known (e.g. for proprietary
             models).
         embed_dim: The dimension of the embeddings produced by the model. Currently all models are assumed to produce fixed-size embeddings.
+        embd_dtype: The data type of the embeddings produced by the model (e.g., "float32", "int8", "binary"). Defaults to "float32".
         revision: The revision number of the model. If None, it is assumed that the metadata (including the loader) is valid for all revisions of the model.
         release_date: The date the model's revision was released.
         license: The license under which the model is released. Required if open_weights is True.
@@ -116,6 +146,10 @@ class ModelMeta(BaseModel):
     superseded_by: str | None = None
     is_cross_encoder: bool | None = None
     modalities: list[MODALITIES] = ["text"]
+    # Attribute merged from rteb/ebr/core/meta.py
+    embd_dtype: EMBEDDING_DTYPES = (
+        "float32"  # Defaulting to float32 as requested, type hint updated
+    )
 
     def to_dict(self):
         dict_repr = self.model_dump()
@@ -263,6 +297,15 @@ def calculate_memory_usage_mb(self) -> int | None:
         model_memory_mb = model_memory_bytes / MB
         return round(model_memory_mb)
 
+    @property
+    def _id(self) -> str:
+        """Generates a unique ID for the model based on name, dtype, and dimension."""
+        if self.name is None:
+            raise ValueError("Model name is required to generate an ID.")
+        # Note: Using target's embed_dim and the newly added embd_dtype
+        # self.embd_dtype will be 'float32' by default if not specified otherwise
+        return model_id(self.name, self.embd_dtype, self.embed_dim)
+
 
 def collect_similar_tasks(dataset: str, visited: set[str]) -> set[str]:
     """Recursively collect all similar tasks for a given dataset."""

mteb/rteb/ebr/__main__.py

@@ -0,0 +1,220 @@
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+from collections import defaultdict
+from pathlib import Path
+
+import pytorch_lightning as pl
+from ebr.core import Encoder, Retriever
+from ebr.datasets import DATASET_REGISTRY, DatasetMeta
+from ebr.models import MODEL_REGISTRY, ModelMeta
+from ebr.retrieve import run_retrieve_task
+from pytorch_lightning.strategies.ddp import DDPStrategy
+
+logger = logging.getLogger(__name__)
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+def get_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+
+    # Evaluation
+    parser.add_argument(
+        "--gpus", type=int, default=0, help="Number of gpus used for encoding."
+    )
+    parser.add_argument(
+        "--cpus",
+        type=int,
+        default=1,
+        help="Number of cpus used for computation (this is only for models that are not using gpus).",
+    )
+    parser.add_argument("--bf16", action="store_true", help="`Use bf16 precision.")
+    parser.add_argument(
+        "--batch_size", type=int, default=16, help="Batch size for encoding."
+    )
+    parser.add_argument(
+        "--embd_batch_size",
+        type=int,
+        default=1024,
+        help="Batch size for computing similarity of embeddings.",
+    )
+    parser.add_argument(
+        "--embd_in_memory_threshold",
+        type=int,
+        default=200000,
+        help="Embeddings will be stored in memory if the amount is below this threshold.",
+    )
+
+    # Model
+    # parser.add_argument(
+    #    "--model_name", type=str, default=None, help="Model name or path.")
+    # parser.add_argument(
+    #    "--embd_dtype", type=str, default="float", help="Embedding type. Options: float32, int8, binary.")
+    # parser.add_argument(
+    #    "--embd_dim", type=int, default=None, help="Embedding dimension.")
+    # parser.add_argument(
+    #    "--max_length", type=int, default=None, help="Maximum length of model input.")
+
+    # Data
+    parser.add_argument(
+        "--data_path",
+        type=str,
+        default="data/",
+        help="Path of the dataset, must be specified for custom tasks.",
+    )
+    parser.add_argument(
+        "--task_name",
+        type=str,
+        default=None,
+        help="Name of the task. Can be multiple tasks splitted by `,`.",
+    )
+    parser.add_argument(
+        "--data_type",
+        default="eval",
+        choices=["eval", "train", "chunk", "merge"],
+        help="Dataset type.",
+    )
+    parser.add_argument(
+        "--num_workers", type=int, default=4, help="Number of workers for dataloader."
+    )
+
+    # Output
+    parser.add_argument(
+        "--save_path", type=str, default="output/", help="Path to save the output."
+    )
+    parser.add_argument(
+        "--save_embds", action="store_true", help="Whether to save the embeddings."
+    )
+    parser.add_argument(
+        "--load_embds",
+        action="store_true",
+        help="Whether to load the computed embeddings.",
+    )
+    parser.add_argument(
+        "--save_prediction",
+        action="store_true",
+        help="Whether to save the predictions.",
+    )
+    parser.add_argument(
+        "--topk", type=int, default=100, help="Number of top documents per query."
+    )
+    parser.add_argument(
+        "--overwrite", action="store_true", help="Whether to overwrite the results."
+    )
+
+    args = parser.parse_args()
+    return args
+
+
+def _dump_model_meta(
+    results_dir: str = "results",
+    model_registry: dict[str, ModelMeta] = MODEL_REGISTRY,
+):
+    models = [meta.model_dump() for meta in model_registry.values()]
+    with open(Path(results_dir) / "models.json", "w") as f:
+        f.write(json.dumps(models, indent=4))
+
+
+def _dump_dataset_info(
+    results_dir: str = "results",
+    dataset_registry: dict[str, DatasetMeta] = DATASET_REGISTRY,
+):
+    group_data = defaultdict(list)
+    for dataset_meta in dataset_registry.values():
+        for group_name in dataset_meta.groups.keys():
+            leaderboard = dataset_meta.loader.LEADERBOARD
+            group_data[(leaderboard, group_name)].append(dataset_meta.dataset_name)
+
+    groups = []
+    for (leaderboard, group_name), datasets in group_data.items():
+        groups.append(
+            {"name": group_name, "datasets": datasets, "leaderboard": leaderboard}
+        )
+    with open(Path(results_dir) / "datasets.json", "w") as f:
+        f.write(json.dumps(groups, indent=4))
+
+
+def _compile_results(results_dir: str = "results", output_dir: str = "output"):
+    results = []
+    for dataset_output_dir in Path(output_dir).iterdir():
+        dataset_results = []
+        for one_result in dataset_output_dir.iterdir():
+            eval_file = one_result / "retrieve_eval.json"
+            if eval_file.exists():
+                with open(eval_file) as f:
+                    dataset_results.append(json.load(f))
+
+        results.append(
+            {
+                **DATASET_REGISTRY[dataset_output_dir.name].model_dump(),
+                "results": dataset_results,
+                "is_closed": DATASET_REGISTRY[dataset_output_dir.name].tier != 3,
+            }
+        )
+
+    with open(Path(results_dir) / "results.json", "w") as f:
+        f.write(json.dumps(results, indent=4))
+
+
+def main(args: argparse.Namespace):
+    _dump_model_meta()
+    _dump_dataset_info()
+
+    if args.gpus:
+        trainer = pl.Trainer(
+            strategy=DDPStrategy(find_unused_parameters=False),
+            accelerator="gpu",
+            devices=args.gpus,
+            precision="bf16" if args.bf16 else "32",
+        )
+    else:
+        trainer = pl.Trainer(
+            strategy=DDPStrategy(),
+            accelerator="cpu",
+            devices=args.cpus,
+        )
+
+    if not trainer.is_global_zero:
+        logging.basicConfig(level=logging.ERROR)
+
+    # Evaluate each model on the specified datasets
+    for model_meta in MODEL_REGISTRY.values():
+        encoder = Encoder(
+            model_meta.load_model(),
+            save_embds=args.save_embds,
+            load_embds=args.load_embds,
+        )
+        retriever = Retriever(
+            topk=args.topk,
+            similarity=model_meta.similarity,
+            save_prediction=args.save_prediction,
+        )
+
+        eval_results = {}
+        for dataset_meta in DATASET_REGISTRY.values():
+            # if trainer.is_global_zero:
+            #    trainer.print(f"Evaluating {model_meta.model_name} on {dataset_meta.dataset_name}")
+
+            result = run_retrieve_task(dataset_meta, trainer, encoder, retriever, args)
+            eval_results[dataset_meta.dataset_name] = result
+
+        metric = "ndcg_at_10"
+
+        # Print the results
+        if trainer.is_global_zero:
+            trainer.print("=" * 40)
+            trainer.print(args.save_path)
+            trainer.print("=" * 40)
+            for task in eval_results.keys():
+                if metric in eval_results[task]:
+                    trainer.print(f"{task:<32}{eval_results[task][metric]:.4f}")
+
+    _compile_results()
+
+
+if __name__ == "__main__":
+    args = get_args()
+    main(args)

mteb/rteb/ebr/core/__init__.py

@@ -0,0 +1 @@
+from __future__ import annotations

mteb/rteb/ebr/core/base/__init__.py

@@ -0,0 +1 @@
+from __future__ import annotations

mteb/rteb/ebr/core/base/dataset.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+from abc import ABC
+from functools import cache
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from torch.utils.data import Dataset
+
+if TYPE_CHECKING:
+    from ebr.core.meta import DatasetMeta
+
+
+def add_instruct(dataset: Dataset, instruct: str, input_type: str):
+    for item in dataset.data:
+        if instruct:
+            item["text"] = instruct + item["text"]
+        item["input_type"] = input_type
+
+    return dataset
+
+
+class RetrievalDataset(ABC):
+    LEADERBOARD: str = None
+
+    def __init__(
+        self,
+        data_path: str,
+        dataset_meta: DatasetMeta,
+        query_instruct: str | None = None,
+        corpus_instruct: str | None = None,
+        **kwargs,
+    ):
+        assert type(self).LEADERBOARD, "leaderboard must be defined"
+        super().__init__()
+        self._dataset_meta = dataset_meta
+        self._query_instruct = query_instruct
+        self._corpus_instruct = corpus_instruct
+        self._task_path = (Path(data_path) / dataset_meta.dataset_name).resolve()
+
+    # def __getattr__(self, name: str) -> Any:
+    #    try:
+    #        return super().__getattr__(name)
+    #    except AttributeError:
+    #        return getattr(self._dataset_meta, name)
+
+    @property
+    @cache
+    def corpus(self) -> Dataset:
+        corpus = self._corpus()
+        corpus = add_instruct(corpus, self._corpus_instruct, "document")
+        return corpus
+
+    def _corpus(self) -> Dataset:
+        raise NotImplementedError
+
+    @property
+    @cache
+    def queries(self) -> Dataset:
+        queries = self._queries()
+        queries = add_instruct(queries, self._query_instruct, "query")
+        return queries
+
+    def _queries(self) -> Dataset:
+        raise NotImplementedError
+
+    @property
+    @cache
+    def relevance(self) -> dict:
+        # Dict of dict: relevance[query_id][corpus_id] = score
+        pass
+
+    def prepare_data(self):
+        _ = self.corpus
+        _ = self.queries
+        _ = self.relevance