Candidate relevance

Modeling a person’s interests

Ranking scores C candidates for each of B examples, so candidate_ids is [B, C]. The model itself doesn't change: flatten the candidate axis into the batch axis, score all B×C pairs in one call, then reshape the results back to [B, C]. Each candidate still gets its own history weights and interest vector. Each history has N positions, with ID 0 reserved for padding. The included model learns D-feature item vectors, scores each history item against its candidate, sums the weighted vectors, and maps that summary plus the candidate to a click logit. Sigmoid converts each logit to a probability. The file includes this entire model; no earlier example is required.

View the full file · rank_candidates.py

1import torch2from torch import nn3 4class TargetAwareEmbedding(nn.Module):5    def __init__(self, num_items, dim, padding_idx=0):6        super().__init__()7        self.padding_idx = padding_idx8        self.embedding = nn.Embedding(num_items, dim, padding_idx=padding_idx)9        self.score = nn.Sequential(10            nn.Linear(4 * dim, 32), nn.PReLU(), nn.Linear(32, 1)11        )12 13    def forward(self, history_ids, target_ids):14        # history_ids: [B, N]; target_ids: [B]. Targets must be real items.15        valid = history_ids.ne(self.padding_idx)        # [B, N]16        items = self.embedding(history_ids)             # [B, N, D]17        target = self.embedding(target_ids)             # [B, D]18        query = target[:, None, :].expand_as(items)     # [B, N, D]19        features = torch.cat(20            [items, query, items - query, items * query], dim=-121        )                                               # [B, N, 4D]22        weights = self.score(features).squeeze(-1)      # [B, N]23        weights = weights.masked_fill(~valid, 0.0)      # No softmax24        interest = (weights[..., None] * items).sum(dim=1)  # [B, D]25        return interest, target, weights26 27class ClickPredictor(nn.Module):28    def __init__(self, num_items, dim):29        super().__init__()30        self.history = TargetAwareEmbedding(num_items, dim)31        self.head = nn.Sequential(32            nn.Linear(2 * dim, 32), nn.ReLU(), nn.Linear(32, 1)33        )34 35    def forward(self, history_ids, target_ids):36        interest, target, weights = self.history(history_ids, target_ids)37        features = torch.cat([interest, target], dim=-1)  # [B, 2D]38        logits = self.head(features).squeeze(-1)          # [B]39        return logits, weights40 41if __name__ == "__main__":42    torch.manual_seed(0)43    model = ClickPredictor(num_items=6, dim=8).eval()44    history_ids = torch.tensor([[1, 2, 0], [3, 4, 5]])  # [B=2, N=3]45    candidate_ids = torch.tensor([[3, 4], [1, 2]])      # [B=2, C=2]46    batch, candidates = candidate_ids.shape47    length = history_ids.shape[1]48    histories = history_ids[:, None, :].expand(batch, candidates, length)49    flat_histories = histories.reshape(batch * candidates, length)50    flat_targets = candidate_ids.reshape(batch * candidates)51    with torch.no_grad():52        flat_logits, flat_weights = model(flat_histories, flat_targets)53        probabilities = flat_logits.sigmoid().reshape(batch, candidates)54        weights = flat_weights.reshape(batch, candidates, length)55    print("Candidate probabilities:", probabilities)  # [B, C]56    print("Candidate-specific history weights:", weights)  # [B, C, N]57    # Use a trained checkpoint for real ranking; this demonstrates batching.

Lines 44–47

Separate examples, candidates, and history items

history_ids [B, N] holds one history per example; candidate_ids [B, C] holds C candidates for each of those histories. With B=2 and C=2 the model scores four history/candidate pairs, not four history items. C is the ranking axis; N is the axis that pooling sums over.

42    torch.manual_seed(0)43    model = ClickPredictor(num_items=6, dim=8).eval()44    history_ids = torch.tensor([[1, 2, 0], [3, 4, 5]])  # [B=2, N=3]45    candidate_ids = torch.tensor([[3, 4], [1, 2]])      # [B=2, C=2]46    batch, candidates = candidate_ids.shape47    length = history_ids.shape[1]48    histories = history_ids[:, None, :].expand(batch, candidates, length)49    flat_histories = histories.reshape(batch * candidates, length)

Quick checkB = 2 histories of N = 4 items, with C = 3 candidates each. How many history/candidate pairs does the model score?

rank_candidates.py

1import torch2from torch import nn3 4class TargetAwareEmbedding(nn.Module):5    def __init__(self, num_items, dim, padding_idx=0):6        super().__init__()7        self.padding_idx = padding_idx8        self.embedding = nn.Embedding(num_items, dim, padding_idx=padding_idx)9        self.score = nn.Sequential(10            nn.Linear(4 * dim, 32), nn.PReLU(), nn.Linear(32, 1)11        )12 13    def forward(self, history_ids, target_ids):14        # history_ids: [B, N]; target_ids: [B]. Targets must be real items.15        valid = history_ids.ne(self.padding_idx)        # [B, N]16        items = self.embedding(history_ids)             # [B, N, D]17        target = self.embedding(target_ids)             # [B, D]18        query = target[:, None, :].expand_as(items)     # [B, N, D]19        features = torch.cat(20            [items, query, items - query, items * query], dim=-121        )                                               # [B, N, 4D]22        weights = self.score(features).squeeze(-1)      # [B, N]23        weights = weights.masked_fill(~valid, 0.0)      # No softmax24        interest = (weights[..., None] * items).sum(dim=1)  # [B, D]25        return interest, target, weights26 27class ClickPredictor(nn.Module):28    def __init__(self, num_items, dim):29        super().__init__()30        self.history = TargetAwareEmbedding(num_items, dim)31        self.head = nn.Sequential(32            nn.Linear(2 * dim, 32), nn.ReLU(), nn.Linear(32, 1)33        )34 35    def forward(self, history_ids, target_ids):36        interest, target, weights = self.history(history_ids, target_ids)37        features = torch.cat([interest, target], dim=-1)  # [B, 2D]38        logits = self.head(features).squeeze(-1)          # [B]39        return logits, weights40 41if __name__ == "__main__":42    torch.manual_seed(0)43    model = ClickPredictor(num_items=6, dim=8).eval()44    history_ids = torch.tensor([[1, 2, 0], [3, 4, 5]])  # [B=2, N=3]45    candidate_ids = torch.tensor([[3, 4], [1, 2]])      # [B=2, C=2]46    batch, candidates = candidate_ids.shape47    length = history_ids.shape[1]48    histories = history_ids[:, None, :].expand(batch, candidates, length)49    flat_histories = histories.reshape(batch * candidates, length)50    flat_targets = candidate_ids.reshape(batch * candidates)51    with torch.no_grad():52        flat_logits, flat_weights = model(flat_histories, flat_targets)53        probabilities = flat_logits.sigmoid().reshape(batch, candidates)54        weights = flat_weights.reshape(batch, candidates, length)55    print("Candidate probabilities:", probabilities)  # [B, C]56    print("Candidate-specific history weights:", weights)  # [B, C, N]57    # Use a trained checkpoint for real ranking; this demonstrates batching.