dataset

sentence_transformers.util.dataset.resolve_ids(lookups: dict[str, Dataset | tuple[Dataset, str] | tuple[Dataset, str, str]], keep_columns: list[str] | tuple[str, ...] = ('label', 'labels', 'score', 'scores'), max_list_length: int | None = None, output_format: Literal['columns', 'lists'] = 'columns') → Callable[[dict[str, Any]], dict[str, Any]]

Build a batched transform that resolves ID columns to their values via a join against lookup datasets. Useful for IR datasets that store query / document IDs alongside separate text (or image) datasets, e.g. lightonai/ms-marco-en-bge.

An input column may hold one ID per row (e.g. positive_id: str) or a list of IDs per row (e.g. document_ids: list[str]), detected from the row values. An ID that is missing from its lookup dataset raises a KeyError when the batch is read.

Pass the returned callable to set_transform() (lazy, no caching) or map() with batched=True (eager, cached). With map, also pass remove_columns=list(lookups) so the raw ID columns are dropped from the result: unlike set_transform, map merges the transform output with the untouched input columns. map also hands the transform batches in the dataset’s current format, so remove any format first with with_format(None).

Parameters:

• lookups –

  {input_col: lookup} for every input column to resolve, where lookup is:

  • a datasets.Dataset: the ID column (the single column named id or ending in _id) and the value column (the single remaining column) are inferred. Ambiguity raises at build time.

  • (dataset, id_col): explicit ID column, value column inferred.

  • (dataset, id_col, value_col): fully explicit.

  All other input columns are dropped (except keep_columns), as the losses read batch columns positionally. Insertion order sets the output column order, so list the anchor (query) column first.

• keep_columns – Input columns to pass through unresolved. Defaults to the label columns the data collators recognize by default (label, labels, score, scores). Mirror your collator’s valid_label_columns if you customized it. Absent columns are ignored.

• max_list_length – Truncate every list-per-row column (ID lists and list-valued kept columns) to the first max_list_length entries. Must be a positive integer or None (default, keep all). Note that this takes the first N as stored, not top-N by teacher score.

• output_format – "columns" (default) expands each list column into numbered columns (document_1, …, document_N), requiring a uniform list length per row. "lists" keeps list columns nested under a plural name (documents: list), allowing ragged rows. Use "lists" for CrossEncoder listwise training.

Returns:

A picklable batched transform mapping a batch dict of ID columns to a batch dict of resolved columns, for set_transform() or map() with batched=True.

Example (KD, LightOn ms-marco-en-bge):

from datasets import load_dataset
from sentence_transformers.util import resolve_ids

train = load_dataset("lightonai/ms-marco-en-bge", "train", split="train")
queries = load_dataset("lightonai/ms-marco-en-bge", "queries", split="train")
documents = load_dataset("lightonai/ms-marco-en-bge", "documents", split="train")

train.set_transform(resolve_ids({
    "query_id": queries,
    "document_ids": documents,
}, max_list_length=32))
# -> rows of {"query": str, "document_1": str, ..., "document_32": str, "scores": list[float]}

Example (triplet with IDs):

train.set_transform(resolve_ids({
    "query_id": queries,
    "positive_id": documents,
    "negative_id": documents,
}))
# -> rows of {"query": str, "positive": str, "negative": str}

Example (CrossEncoder listwise, nested lists):

train.set_transform(resolve_ids({
    "query_id": queries,
    "document_ids": documents,
}, output_format="lists"))
# -> rows of {"query": str, "documents": list[str], "scores": list[float]}

Example (streaming):

from datasets import Features, Sequence, Value

# Only the train split streams. The lookup datasets are random-access joins, so they
# must stay regular (materialized) Datasets.
train = load_dataset("lightonai/ms-marco-en-bge", "train", split="train", streaming=True)
train = train.map(
    resolve_ids({"query_id": queries, "document_ids": documents}, max_list_length=32),
    batched=True,
    remove_columns=["query_id", "document_ids"],
    features=Features({
        "query": Value("string"),
        **{f"document_{i}": Value("string") for i in range(1, 33)},
        "scores": Sequence(Value("float32")),
    }),
)
# Streaming map cannot infer the output schema, and the trainers require it, so pass
# ``features=`` explicitly.