| """Common interface for every recommender model. |
| |
| Trainer and evaluator depend only on this class — concrete models (MF, |
| TwoTower, future GNNs) plug in without any changes upstream. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from abc import ABC, abstractmethod |
|
|
| import torch |
| from torch import Tensor, nn |
|
|
|
|
| class BaseRecommender(nn.Module, ABC): |
| """Abstract base. Subclasses implement `score` and `score_all_items`.""" |
|
|
| num_users: int |
| num_items: int |
|
|
| def __init__(self, num_users: int, num_items: int) -> None: |
| super().__init__() |
| if num_users < 1 or num_items < 1: |
| raise ValueError("num_users and num_items must be >= 1") |
| self.num_users = int(num_users) |
| self.num_items = int(num_items) |
|
|
| @abstractmethod |
| def score(self, users: Tensor, items: Tensor) -> Tensor: |
| """Score (user, item) pairs. |
| |
| Args: |
| users: int64 tensor of shape [B] or broadcastable to items. |
| items: int64 tensor of shape [B] or [B, K]. |
| |
| Returns: |
| float tensor with shape matching `items`. |
| """ |
|
|
| @abstractmethod |
| def score_all_items(self, users: Tensor) -> Tensor: |
| """Score a batch of users against every item in the catalog. |
| |
| Args: |
| users: int64 tensor of shape [B]. |
| |
| Returns: |
| float tensor of shape [B, num_items]. |
| """ |
|
|
| def forward( |
| self, users: Tensor, pos_items: Tensor, neg_items: Tensor |
| ) -> tuple[Tensor, Tensor]: |
| """Shared forward used by BPR training. |
| |
| Args: |
| users: [B] int64. |
| pos_items: [B] int64. |
| neg_items: [B, K] int64. |
| |
| Returns: |
| (pos_scores [B], neg_scores [B, K]). |
| """ |
| pos_scores = self.score(users, pos_items) |
| |
| users_expanded = users.unsqueeze(-1).expand_as(neg_items) |
| neg_scores = self.score(users_expanded, neg_items) |
| return pos_scores, neg_scores |
|
|
