File size: 2,684 Bytes
188f0cf | 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 | """PyTorch Dataset for BPR training with on-the-fly negative sampling.
Each item yields (user_idx, pos_item_idx, neg_item_idxs[K]) where `neg_item_idxs`
are uniformly sampled item indices that the user has NOT positively interacted
with in any split (train OR val OR test). Rejection sampling is fine here:
typical sparsity is |seen_u| / num_items ≪ 1%, so expected retries are ~0.01.
"""
from __future__ import annotations
from typing import Final
import numpy as np
import torch
from torch.utils.data import Dataset
_MAX_REJECTION_ATTEMPTS: Final[int] = 50 # safety bound; sparsity makes this fine
class BPRDataset(Dataset):
"""Yields (user, pos, negs) triples for pairwise ranking training."""
def __init__(
self,
train_pairs: np.ndarray,
user_positives: list[set[int]],
num_items: int,
num_negatives: int,
seed: int,
) -> None:
if train_pairs.ndim != 2 or train_pairs.shape[1] != 2:
raise ValueError(
f"train_pairs must have shape [N, 2], got {train_pairs.shape}"
)
if num_negatives < 1:
raise ValueError(f"num_negatives must be >= 1, got {num_negatives}")
self._pairs = train_pairs.astype(np.int64, copy=False)
self._user_positives = user_positives
self._num_items = int(num_items)
self._num_negatives = int(num_negatives)
# Per-worker RNGs are re-seeded in `worker_init_fn`. This is only used
# in single-worker mode; multi-worker overrides via worker_init_fn.
self._rng = np.random.default_rng(seed)
def __len__(self) -> int:
return self._pairs.shape[0]
def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
u, i_pos = self._pairs[idx]
seen = self._user_positives[int(u)]
negs = np.empty(self._num_negatives, dtype=np.int64)
for k in range(self._num_negatives):
negs[k] = self._sample_negative(seen)
return (
torch.from_numpy(np.asarray(u, dtype=np.int64)),
torch.from_numpy(np.asarray(i_pos, dtype=np.int64)),
torch.from_numpy(negs),
)
def _sample_negative(self, seen: set[int]) -> int:
"""Rejection-sample a single item index not in `seen`."""
for _ in range(_MAX_REJECTION_ATTEMPTS):
j = int(self._rng.integers(0, self._num_items))
if j not in seen:
return j
# Extremely rare fallback for pathological users who've seen nearly
# everything; just return any index and let training absorb the noise.
return int(self._rng.integers(0, self._num_items))
|