[v2] Change `corpus` and `queries` to use `dataset` #2885

Samoed · 2025-07-06T19:50:37Z

I've updated retrieval split as following:

class RetrievalSplitData(TypedDict):
-    corpus: dict[str, dict[str, str]]
-    queries: dict[str, str]
+    corpus: Dataset
+    queries: Dataset
    relevant_docs: Mapping[str, Mapping[str, float]]
-    instructions: dict[str, str] | None
+    instructions: Dataset | None
    top_ranked: Mapping[str, list[str]] | None

I didn't change relevant_docs and top_ranked to a dataset because we're accessing these elements by ID, and using a dict is simpler for that. However, I think we can change it if needed.

I also modified the test a bit because we’re creating mock embeddings independently of the input. Previously, the corpus dict was sorted by keys and, for some reason, ended up as [d2, d1], so the first embedding was used. Maybe we should generate a random embedding based on the input instead. I’ve tested the branches on SCIDOCS, and the results didn’t change.

orionw

LGTM, thanks for the fix. I don't love that we have different types (dict and Dataset) but you're right it doesn't seem to make sense to force them to be datasets when we use them as dicts. And we have type annotations, so users can figure it out if they need to.

KennethEnevoldsen

~~Can you check that this still lead to the same results?~~ I see that you did that

Maybe we should generate a random embedding based on the input instead.

Yes. I have had that problem before as well. I have actually considered replacing the mockencoder with a small model2vec model (small enough to test with and is a more valid test model)

KennethEnevoldsen · 2025-07-25T15:17:29Z

mteb/abstasks/AbsTaskRetrieval.py

@@ -116,11 +120,13 @@ def __init__(self, **kwargs):
    def convert_v1_dataset_format_to_v2(self):


I am unsure how our changes influence the loading of v2 retrieval dataset (do we have any of those? It would be great to see how that would influence the load data setup.

How does this influence pushing datasets to the hub? Does it push in a v2 compatible manner, and do we need to update now?

KennethEnevoldsen · 2025-07-25T15:18:09Z

mteb/abstasks/AbsTaskRetrieval.py

@@ -116,11 +120,13 @@ def __init__(self, **kwargs):
    def convert_v1_dataset_format_to_v2(self):
        if not hasattr(self, "queries"):


Is this a way to check if it is v2? Might be worth adding a comment on this

KennethEnevoldsen · 2025-07-25T15:19:58Z

mteb/abstasks/AbsTaskRetrieval.py

+                cur_queries = split_data["queries"].map(
+                    map_indexes_for_ds, fn_kwargs=dict(hf_subset=hf_subset, split=split)
+                )
+                cur_corpus = split_data["corpus"].map(
+                    map_indexes_for_ds, fn_kwargs=dict(hf_subset=hf_subset, split=split)
+                )


isn't there already an index, why do we need to create a new one?

KennethEnevoldsen · 2025-07-25T15:21:57Z

mteb/abstasks/AbsTaskRetrieval.py

@@ -577,17 +636,24 @@ def calculate_corpus_length(


 def process_docs(
-    collection: dict[str, dict[str, str] | str], hf_subset: str, split: str
-) -> dict[str, str]:
+    collection: dict[str, list[str]], hf_subset: str, split: str


Do we still need this function? (seems like the following superseedes it)

KennethEnevoldsen · 2025-07-25T15:24:24Z

mteb/create_dataloaders.py

@@ -62,7 +62,7 @@ def corpus_to_dict(


 def create_dataloader_for_retrieval_corpus(


in this function we take a dataset, convert it to a dictionary and then back to a dataset - could we avoid the transformation?

KennethEnevoldsen · 2025-07-25T15:26:06Z

mteb/evaluation/evaluators/model_classes.py

@@ -70,15 +71,15 @@ def __init__(

    def search(
        self,
-        corpus: dict[str, dict[str, str]],
-        queries: dict[str, str | list[str] | list[dict[str, str]]],
+        corpus: Dataset,


We could introduce two new types:

QueryDataset = Dataset CorpusDataset = Dataset

To differentiate between the two. It also allows us to add some documentation on what those entail, but I am not sure if it is worth it.

I think this would be good addition

KennethEnevoldsen · 2025-07-25T15:29:11Z

tests/test_benchmark/mock_tasks.py

@@ -1419,12 +1439,12 @@ class MockRerankingTask(AbsTaskRetrieval):
    expected_stats = {
        "test": {
            "num_samples": 4,
-            "number_of_characters": 112,
+            "number_of_characters": 136,


were these incorrect before?

I've added title to the corpus

KennethEnevoldsen · 2025-07-25T15:30:19Z

tests/test_tasks/test_mteb_rerank.py

+    results = sorted(
+        results["1"].keys(), key=lambda x: (results["1"][x], x), reverse=True
+    )[:2]


reversed=True. Isn't that a bit worrying?

KennethEnevoldsen · 2025-07-25T15:30:39Z

tests/test_benchmark/test_models.py

@@ -27,7 +27,7 @@ def test_colbert_model_e2e(task: AbsTask, model: str, tmp_path: Path):
    )
    result = results[0]

-    assert result.scores["test"][0]["ndcg_at_1"] == 1.0
+    assert result.scores["test"][0]["ndcg_at_1"] == 0.0


Why would this change?

As I said in description, there is issue due to sorting and embedding creation

I also modified the test a bit because we’re creating mock embeddings independently of the input. Previously, the corpus dict was sorted by keys and, for some reason, ended up as [d2, d1], so the first embedding was used.

Samoed · 2025-07-25T15:42:39Z

Yes. I have had that problem before as well. I have actually considered replacing the mockencoder with a small model2vec model

We can use approach as is in langchain fake embedding by setting seed from the text https://github.com/langchain-ai/langchain/blob/master/libs%2Fcore%2Flangchain_core%2Fembeddings%2Ffake.py#L120

KennethEnevoldsen · 2025-07-25T15:49:08Z

We can use approach as is in langchain fake embedding by setting seed from the text https://github.com/langchain-ai/langchain/blob/master/libs%2Fcore%2Flangchain_core%2Fembeddings%2Ffake.py#L120

Great idea! Let us do that instead.

Samoed added 3 commits July 6, 2025 22:29

change corpus and queries to dataset

d0eb2b2

remove commented out code

9ac554d

add convertion for v1 datasets

6b3d00c

Samoed changed the title ~~Retrieval ds~~ Change corpus and queries to use dataset Jul 6, 2025

Samoed added the v2 Issues and PRs related to `v2` branch label Jul 6, 2025

Samoed changed the title ~~Change corpus and queries to use dataset~~ [v2] Change corpus and queries to use dataset Jul 6, 2025

Samoed added 9 commits July 20, 2025 14:46

Merge branch 'v2.0.0' into retrieval_ds

ed38ff4

fix descriptive stats

b41c5e5

update reranking

e5db6e5

Merge branch 'v2.0.0' into retrieval_ds

6a61e17

format

9d28075

fix tests

670aad5

lint

6ae75a6

change ids of mock dataset

59d3959

change score for colbert

0708f32

Samoed marked this pull request as ready for review July 25, 2025 14:05

Samoed requested review from orionw and KennethEnevoldsen July 25, 2025 14:05

orionw approved these changes Jul 25, 2025

View reviewed changes

KennethEnevoldsen reviewed Jul 25, 2025

View reviewed changes

		@@ -116,11 +120,13 @@ def __init__(self, **kwargs):
		def convert_v1_dataset_format_to_v2(self):

		@@ -116,11 +120,13 @@ def __init__(self, **kwargs):
		def convert_v1_dataset_format_to_v2(self):
		if not hasattr(self, "queries"):

		@@ -62,7 +62,7 @@ def corpus_to_dict(


		def create_dataloader_for_retrieval_corpus(

[v2] Change corpus and queries to use dataset #2885

Are you sure you want to change the base?

[v2] Change corpus and queries to use dataset #2885

Uh oh!

Conversation

Samoed commented Jul 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

orionw left a comment

Choose a reason for hiding this comment

Uh oh!

KennethEnevoldsen left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Samoed Jul 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Samoed commented Jul 25, 2025

Uh oh!

KennethEnevoldsen commented Jul 25, 2025

Uh oh!

Uh oh!

[v2] Change `corpus` and `queries` to use `dataset` #2885

[v2] Change `corpus` and `queries` to use `dataset` #2885

Samoed commented Jul 6, 2025 •

edited

Loading

KennethEnevoldsen left a comment •

edited

Loading

Samoed Jul 25, 2025 •

edited

Loading