models/two_tower.py

"""
Two-Tower (Dual Encoder) Model for Pinterest Personalized Retrieval.

Architecture:
  UserTower: user_features → MLP → L2-normalized embedding (d)
  ItemTower: item_features → MLP → L2-normalized embedding (d)

Training objective:
  InfoNCE / in-batch softmax contrastive loss.
  For a batch of size B, each (user_i, item_i) pair is positive;
  all other items in the batch serve as negatives.

At inference:
  - Pre-compute item embeddings → build FAISS index
  - For a query user → encode → ANN search in FAISS
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from loguru import logger


# ─── MLP Tower ───────────────────────────────────────────────────────────────

class MLP(nn.Module):
    """Shared MLP backbone used by both towers."""

    def __init__(
        self,
        input_dim: int,
        hidden_dims: list[int],
        output_dim: int,
        dropout: float = 0.2,
    ):
        super().__init__()
        layers = []
        in_dim = input_dim
        for h in hidden_dims:
            layers += [
                nn.Linear(in_dim, h),
                nn.LayerNorm(h),
                nn.GELU(),
                nn.Dropout(dropout),
            ]
            in_dim = h
        layers.append(nn.Linear(in_dim, output_dim))
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


# ─── Two-Tower Model ─────────────────────────────────────────────────────────

class TwoTowerModel(nn.Module):
    """
    Dual encoder producing L2-normalized embeddings in a shared space.

    Args:
        user_feat_dim  : dimension of raw user feature vector
        item_feat_dim  : dimension of raw item feature vector
        embedding_dim  : shared embedding space dimension
        hidden_dims    : list of hidden layer sizes for each MLP
        dropout        : dropout probability
        temperature    : InfoNCE softmax temperature
    """

    def __init__(
        self,
        user_feat_dim: int,
        item_feat_dim: int,
        embedding_dim: int = 64,
        hidden_dims: list[int] = [256, 128],
        dropout: float = 0.2,
        temperature: float = 0.07,
    ):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.temperature = nn.Parameter(
            torch.tensor(temperature), requires_grad=True
        )

        self.user_tower = MLP(user_feat_dim, hidden_dims, embedding_dim, dropout)
        self.item_tower = MLP(item_feat_dim, hidden_dims, embedding_dim, dropout)

        self._init_weights()
        logger.info(
            f"TwoTowerModel | user_dim={user_feat_dim} → item_dim={item_feat_dim} "
            f"→ embed_dim={embedding_dim} | temp={temperature}"
        )

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def encode_users(self, user_feats: torch.Tensor) -> torch.Tensor:
        """Returns L2-normalized user embeddings. Shape: (B, D)"""
        return F.normalize(self.user_tower(user_feats), p=2, dim=-1)

    def encode_items(self, item_feats: torch.Tensor) -> torch.Tensor:
        """Returns L2-normalized item embeddings. Shape: (B, D)"""
        return F.normalize(self.item_tower(item_feats), p=2, dim=-1)

    def forward(
        self,
        user_feats: torch.Tensor,
        item_feats: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Returns (user_embeddings, item_embeddings) — both L2-normalized.
        Loss is computed externally for flexibility.
        """
        u_emb = self.encode_users(user_feats)
        i_emb = self.encode_items(item_feats)
        return u_emb, i_emb

    def num_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


# ─── Loss Functions ───────────────────────────────────────────────────────────

class InfoNCELoss(nn.Module):
    """
    In-batch InfoNCE (NT-Xent) contrastive loss.

    Given a batch of (user, pos_item) pairs, each pos_item is the positive
    for its user and all other items in the batch act as negatives.

    Loss = -mean(log softmax(sim(u_i, i_i) / τ))
    """

    def __init__(self):
        super().__init__()

    def forward(
        self,
        user_emb: torch.Tensor,   # (B, D)
        item_emb: torch.Tensor,   # (B, D)
        temperature: torch.Tensor,
        weights: torch.Tensor | None = None,  # (B,) optional interaction weights
    ) -> torch.Tensor:
        # Clamp temperature to avoid numerical instability
        temp = torch.clamp(temperature, min=0.01, max=1.0)

        # Similarity matrix (B, B): rows=users, cols=items
        logits = torch.matmul(user_emb, item_emb.T) / temp  # (B, B)

        # Diagonal entries are positives
        labels = torch.arange(logits.size(0), device=logits.device)

        if weights is not None:
            # Weight the loss by interaction quality (saves > clicks)
            loss_per_sample = F.cross_entropy(logits, labels, reduction="none")
            w = weights / weights.sum()
            return (loss_per_sample * w).sum()
        else:
            return F.cross_entropy(logits, labels)


class HardNegativeMiner:
    """
    Semi-hard negative mining within a batch.

    Identifies negatives that are closer to the query than the positive
    but not so hard they dominate training.
    """

    def __init__(self, hard_ratio: float = 0.3):
        self.hard_ratio = hard_ratio

    @torch.no_grad()
    def mine(
        self,
        user_emb: torch.Tensor,  # (B, D)
        item_emb: torch.Tensor,  # (B, D)
        temperature: float = 0.07,
    ) -> torch.Tensor:
        """
        Returns a permuted item_emb where hard negatives are brought
        closer to the front of each row's negative list.
        """
        B = user_emb.size(0)
        sim = torch.matmul(user_emb, item_emb.T)  # (B, B)

        # For each user, rank non-positive items by similarity (desc)
        pos_sim = sim.diag().unsqueeze(1)   # (B, 1)
        mask = torch.eye(B, dtype=torch.bool, device=sim.device)
        sim_neg = sim.masked_fill(mask, -1e9)

        # Hard negatives: sim_neg > pos_sim (violators of margin)
        hard_mask = sim_neg > pos_sim
        n_hard = int(B * self.hard_ratio)

        # Sort negatives: hard first, then easy
        _, sorted_idx = sim_neg.sort(dim=1, descending=True)

        # Re-index item embeddings to surface hard negatives per query
        # (used for logging/analysis; actual loss uses full in-batch)
        hard_neg_idx = sorted_idx[:, :n_hard]  # (B, n_hard)

        return hard_neg_idx, hard_mask.sum().item()


# ─── Model Factory ────────────────────────────────────────────────────────────

def build_model(cfg: dict, user_feat_dim: int, item_feat_dim: int) -> TwoTowerModel:
    mc = cfg["model"]
    return TwoTowerModel(
        user_feat_dim=user_feat_dim,
        item_feat_dim=item_feat_dim,
        embedding_dim=mc["embedding_dim"],
        hidden_dims=mc["hidden_dims"],
        dropout=mc["dropout"],
        temperature=mc["temperature"],
    )