File size: 5,879 Bytes
fef33ce 2647d97 fef33ce | 1 2 3 4 5 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 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 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 90 91 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 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 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 | """
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:
# Per-user pairing: for each user's positive, sample from THAT user's negatives.
# This ensures user-specific gradients (critical for non-collapsed user tower).
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:
# Fallback: global random item
assert n_items is not None
sampled_negs[i] = rng.integers(0, n_items)
else:
# No pre-sampled negatives (val/test split) — sample globally
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)
# Store full item feature matrix for indexed lookup
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
|