| """ |
| PyTorch Dataset classes for the Two-Tower neural recommender. |
| |
| Provides: |
| - InteractionDataset — for training (user_idx, item_idx, label) |
| - InferenceDataset — for building item tower embeddings at inference time |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| from torch.utils.data import Dataset |
|
|
|
|
| class InteractionDataset(Dataset): |
| """ |
| Dataset of (user_idx, item_idx, label) triples for training the Two-Tower model. |
| |
| Optionally attaches item feature vectors when `item_features` is provided. |
| |
| Args: |
| df: DataFrame with columns [user_idx, item_idx, label]. |
| item_features: Optional numpy array of shape (n_items, n_item_features). |
| """ |
|
|
| def __init__( |
| self, |
| df: pd.DataFrame, |
| item_features: np.ndarray | None = None, |
| ) -> None: |
| self.user_idxs = torch.LongTensor(df["user_idx"].values) |
| self.item_idxs = torch.LongTensor(df["item_idx"].values) |
| self.labels = torch.FloatTensor(df["label"].values) |
| self.item_features = ( |
| torch.FloatTensor(item_features) if item_features is not None else None |
| ) |
|
|
| def __len__(self) -> int: |
| return len(self.labels) |
|
|
| def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: |
| sample: dict[str, torch.Tensor] = { |
| "user_idx": self.user_idxs[idx], |
| "item_idx": self.item_idxs[idx], |
| "label": self.labels[idx], |
| } |
| if self.item_features is not None: |
| sample["item_features"] = self.item_features[self.item_idxs[idx]] |
| return sample |
|
|
|
|
| class BPRDataset(Dataset): |
| """ |
| Bayesian Personalized Ranking dataset: (user, pos_item, neg_item) triples. |
| |
| For each positive interaction, a random negative item is paired. |
| If the DataFrame has pre-sampled negatives (label=0 rows from train split), |
| those are used. Otherwise (e.g., val split with only positives), random |
| item indices are sampled globally. |
| |
| Args: |
| df: DataFrame with columns [user_idx, item_idx, label]. |
| item_features: Optional numpy array of shape (n_items, n_item_features). |
| n_items: Total number of items (required when df has no negatives). |
| """ |
|
|
| def __init__( |
| self, |
| df: pd.DataFrame, |
| item_features: np.ndarray | None = None, |
| n_items: int | None = None, |
| ) -> None: |
| label_col = "label" if "label" in df.columns else None |
| pos_df = ( |
| df[df["label"] == 1].reset_index(drop=True) |
| if label_col |
| else df.reset_index(drop=True) |
| ) |
| neg_df = ( |
| df[df["label"] == 0].reset_index(drop=True) |
| if label_col |
| else pd.DataFrame() |
| ) |
|
|
| n_pos = len(pos_df) |
| rng = np.random.default_rng(42) |
| sampled_negs = np.empty(n_pos, dtype=np.int64) |
|
|
| if len(neg_df) > 0: |
| |
| |
| user_neg_map: dict[int, np.ndarray] = { |
| uid: grp["item_idx"].values |
| for uid, grp in neg_df.groupby("user_idx") |
| } |
| pos_users = pos_df["user_idx"].values |
| for i, uid in enumerate(pos_users): |
| neg_pool = user_neg_map.get(uid) |
| if neg_pool is not None and len(neg_pool) > 0: |
| sampled_negs[i] = neg_pool[rng.integers(0, len(neg_pool))] |
| else: |
| |
| assert n_items is not None |
| sampled_negs[i] = rng.integers(0, n_items) |
| else: |
| |
| assert n_items is not None, "n_items required when df has no label=0 rows" |
| sampled_negs = rng.integers(0, n_items, size=n_pos).astype(np.int64) |
|
|
| self.user_idxs = torch.LongTensor(pos_df["user_idx"].values) |
| self.pos_item_idxs = torch.LongTensor(pos_df["item_idx"].values) |
| self.neg_item_idxs = torch.LongTensor(sampled_negs) |
|
|
| |
| self.item_features = ( |
| torch.FloatTensor(item_features) if item_features is not None else None |
| ) |
|
|
| def __len__(self) -> int: |
| return len(self.user_idxs) |
|
|
| def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: |
| pos_iid = self.pos_item_idxs[idx] |
| neg_iid = self.neg_item_idxs[idx] |
| sample: dict[str, torch.Tensor] = { |
| "user_idx": self.user_idxs[idx], |
| "pos_item_idx": pos_iid, |
| "neg_item_idx": neg_iid, |
| } |
| if self.item_features is not None: |
| sample["pos_item_features"] = self.item_features[pos_iid] |
| sample["neg_item_features"] = self.item_features[neg_iid] |
| return sample |
|
|
|
|
| class InferenceDataset(Dataset): |
| """ |
| Dataset over all items for computing item tower embeddings in batch. |
| |
| Args: |
| n_items: Total number of items. |
| item_features: Optional numpy array of shape (n_items, n_item_features). |
| """ |
|
|
| def __init__( |
| self, |
| n_items: int, |
| item_features: np.ndarray | None = None, |
| ) -> None: |
| self.item_idxs = torch.arange(n_items, dtype=torch.long) |
| self.item_features = ( |
| torch.FloatTensor(item_features) if item_features is not None else None |
| ) |
|
|
| def __len__(self) -> int: |
| return len(self.item_idxs) |
|
|
| def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: |
| sample: dict[str, torch.Tensor] = {"item_idx": self.item_idxs[idx]} |
| if self.item_features is not None: |
| sample["item_features"] = self.item_features[idx] |
| return sample |
|
|
