Skip to content

Data And Datasets

RecDistillery trains and evaluates recommendation models from interaction splits stored under data/<dataset>/. Dataset preprocessing is based on DataRec, while the training pipeline consumes encoded PyTorch-ready interaction objects from recdistill.data.

Dataset Layout

Each dataset is expected to expose the canonical split files:

data/<dataset>/
  train.tsv
  val.tsv
  test.tsv

The example dataset configs live in:

config/dataset/
  amazon_cd.yaml
  bookcrossing.yaml
  citeulike.yaml

Preparation Scripts

Dataset-specific preparation entry points are stored in scripts/data_preparation/:

scripts/data_preparation/amazon_cd_2014.py
scripts/data_preparation/bookcrossing.py
scripts/data_preparation/citeulike.py

These scripts prepare the split files consumed by the training and evaluation loaders. They are intended as examples, while the framework can be extended to any dataset either directly available in DataRec or loaded through its multi-format data interfaces.

Runtime Data Objects

InteractionBatch dataclass

Source code in recdistill/data/batch.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
@dataclass
class InteractionBatch:
    users: torch.Tensor
    pos_items: torch.Tensor
    neg_items: torch.Tensor

    @property
    def unique_users(self) -> torch.Tensor:
        return self.users.unique()

    @property
    def unique_items(self) -> torch.Tensor:
        return torch.cat([self.pos_items, self.neg_items]).unique()

users: torch.Tensor instance-attribute

pos_items: torch.Tensor instance-attribute

neg_items: torch.Tensor instance-attribute

unique_users: torch.Tensor property

unique_items: torch.Tensor property

__init__(users: torch.Tensor, pos_items: torch.Tensor, neg_items: torch.Tensor) -> None

RRDAuxBatch dataclass

Source code in recdistill/data/batch.py 
23
24
25
26
27
@dataclass
class RRDAuxBatch:
    users: torch.Tensor
    interesting_items: torch.Tensor
    uninteresting_items: torch.Tensor

users: torch.Tensor instance-attribute

interesting_items: torch.Tensor instance-attribute

uninteresting_items: torch.Tensor instance-attribute

__init__(users: torch.Tensor, interesting_items: torch.Tensor, uninteresting_items: torch.Tensor) -> None

UnKDAuxBatch dataclass

Source code in recdistill/data/batch.py 
30
31
32
33
34
@dataclass
class UnKDAuxBatch:
    users: torch.Tensor
    pos_items: torch.Tensor
    neg_items: torch.Tensor

users: torch.Tensor instance-attribute

pos_items: torch.Tensor instance-attribute

neg_items: torch.Tensor instance-attribute

__init__(users: torch.Tensor, pos_items: torch.Tensor, neg_items: torch.Tensor) -> None

InteractionDataset dataclass

Source code in recdistill/data/interactions.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
@dataclass
class InteractionDataset:
    num_users: int
    num_items: int
    train_dict: dict[int, set[int]]
    interactions: list[tuple[int, int]]

    @classmethod
    def from_train_dict(
        cls,
        train_dict: dict[int, set[int]] | dict[int, list[int]],
        num_users: int,
        num_items: int,
    ) -> "InteractionDataset":
        normalized = {int(u): set(map(int, items)) for u, items in train_dict.items()}
        interactions = [
            (user, item) for user, items in normalized.items() for item in sorted(items)
        ]
        return cls(
            num_users=num_users,
            num_items=num_items,
            train_dict=normalized,
            interactions=interactions,
        )

    def seen_items(self, user: int) -> set[int]:
        return self.train_dict.get(int(user), set())

num_users: int instance-attribute

num_items: int instance-attribute

train_dict: dict[int, set[int]] instance-attribute

interactions: list[tuple[int, int]] instance-attribute

__init__(num_users: int, num_items: int, train_dict: dict[int, set[int]], interactions: list[tuple[int, int]]) -> None

from_train_dict(train_dict: dict[int, set[int]] | dict[int, list[int]], num_users: int, num_items: int) -> 'InteractionDataset' classmethod

Source code in recdistill/data/interactions.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
@classmethod
def from_train_dict(
    cls,
    train_dict: dict[int, set[int]] | dict[int, list[int]],
    num_users: int,
    num_items: int,
) -> "InteractionDataset":
    normalized = {int(u): set(map(int, items)) for u, items in train_dict.items()}
    interactions = [
        (user, item) for user, items in normalized.items() for item in sorted(items)
    ]
    return cls(
        num_users=num_users,
        num_items=num_items,
        train_dict=normalized,
        interactions=interactions,
    )

seen_items(user: int) -> set[int]

Source code in recdistill/data/interactions.py 
31
32
def seen_items(self, user: int) -> set[int]:
    return self.train_dict.get(int(user), set())

DataRec Loading

USER_COLUMNS = ('user', 'userId', 'user_id', 'uid', 0) module-attribute

ITEM_COLUMNS = ('item', 'itemId', 'item_id', 'iid', 1) module-attribute

RATING_COLUMNS = ('rating', 'ratings', 'score', 2) module-attribute

TIMESTAMP_COLUMNS = ('timestamp', 'time', 'ts', 3) module-attribute

SPLIT_ORDER = ('train', 'val', 'test') module-attribute

LoadedSplit dataclass

Source code in recdistill/data/datarec_loader.py 
21
22
23
24
25
26
@dataclass(frozen=True)
class LoadedSplit:
    name: str
    path: Path
    frame: pd.DataFrame
    backend: str

name: str instance-attribute

path: Path instance-attribute

frame: pd.DataFrame instance-attribute

backend: str instance-attribute

__init__(name: str, path: Path, frame: pd.DataFrame, backend: str) -> None

EncodedDataset dataclass

Source code in recdistill/data/datarec_loader.py 
29
30
31
32
33
34
@dataclass(frozen=True)
class EncodedDataset:
    frames: dict[str, pd.DataFrame]
    num_users: int
    num_items: int
    backend: str

frames: dict[str, pd.DataFrame] instance-attribute

num_users: int instance-attribute

num_items: int instance-attribute

backend: str instance-attribute

__init__(frames: dict[str, pd.DataFrame], num_users: int, num_items: int, backend: str) -> None

datarec_available() -> bool

Source code in recdistill/data/datarec_loader.py 
37
38
def datarec_available() -> bool:
    return _datarec_reader() is not None

load_split_frame(dataset_name: str, split_name: str, *, repo_root: Path | str = '.', use_datarec: bool = True) -> LoadedSplit

Load one dataset split through DataRec, falling back to pandas if unavailable.

Source code in recdistill/data/datarec_loader.py 
41
42
43
44
45
46
47
48
49
50
51
def load_split_frame(
    dataset_name: str,
    split_name: str,
    *,
    repo_root: Path | str = ".",
    use_datarec: bool = True,
) -> LoadedSplit:
    """Load one dataset split through DataRec, falling back to pandas if unavailable."""
    path = _split_path(dataset_name=dataset_name, split_name=split_name, repo_root=repo_root)
    columns = _column_names(dataset_name)
    return load_frame_from_path(path, dataset_name=dataset_name, split_name=split_name, columns=columns, use_datarec=use_datarec)

load_frame_from_path(path: Path | str, *, dataset_name: str, split_name: str, columns: list[str] | None = None, use_datarec: bool = True) -> LoadedSplit

Source code in recdistill/data/datarec_loader.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def load_frame_from_path(
    path: Path | str,
    *,
    dataset_name: str,
    split_name: str,
    columns: list[str] | None = None,
    use_datarec: bool = True,
) -> LoadedSplit:
    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"{split_name} split not found at {path}")
    columns = columns or ["userId", "itemId", "rating"]

    reader = _datarec_reader() if use_datarec else None
    if reader is not None:
        data = reader(
            str(path),
            sep="\t",
            user_col=columns[0],
            item_col=columns[1],
            rating_col=columns[2] if len(columns) > 2 else None,
            timestamp_col=columns[3] if len(columns) > 3 else None,
            header=None,
            cols=columns,
            dataset_name=dataset_name,
            version_name=split_name,
        )
        return LoadedSplit(
            name=split_name,
            path=path,
            frame=_datarec_to_frame(data),
            backend="datarec",
        )

    frame = pd.read_csv(path, sep="\t", header=None, names=columns)
    return LoadedSplit(name=split_name, path=path, frame=frame, backend="pandas")

load_train_dataset(dataset_name: str, teacher_num_users: int, teacher_num_items: int, user_mapping: dict[int, int] | dict[str, int] | None = None, item_mapping: dict[int, int] | dict[str, int] | None = None, id_space: str | None = None) -> tuple[InteractionDataset, int]

Source code in recdistill/data/datarec_loader.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def load_train_dataset(
    dataset_name: str,
    teacher_num_users: int,
    teacher_num_items: int,
    user_mapping: dict[int, int] | dict[str, int] | None = None,
    item_mapping: dict[int, int] | dict[str, int] | None = None,
    id_space: str | None = None,
) -> tuple[InteractionDataset, int]:
    train_dict, dropped = load_ground_truth_split(
        dataset_name=dataset_name,
        split_name="train",
        num_users=teacher_num_users,
        num_items=teacher_num_items,
        user_mapping=user_mapping,
        item_mapping=item_mapping,
        id_space=id_space,
    )
    return (
        InteractionDataset.from_train_dict(
            train_dict=train_dict,
            num_users=teacher_num_users,
            num_items=teacher_num_items,
        ),
        dropped,
    )

load_eval_split(dataset_name: str, split_name: str, teacher_num_users: int, teacher_num_items: int, user_mapping: dict[int, int] | dict[str, int] | None = None, item_mapping: dict[int, int] | dict[str, int] | None = None, id_space: str | None = None) -> tuple[dict[int, set[int]], int]

Source code in recdistill/data/datarec_loader.py 
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
def load_eval_split(
    dataset_name: str,
    split_name: str,
    teacher_num_users: int,
    teacher_num_items: int,
    user_mapping: dict[int, int] | dict[str, int] | None = None,
    item_mapping: dict[int, int] | dict[str, int] | None = None,
    id_space: str | None = None,
) -> tuple[dict[int, set[int]], int]:
    return load_ground_truth_split(
        dataset_name=dataset_name,
        split_name=split_name,
        num_users=teacher_num_users,
        num_items=teacher_num_items,
        user_mapping=user_mapping,
        item_mapping=item_mapping,
        id_space=id_space,
    )

load_interaction_dataset(dataset_name: str, user_mapping: dict[int, int] | dict[str, int] | None = None, item_mapping: dict[int, int] | dict[str, int] | None = None, num_users: int | None = None, num_items: int | None = None, id_space: str | None = None) -> InteractionDataset

Source code in recdistill/data/datarec_loader.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
def load_interaction_dataset(
    dataset_name: str,
    user_mapping: dict[int, int] | dict[str, int] | None = None,
    item_mapping: dict[int, int] | dict[str, int] | None = None,
    num_users: int | None = None,
    num_items: int | None = None,
    id_space: str | None = None,
) -> InteractionDataset:
    if _normalize_id_space(id_space) == "dataset_integer":
        split, _ = _frame_to_split(
            load_split_frame(dataset_name, "train", use_datarec=False).frame,
            user_mapping=None,
            item_mapping=None,
            num_users=num_users,
            num_items=num_items,
        )
        if num_users is None:
            num_users = max(split.keys(), default=-1) + 1
        if num_items is None:
            num_items = max((item for items in split.values() for item in items), default=-1) + 1
    elif user_mapping is None and item_mapping is None:
        encoded = load_encoded_dataset(dataset_name)
        if num_users is None:
            num_users = encoded.num_users
        if num_items is None:
            num_items = encoded.num_items
        split, _ = _encoded_frame_to_split(encoded.frames["train"], num_users=num_users, num_items=num_items)
    else:
        split, _ = _frame_to_split(
            load_split_frame(dataset_name, "train").frame,
            user_mapping=user_mapping,
            item_mapping=item_mapping,
            num_users=num_users,
            num_items=num_items,
        )
        if num_users is None:
            num_users = max(split.keys(), default=-1) + 1
        if num_items is None:
            num_items = max((item for items in split.values() for item in items), default=-1) + 1
    return InteractionDataset.from_train_dict(split, num_users=num_users, num_items=num_items)

resolve_teacher_dataset_mappings(metadata: dict[str, Any] | None, *, dataset_name: str) -> tuple[dict[int, int] | dict[str, int] | None, dict[int, int] | dict[str, int] | None, str]

Source code in recdistill/data/datarec_loader.py 
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
def resolve_teacher_dataset_mappings(
    metadata: dict[str, Any] | None,
    *,
    dataset_name: str,
) -> tuple[dict[int, int] | dict[str, int] | None, dict[int, int] | dict[str, int] | None, str]:
    metadata = metadata or {}
    user_mapping = metadata.get("public_to_local_user_id")
    item_mapping = metadata.get("public_to_local_item_id")
    if user_mapping or item_mapping:
        return user_mapping, item_mapping, "teacher_metadata"

    id_space = _normalize_id_space(metadata.get("id_space"))
    if id_space == "dataset_integer":
        return None, None, "dataset_integer"
    if id_space in {"internal_integer", "framework_integer", "local_integer"}:
        raise ValueError(
            f"External teacher artifact declares id_space={id_space!r}, so its user/item IDs are framework-local "
            "indices rather than dataset-original IDs. Evaluation/training against raw dataset splits requires "
            "public_to_local_user_id and public_to_local_item_id metadata, or a prediction export converted back "
            "to the original user/item IDs."
        )

    framework = str(metadata.get("framework") or metadata.get("source") or "").strip().lower()
    source_framework = str(metadata.get("import_source_framework") or "").strip().lower()
    if framework == "external" or source_framework == "external":
        raise ValueError(
            "External teacher artifacts must define either public_to_local_user_id/public_to_local_item_id "
            "metadata or id_space='dataset_integer'. Refusing to evaluate with an implicit DataRec re-encoding."
        )

    if not user_mapping and not item_mapping:
        return None, None, "datarec"

load_ground_truth_split(dataset_name: str, split_name: str, num_users: int | None = None, num_items: int | None = None, user_mapping: dict[int, int] | dict[str, int] | None = None, item_mapping: dict[int, int] | dict[str, int] | None = None, id_space: str | None = None) -> tuple[dict[int, set[int]], int]

Source code in recdistill/data/datarec_loader.py 
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
def load_ground_truth_split(
    dataset_name: str,
    split_name: str,
    num_users: int | None = None,
    num_items: int | None = None,
    user_mapping: dict[int, int] | dict[str, int] | None = None,
    item_mapping: dict[int, int] | dict[str, int] | None = None,
    id_space: str | None = None,
) -> tuple[dict[int, set[int]], int]:
    if _normalize_id_space(id_space) == "dataset_integer":
        loaded = load_split_frame(dataset_name, split_name, use_datarec=False)
        split, dropped = _frame_to_split(
            loaded.frame,
            user_mapping=None,
            item_mapping=None,
            num_users=num_users,
            num_items=num_items,
        )
        if dropped == 0:
            return split, dropped
        encoded = load_encoded_dataset(dataset_name)
        normalized_split = "val" if split_name == "validation" else split_name
        if normalized_split not in encoded.frames:
            raise ValueError(f"Unknown split: {split_name}")
        return _encoded_frame_to_split(
            encoded.frames[normalized_split],
            num_users=num_users,
            num_items=num_items,
        )

    if user_mapping is None and item_mapping is None:
        encoded = load_encoded_dataset(dataset_name)
        normalized_split = "val" if split_name == "validation" else split_name
        if normalized_split not in encoded.frames:
            raise ValueError(f"Unknown split: {split_name}")
        return _encoded_frame_to_split(
            encoded.frames[normalized_split],
            num_users=num_users,
            num_items=num_items,
        )

    loaded = load_split_frame(dataset_name, split_name)
    return _frame_to_split(
        loaded.frame,
        user_mapping=user_mapping,
        item_mapping=item_mapping,
        num_users=num_users,
        num_items=num_items,
    )

load_encoded_dataset(dataset_name: str, *, repo_root: Path | str = '.') -> EncodedDataset

Load all splits and apply one shared DataRec user/item encoding.

Source code in recdistill/data/datarec_loader.py 
266
267
268
def load_encoded_dataset(dataset_name: str, *, repo_root: Path | str = ".") -> EncodedDataset:
    """Load all splits and apply one shared DataRec user/item encoding."""
    return _load_encoded_dataset_cached(dataset_name, str(Path(repo_root)))

resolve_local_id(raw_value: Any, mapping: dict[int, int] | dict[str, int] | None) -> int | None

Source code in recdistill/data/datarec_loader.py 
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
def resolve_local_id(raw_value: Any, mapping: dict[int, int] | dict[str, int] | None) -> int | None:
    candidates: list[Any] = [raw_value]
    raw_string = str(raw_value).strip()
    candidates.append(raw_string)
    try:
        parsed_int = int(raw_string)
    except (TypeError, ValueError):
        parsed_int = None
    if parsed_int is not None:
        candidates.extend([parsed_int, str(parsed_int)])

    if mapping is not None:
        mapped_value = None
        for candidate in candidates:
            if candidate in mapping:
                mapped_value = mapping[candidate]
                break
    else:
        mapped_value = parsed_int if parsed_int is not None else raw_string

    if mapped_value is None:
        return None
    try:
        return int(mapped_value)
    except (TypeError, ValueError):
        return None