nextrec 0.1.11__py3-none-any.whl → 0.2.2__py3-none-any.whl

nextrec/loss/__init__.py

@@ -1,35 +1,42 @@
-from nextrec.loss.match_losses import (
-    BPRLoss,
-    HingeLoss,
-    TripletLoss,
-    SampledSoftmaxLoss,
-    CosineContrastiveLoss,
+from nextrec.loss.listwise import (
+    ApproxNDCGLoss,
     InfoNCELoss,
-    ListNetLoss,
     ListMLELoss,
-    ApproxNDCGLoss,
+    ListNetLoss,
+    SampledSoftmaxLoss,
+)
+from nextrec.loss.pairwise import BPRLoss, HingeLoss, TripletLoss
+from nextrec.loss.pointwise import (
+    ClassBalancedFocalLoss,
+    CosineContrastiveLoss,
+    FocalLoss,
+    WeightedBCELoss,
 )
-
 from nextrec.loss.loss_utils import (
     get_loss_fn,
-    validate_training_mode,
+    get_loss_kwargs,
     VALID_TASK_TYPES,
 )
 
 __all__ = [
-    # Match losses
-    'BPRLoss',
-    'HingeLoss',
-    'TripletLoss',
-    'SampledSoftmaxLoss',
-    'CosineContrastiveLoss',
-    'InfoNCELoss',
-    # Listwise losses
-    'ListNetLoss',
-    'ListMLELoss',
-    'ApproxNDCGLoss',
+    # Pointwise
+    "CosineContrastiveLoss",
+    "WeightedBCELoss",
+    "FocalLoss",
+    "ClassBalancedFocalLoss",
+    # Pairwise
+    "BPRLoss",
+    "HingeLoss",
+    "TripletLoss",
+    # Listwise
+    "SampledSoftmaxLoss",
+    "InfoNCELoss",
+    "ListNetLoss",
+    "ListMLELoss",
+    "ApproxNDCGLoss",
     # Utilities
-    'get_loss_fn',
-    'validate_training_mode',
-    'VALID_TASK_TYPES',
+    "get_loss_fn",
+    "get_loss_kwargs",
+    "validate_training_mode",
+    "VALID_TASK_TYPES",
 ]

nextrec/loss/listwise.py

@@ -0,0 +1,164 @@
+"""
+Listwise loss functions for ranking and contrastive training.
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SampledSoftmaxLoss(nn.Module):
+    """
+    Softmax over one positive and multiple sampled negatives.
+    """
+
+    def __init__(self, reduction: str = "mean"):
+        super().__init__()
+        self.reduction = reduction
+
+    def forward(self, pos_logits: torch.Tensor, neg_logits: torch.Tensor) -> torch.Tensor:
+        pos_logits = pos_logits.unsqueeze(1)
+        all_logits = torch.cat([pos_logits, neg_logits], dim=1)
+        targets = torch.zeros(all_logits.size(0), dtype=torch.long, device=all_logits.device)
+        loss = F.cross_entropy(all_logits, targets, reduction=self.reduction)
+        return loss
+
+
+class InfoNCELoss(nn.Module):
+    """
+    InfoNCE loss for contrastive learning with one positive and many negatives.
+    """
+
+    def __init__(self, temperature: float = 0.07, reduction: str = "mean"):
+        super().__init__()
+        self.temperature = temperature
+        self.reduction = reduction
+
+    def forward(
+        self, query: torch.Tensor, pos_key: torch.Tensor, neg_keys: torch.Tensor
+    ) -> torch.Tensor:
+        pos_sim = torch.sum(query * pos_key, dim=-1) / self.temperature
+        pos_sim = pos_sim.unsqueeze(1)
+        query_expanded = query.unsqueeze(1)
+        neg_sim = torch.sum(query_expanded * neg_keys, dim=-1) / self.temperature
+        logits = torch.cat([pos_sim, neg_sim], dim=1)
+        labels = torch.zeros(logits.size(0), dtype=torch.long, device=logits.device)
+        loss = F.cross_entropy(logits, labels, reduction=self.reduction)
+        return loss
+
+
+class ListNetLoss(nn.Module):
+    """
+    ListNet loss using top-1 probability distribution.
+    Reference: Cao et al. (ICML 2007)
+    """
+
+    def __init__(self, temperature: float = 1.0, reduction: str = "mean"):
+        super().__init__()
+        self.temperature = temperature
+        self.reduction = reduction
+
+    def forward(self, scores: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        pred_probs = F.softmax(scores / self.temperature, dim=1)
+        true_probs = F.softmax(labels / self.temperature, dim=1)
+        loss = -torch.sum(true_probs * torch.log(pred_probs + 1e-10), dim=1)
+
+        if self.reduction == "mean":
+            return loss.mean()
+        if self.reduction == "sum":
+            return loss.sum()
+        return loss
+
+
+class ListMLELoss(nn.Module):
+    """
+    ListMLE (Maximum Likelihood Estimation) loss.
+    Reference: Xia et al. (ICML 2008)
+    """
+
+    def __init__(self, reduction: str = "mean"):
+        super().__init__()
+        self.reduction = reduction
+
+    def forward(self, scores: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        sorted_labels, sorted_indices = torch.sort(labels, descending=True, dim=1)
+        batch_size, list_size = scores.shape
+        batch_indices = torch.arange(batch_size, device=scores.device).unsqueeze(1).expand(-1, list_size)
+        sorted_scores = scores[batch_indices, sorted_indices]
+
+        loss = torch.tensor(0.0, device=scores.device)
+        for i in range(list_size):
+            remaining_scores = sorted_scores[:, i:]
+            log_sum_exp = torch.logsumexp(remaining_scores, dim=1)
+            loss = loss + (log_sum_exp - sorted_scores[:, i]).sum()
+
+        if self.reduction == "mean":
+            return loss / batch_size
+        if self.reduction == "sum":
+            return loss
+        return loss / batch_size
+
+
+class ApproxNDCGLoss(nn.Module):
+    """
+    Approximate NDCG loss for learning to rank.
+    Reference: Qin et al. (2010)
+    """
+
+    def __init__(self, temperature: float = 1.0, reduction: str = "mean"):
+        super().__init__()
+        self.temperature = temperature
+        self.reduction = reduction
+
+    def _ideal_dcg(self, labels: torch.Tensor, k: Optional[int]) -> torch.Tensor:
+        # labels: [B, L]
+        sorted_labels, _ = torch.sort(labels, dim=1, descending=True)
+        if k is not None:
+            sorted_labels = sorted_labels[:, :k]
+
+        gains = torch.pow(2.0, sorted_labels) - 1.0  # [B, K]
+        positions = torch.arange(
+            1, gains.size(1) + 1, device=gains.device, dtype=torch.float32
+        )  # [K]
+        discounts = 1.0 / torch.log2(positions + 1.0)  # [K]
+        ideal_dcg = torch.sum(gains * discounts, dim=1)  # [B]
+        return ideal_dcg
+
+    def forward(
+        self, scores: torch.Tensor, labels: torch.Tensor, k: Optional[int] = None
+    ) -> torch.Tensor:
+        """
+        scores: [B, L]
+        labels: [B, L]
+        """
+        batch_size, list_size = scores.shape
+        device = scores.device
+
+        # diff[b, i, j] = (s_j - s_i) / T
+        scores_i = scores.unsqueeze(2)          # [B, L, 1]
+        scores_j = scores.unsqueeze(1)          # [B, 1, L]
+        diff = (scores_j - scores_i) / self.temperature  # [B, L, L]
+
+        P_ji = torch.sigmoid(diff)             # [B, L, L]
+        eye = torch.eye(list_size, device=device).unsqueeze(0)  # [1, L, L]
+        P_ji = P_ji * (1.0 - eye)
+
+        exp_rank = 1.0 + P_ji.sum(dim=-1)      # [B, L]
+
+        discounts = 1.0 / torch.log2(exp_rank + 1.0)  # [B, L]
+
+        gains = torch.pow(2.0, labels) - 1.0   # [B, L]
+        approx_dcg = torch.sum(gains * discounts, dim=1)  # [B]
+
+        ideal_dcg = self._ideal_dcg(labels, k)  # [B]
+
+        ndcg = approx_dcg / (ideal_dcg + 1e-10)  # [B]
+        loss = 1.0 - ndcg
+
+        if self.reduction == "mean":
+            return loss.mean()
+        if self.reduction == "sum":
+            return loss.sum()
+        return loss

nextrec/loss/loss_utils.py

@@ -1,136 +1,163 @@
 """
-Loss utilities for NextRec
-
-Date: create on 09/11/2025
-Author:
-    Yang Zhou,zyaztec@gmail.com
+Loss utilities for NextRec.
 """
-import torch
-import torch.nn as nn
+
 from typing import Literal
 
-from nextrec.loss.match_losses import (
-    BPRLoss, 
-    HingeLoss, 
-    TripletLoss, 
+import torch.nn as nn
+
+from nextrec.loss.listwise import (
+    ApproxNDCGLoss,
+    InfoNCELoss,
+    ListMLELoss,
+    ListNetLoss,
     SampledSoftmaxLoss,
-    CosineContrastiveLoss, 
-    InfoNCELoss
+)
+from nextrec.loss.pairwise import BPRLoss, HingeLoss, TripletLoss
+from nextrec.loss.pointwise import (
+    ClassBalancedFocalLoss,
+    CosineContrastiveLoss,
+    FocalLoss,
+    WeightedBCELoss,
 )
 
 # Valid task types for validation
-VALID_TASK_TYPES = ['binary', 'multiclass', 'regression', 'multivariate_regression', 'match', 'ranking', 'multitask', 'multilabel']
+VALID_TASK_TYPES = [
+    "binary",
+    "multiclass",
+    "regression",
+    "multivariate_regression",
+    "match",
+    "ranking",
+    "multitask",
+    "multilabel",
+]
 
 
 def get_loss_fn(
     task_type: str = "binary",
     training_mode: str | None = None,
     loss: str | nn.Module | None = None,
-    **loss_kwargs
+    **loss_kwargs,
 ) -> nn.Module:
     """
     Get loss function based on task type and training mode.
-    
-    Examples:
-        # Ranking task (binary classification)
-        >>> loss_fn = get_loss_fn(task_type="binary", loss="bce")
-        
-        # Match task with pointwise training
-        >>> loss_fn = get_loss_fn(task_type="match", training_mode="pointwise")
-        
-        # Match task with pairwise training
-        >>> loss_fn = get_loss_fn(task_type="match", training_mode="pairwise", loss="bpr")
-        
-        # Match task with listwise training
-        >>> loss_fn = get_loss_fn(task_type="match", training_mode="listwise", loss="sampled_softmax")
     """
 
     if isinstance(loss, nn.Module):
         return loss
 
+    # Common mappings
     if task_type == "match":
-        if training_mode == "pointwise":
-            # Pointwise training uses binary cross entropy
-            if loss is None or loss == "bce" or loss == "binary_crossentropy":
-                return nn.BCELoss(**loss_kwargs)
-            elif loss == "cosine_contrastive":
-                return CosineContrastiveLoss(**loss_kwargs)
-            elif isinstance(loss, str):
-                raise ValueError(f"Unsupported pointwise loss: {loss}")
-        
-        elif training_mode == "pairwise":
-            if loss is None or loss == "bpr":
-                return BPRLoss(**loss_kwargs)
-            elif loss == "hinge":
-                return HingeLoss(**loss_kwargs)
-            elif loss == "triplet":
-                return TripletLoss(**loss_kwargs)
-            elif isinstance(loss, str):
-                raise ValueError(f"Unsupported pairwise loss: {loss}")
-        
-        elif training_mode == "listwise":
-            if loss is None or loss == "sampled_softmax" or loss == "softmax":
-                return SampledSoftmaxLoss(**loss_kwargs)
-            elif loss == "infonce":
-                return InfoNCELoss(**loss_kwargs)
-            elif loss == "crossentropy" or loss == "ce":
-                return nn.CrossEntropyLoss(**loss_kwargs)
-            elif isinstance(loss, str):
-                raise ValueError(f"Unsupported listwise loss: {loss}")
-        
-        else:
-            raise ValueError(f"Unknown training_mode: {training_mode}")
-
-    elif task_type in ["ranking", "multitask", "binary"]:
-        if loss is None or loss == "bce" or loss == "binary_crossentropy":
-            return nn.BCELoss(**loss_kwargs)
-        elif loss == "mse":
-            return nn.MSELoss(**loss_kwargs)
-        elif loss == "mae":
-            return nn.L1Loss(**loss_kwargs)
-        elif loss == "crossentropy" or loss == "ce":
-            return nn.CrossEntropyLoss(**loss_kwargs)
-        elif isinstance(loss, str):
-            raise ValueError(f"Unsupported loss function: {loss}")
+        return _get_match_loss(training_mode, loss, **loss_kwargs)
 
-    elif task_type == "multiclass":
-        if loss is None or loss == "crossentropy" or loss == "ce":
-            return nn.CrossEntropyLoss(**loss_kwargs)
-        elif isinstance(loss, str):
-            raise ValueError(f"Unsupported multiclass loss: {loss}")
-    
-    elif task_type == "regression":
+    if task_type in ["ranking", "multitask", "binary", "multilabel"]:
+        return _get_classification_loss(loss, **loss_kwargs)
+
+    if task_type == "multiclass":
+        return _get_multiclass_loss(loss, **loss_kwargs)
+
+    if task_type == "regression":
         if loss is None or loss == "mse":
             return nn.MSELoss(**loss_kwargs)
-        elif loss == "mae":
+        if loss == "mae":
             return nn.L1Loss(**loss_kwargs)
-        elif isinstance(loss, str):
+        if isinstance(loss, str):
             raise ValueError(f"Unsupported regression loss: {loss}")
-    
-    else:
-        raise ValueError(f"Unsupported task_type: {task_type}")
-    
-    return loss
-
-
-def validate_training_mode(
-    training_mode: str,
-    support_training_modes: list[str],
-    model_name: str = "Model"
-) -> None:
+
+    raise ValueError(f"Unsupported task_type: {task_type}")
+
+
+def _get_match_loss(training_mode: str | None, loss: str | None, **loss_kwargs) -> nn.Module:
+    if training_mode == "pointwise":
+        if loss is None or loss in {"bce", "binary_crossentropy"}:
+            return nn.BCELoss(**loss_kwargs)
+        if loss == "weighted_bce":
+            return WeightedBCELoss(**loss_kwargs)
+        if loss == "focal":
+            return FocalLoss(**loss_kwargs)
+        if loss == "class_balanced_focal":
+            return _build_cb_focal(loss_kwargs)
+        if loss == "cosine_contrastive":
+            return CosineContrastiveLoss(**loss_kwargs)
+        if isinstance(loss, str):
+            raise ValueError(f"Unsupported pointwise loss: {loss}")
+
+    if training_mode == "pairwise":
+        if loss is None or loss == "bpr":
+            return BPRLoss(**loss_kwargs)
+        if loss == "hinge":
+            return HingeLoss(**loss_kwargs)
+        if loss == "triplet":
+            return TripletLoss(**loss_kwargs)
+        if isinstance(loss, str):
+            raise ValueError(f"Unsupported pairwise loss: {loss}")
+
+    if training_mode == "listwise":
+        if loss is None or loss in {"sampled_softmax", "softmax"}:
+            return SampledSoftmaxLoss(**loss_kwargs)
+        if loss == "infonce":
+            return InfoNCELoss(**loss_kwargs)
+        if loss == "listnet":
+            return ListNetLoss(**loss_kwargs)
+        if loss == "listmle":
+            return ListMLELoss(**loss_kwargs)
+        if loss == "approx_ndcg":
+            return ApproxNDCGLoss(**loss_kwargs)
+        if loss in {"crossentropy", "ce"}:
+            return nn.CrossEntropyLoss(**loss_kwargs)
+        if isinstance(loss, str):
+            raise ValueError(f"Unsupported listwise loss: {loss}")
+
+    raise ValueError(f"Unknown training_mode: {training_mode}")
+
+
+def _get_classification_loss(loss: str | None, **loss_kwargs) -> nn.Module:
+    if loss is None or loss in {"bce", "binary_crossentropy"}:
+        return nn.BCELoss(**loss_kwargs)
+    if loss == "weighted_bce":
+        return WeightedBCELoss(**loss_kwargs)
+    if loss == "focal":
+        return FocalLoss(**loss_kwargs)
+    if loss == "class_balanced_focal":
+        return _build_cb_focal(loss_kwargs)
+    if loss == "mse":
+        return nn.MSELoss(**loss_kwargs)
+    if loss == "mae":
+        return nn.L1Loss(**loss_kwargs)
+    if loss in {"crossentropy", "ce"}:
+        return nn.CrossEntropyLoss(**loss_kwargs)
+    if isinstance(loss, str):
+        raise ValueError(f"Unsupported loss function: {loss}")
+    raise ValueError("Loss must be specified for classification task.")
+
+
+def _get_multiclass_loss(loss: str | None, **loss_kwargs) -> nn.Module:
+    if loss is None or loss in {"crossentropy", "ce"}:
+        return nn.CrossEntropyLoss(**loss_kwargs)
+    if loss == "focal":
+        return FocalLoss(**loss_kwargs)
+    if loss == "class_balanced_focal":
+        return _build_cb_focal(loss_kwargs)
+    if isinstance(loss, str):
+        raise ValueError(f"Unsupported multiclass loss: {loss}")
+    raise ValueError("Loss must be specified for multiclass task.")
+
+
+def _build_cb_focal(loss_kwargs: dict) -> ClassBalancedFocalLoss:
+    if "class_counts" not in loss_kwargs:
+        raise ValueError("class_balanced_focal requires `class_counts` argument.")
+    return ClassBalancedFocalLoss(**loss_kwargs)
+
+
+def get_loss_kwargs(loss_params: dict | list[dict] | None, index: int = 0) -> dict:
     """
-    Validate that the requested training mode is supported by the model.
-    
-    Args:
-        training_mode: Requested training mode
-        support_training_modes: List of supported training modes
-        model_name: Name of the model (for error messages)
-    
-    Raises:
-        ValueError: If training mode is not supported
+    Resolve per-task loss kwargs from a dict or list of dicts.
     """
-    if training_mode not in support_training_modes:
-        raise ValueError(
-            f"{model_name} does not support training_mode='{training_mode}'. "
-            f"Supported modes: {support_training_modes}"
-        )
+    if loss_params is None:
+        return {}
+    if isinstance(loss_params, list):
+        if index < len(loss_params) and loss_params[index] is not None:
+            return loss_params[index]
+        return {}
+    return loss_params

nextrec/loss/pairwise.py

@@ -0,0 +1,105 @@
+"""
+Pairwise loss functions for learning-to-rank and matching tasks.
+"""
+
+from typing import Literal
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BPRLoss(nn.Module):
+    """
+    Bayesian Personalized Ranking loss with support for multiple negatives.
+    """
+
+    def __init__(self, reduction: str = "mean"):
+        super().__init__()
+        self.reduction = reduction
+
+    def forward(self, pos_score: torch.Tensor, neg_score: torch.Tensor) -> torch.Tensor:
+        if neg_score.dim() == 2:
+            pos_score = pos_score.unsqueeze(1)
+            diff = pos_score - neg_score
+        else:
+            diff = pos_score - neg_score
+
+        loss = -torch.log(torch.sigmoid(diff) + 1e-8)
+        if self.reduction == "mean":
+            return loss.mean()
+        if self.reduction == "sum":
+            return loss.sum()
+        return loss
+
+
+class HingeLoss(nn.Module):
+    """
+    Hinge loss for pairwise ranking.
+    """
+
+    def __init__(self, margin: float = 1.0, reduction: str = "mean"):
+        super().__init__()
+        self.margin = margin
+        self.reduction = reduction
+
+    def forward(self, pos_score: torch.Tensor, neg_score: torch.Tensor) -> torch.Tensor:
+        if neg_score.dim() == 2:
+            pos_score = pos_score.unsqueeze(1)
+
+        diff = pos_score - neg_score
+        loss = torch.clamp(self.margin - diff, min=0)
+
+        if self.reduction == "mean":
+            return loss.mean()
+        if self.reduction == "sum":
+            return loss.sum()
+        return loss
+
+
+class TripletLoss(nn.Module):
+    """
+    Triplet margin loss with cosine or euclidean distance.
+    """
+
+    def __init__(
+        self,
+        margin: float = 1.0,
+        reduction: str = "mean",
+        distance: Literal["euclidean", "cosine"] = "euclidean",
+    ):
+        super().__init__()
+        self.margin = margin
+        self.reduction = reduction
+        self.distance = distance
+
+    def forward(
+        self, anchor: torch.Tensor, positive: torch.Tensor, negative: torch.Tensor
+    ) -> torch.Tensor:
+        if self.distance == "euclidean":
+            pos_dist = torch.sum((anchor - positive) ** 2, dim=-1)
+            if negative.dim() == 3:
+                anchor_expanded = anchor.unsqueeze(1)
+                neg_dist = torch.sum((anchor_expanded - negative) ** 2, dim=-1)
+            else:
+                neg_dist = torch.sum((anchor - negative) ** 2, dim=-1)
+            if neg_dist.dim() == 2:
+                pos_dist = pos_dist.unsqueeze(1)
+        elif self.distance == "cosine":
+            pos_dist = 1 - F.cosine_similarity(anchor, positive, dim=-1)
+            if negative.dim() == 3:
+                anchor_expanded = anchor.unsqueeze(1)
+                neg_dist = 1 - F.cosine_similarity(anchor_expanded, negative, dim=-1)
+            else:
+                neg_dist = 1 - F.cosine_similarity(anchor, negative, dim=-1)
+            if neg_dist.dim() == 2:
+                pos_dist = pos_dist.unsqueeze(1)
+        else:
+            raise ValueError(f"Unsupported distance: {self.distance}")
+
+        loss = torch.clamp(pos_dist - neg_dist + self.margin, min=0)
+        if self.reduction == "mean":
+            return loss.mean()
+        if self.reduction == "sum":
+            return loss.sum()
+        return loss