nextrec 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

nextrec/loss/loss_utils.py

@@ -1,135 +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

@@ -1,6 +1,105 @@
 """
-Loss functions for pairwise tasks
-
-Date: create on 22/11/2025
-Author: Yang Zhou,zyaztec@gmail.com
+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

nextrec/loss/pointwise.py

@@ -1,6 +1,198 @@
 """
-Loss functions for pointwise tasks
-
-Date: create on 22/11/2025
-Author: Yang Zhou,zyaztec@gmail.com
+Pointwise loss functions, including imbalance-aware variants.
 """
+
+from typing import Optional, Sequence
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class CosineContrastiveLoss(nn.Module):
+    """
+    Contrastive loss using cosine similarity for positive/negative pairs.
+    """
+
+    def __init__(self, margin: float = 0.5, reduction: str = "mean"):
+        super().__init__()
+        self.margin = margin
+        self.reduction = reduction
+
+    def forward(
+        self, user_emb: torch.Tensor, item_emb: torch.Tensor, labels: torch.Tensor
+    ) -> torch.Tensor:
+        labels = labels.float()
+        similarity = F.cosine_similarity(user_emb, item_emb, dim=-1)
+        pos_loss = torch.clamp(self.margin - similarity, min=0) * labels
+        neg_loss = torch.clamp(similarity - self.margin, min=0) * (1 - labels)
+        loss = pos_loss + neg_loss
+
+        if self.reduction == "mean":
+            return loss.mean()
+        if self.reduction == "sum":
+            return loss.sum()
+        return loss
+
+
+class WeightedBCELoss(nn.Module):
+    """
+    Binary cross entropy with controllable positive class weight.
+    Supports probability or logit inputs via `logits` flag.
+    If `auto_balance=True` and `pos_weight` is None, the positive weight is
+    computed from the batch as (#neg / #pos) for stable imbalance handling.
+    """
+    def __init__(
+        self,
+        pos_weight: float | torch.Tensor | None = None,
+        reduction: str = "mean",
+        logits: bool = False,
+        auto_balance: bool = False,
+    ):
+        super().__init__()
+        self.reduction = reduction
+        self.logits = logits
+        self.auto_balance = auto_balance
+
+        if pos_weight is not None:
+            self.register_buffer(
+                "pos_weight",
+                torch.as_tensor(pos_weight, dtype=torch.float32),
+            )
+        else:
+            self.pos_weight = None
+
+    def _resolve_pos_weight(self, labels: torch.Tensor) -> torch.Tensor:
+        if self.pos_weight is not None:
+            return self.pos_weight.to(device=labels.device)
+
+        if not self.auto_balance:
+            return torch.tensor(1.0, device=labels.device, dtype=labels.dtype)
+
+        labels_float = labels.float()
+        pos = torch.clamp(labels_float.sum(), min=1.0)
+        neg = torch.clamp(labels_float.numel() - labels_float.sum(), min=1.0)
+        return (neg / pos).to(device=labels.device, dtype=labels.dtype)
+
+    def forward(self, inputs: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        labels = labels.float()
+        current_pos_weight = self._resolve_pos_weight(labels)
+        current_pos_weight = current_pos_weight.to(inputs.dtype)
+
+        if self.logits:
+            loss = F.binary_cross_entropy_with_logits(
+                inputs, labels, pos_weight=current_pos_weight, reduction="none"
+            )
+        else:
+            probs = torch.clamp(inputs, min=1e-6, max=1 - 1e-6)
+            base_loss = F.binary_cross_entropy(probs, labels, reduction="none")
+            loss = torch.where(labels == 1, base_loss * current_pos_weight, base_loss)
+
+        if self.reduction == "mean":
+            return loss.mean()
+        elif self.reduction == "sum":
+            return loss.sum()
+        else:
+            return loss
+
+
+class FocalLoss(nn.Module):
+    """
+    Standard focal loss for binary or multi-class classification.
+    """
+
+    def __init__(
+        self,
+        gamma: float = 2.0,
+        alpha: Optional[float | Sequence[float] | torch.Tensor] = None,
+        reduction: str = "mean",
+        logits: bool = False,
+    ):
+        super().__init__()
+        self.gamma = gamma
+        self.reduction = reduction
+        self.logits = logits
+        self.alpha = alpha
+
+    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        # Multi-class case
+        if inputs.dim() > 1 and inputs.size(1) > 1:
+            log_probs = F.log_softmax(inputs, dim=1)
+            probs = log_probs.exp()
+            targets_one_hot = F.one_hot(targets.long(), num_classes=inputs.size(1)).float()
+
+            alpha = self._get_alpha(inputs)
+            alpha_factor = targets_one_hot * alpha
+            focal_weight = (1.0 - probs) ** self.gamma
+            loss = torch.sum(alpha_factor * focal_weight * (-log_probs), dim=1)
+        else:
+            targets = targets.float()
+            if self.logits:
+                ce_loss = F.binary_cross_entropy_with_logits(
+                    inputs, targets, reduction="none"
+                )
+                probs = torch.sigmoid(inputs)
+            else:
+                ce_loss = F.binary_cross_entropy(inputs, targets, reduction="none")
+                probs = torch.clamp(inputs, min=1e-6, max=1 - 1e-6)
+
+            p_t = probs * targets + (1 - probs) * (1 - targets)
+            alpha_factor = self._get_binary_alpha(targets, inputs.device)
+            focal_weight = (1.0 - p_t) ** self.gamma
+            loss = alpha_factor * focal_weight * ce_loss
+
+        if self.reduction == "mean":
+            return loss.mean()
+        if self.reduction == "sum":
+            return loss.sum()
+        return loss
+
+    def _get_alpha(self, inputs: torch.Tensor) -> torch.Tensor:
+        if self.alpha is None:
+            return torch.ones_like(inputs)
+        if isinstance(self.alpha, torch.Tensor):
+            return self.alpha.to(inputs.device)
+        alpha_tensor = torch.tensor(self.alpha, device=inputs.device, dtype=inputs.dtype)
+        return alpha_tensor
+
+    def _get_binary_alpha(self, targets: torch.Tensor, device: torch.device) -> torch.Tensor:
+        if self.alpha is None:
+            return torch.ones_like(targets)
+        if isinstance(self.alpha, (float, int)):
+            return torch.where(targets == 1, self.alpha, 1 - float(self.alpha)).to(device)
+        alpha_tensor = torch.tensor(self.alpha, device=device, dtype=targets.dtype)
+        return torch.where(targets == 1, alpha_tensor, 1 - alpha_tensor)
+
+
+class ClassBalancedFocalLoss(nn.Module):
+    """
+    Focal loss weighted by effective number of samples per class.
+    Reference: "Class-Balanced Loss Based on Effective Number of Samples"
+    """
+
+    def __init__(
+        self,
+        class_counts: Sequence[int] | torch.Tensor,
+        beta: float = 0.9999,
+        gamma: float = 2.0,
+        reduction: str = "mean",
+    ):
+        super().__init__()
+        self.gamma = gamma
+        self.reduction = reduction
+        class_counts = torch.as_tensor(class_counts, dtype=torch.float32)
+        effective_num = 1.0 - torch.pow(beta, class_counts)
+        weights = (1.0 - beta) / (effective_num + 1e-12)
+        weights = weights / weights.sum() * len(weights)
+        self.register_buffer("class_weights", weights)
+
+    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        focal = FocalLoss(
+            gamma=self.gamma, alpha=self.class_weights, reduction="none", logits=True
+        )
+        loss = focal(inputs, targets)
+        if self.reduction == "mean":
+            return loss.mean()
+        if self.reduction == "sum":
+            return loss.sum()
+        return loss

nextrec/models/match/dssm.py

@@ -19,7 +19,8 @@ class DSSM(BaseMatchModel):
     """
     Deep Structured Semantic Model
     
-    双塔模型，分别对user和item特征编码为embedding，通过余弦相似度或点积计算匹配分数
+    Dual-tower model that encodes user and item features separately and
+    computes similarity via cosine or dot product.
     """
     
     @property
@@ -48,6 +49,12 @@ class DSSM(BaseMatchModel):
                  embedding_l2_reg: float = 0.0,
                  dense_l2_reg: float = 0.0,
                  early_stop_patience: int = 20,
+                 optimizer: str | torch.optim.Optimizer = "adam",
+                 optimizer_params: dict | None = None,
+                 scheduler: str | torch.optim.lr_scheduler._LRScheduler | type[torch.optim.lr_scheduler._LRScheduler] | None = None,
+                 scheduler_params: dict | None = None,
+                 loss: str | nn.Module | list[str | nn.Module] | None = "bce",
+                 loss_params: dict | list[dict] | None = None,
                  **kwargs):
         
         super(DSSM, self).__init__(
@@ -86,7 +93,7 @@ class DSSM(BaseMatchModel):
         if len(user_features) > 0:
             self.user_embedding = EmbeddingLayer(user_features)
             
-            # 计算user tower输入维度
+            # Compute user tower input dimension
             user_input_dim = 0
             for feat in user_dense_features or []:
                 user_input_dim += 1
@@ -117,7 +124,7 @@ class DSSM(BaseMatchModel):
         if len(item_features) > 0:
             self.item_embedding = EmbeddingLayer(item_features)
             
-            # 计算item tower输入维度
+            # Compute item tower input dimension
             item_input_dim = 0
             for feat in item_dense_features or []:
                 item_input_dim += 1
@@ -136,7 +143,6 @@ class DSSM(BaseMatchModel):
                 activation=dnn_activation
             )
         
-        # 注册正则化权重
         self._register_regularization_weights(
             embedding_attr='user_embedding',
             include_modules=['user_dnn']
@@ -146,28 +152,33 @@ class DSSM(BaseMatchModel):
             include_modules=['item_dnn']
         )
         
+        if optimizer_params is None:
+            optimizer_params = {"lr": 1e-3, "weight_decay": 1e-5}
+
         self.compile(
-            optimizer="adam",
-            optimizer_params={"lr": 1e-3, "weight_decay": 1e-5},
+            optimizer=optimizer,
+            optimizer_params=optimizer_params,
+            scheduler=scheduler,
+            scheduler_params=scheduler_params,
+            loss=loss,
+            loss_params=loss_params,
         )
 
         self.to(device)
     
     def user_tower(self, user_input: dict) -> torch.Tensor:
         """
-        User tower: 将user特征编码为embedding
+        User tower encodes user features into embeddings.
         
         Args:
-            user_input: user特征字典
+            user_input: user feature dict
         
         Returns:
             user_emb: [batch_size, embedding_dim]
         """
-        # 获取user特征的embedding
         all_user_features = self.user_dense_features + self.user_sparse_features + self.user_sequence_features
         user_emb = self.user_embedding(user_input, all_user_features, squeeze_dim=True)
         
-        # 通过user DNN
         user_emb = self.user_dnn(user_emb)
         
         # L2 normalize for cosine similarity
@@ -178,19 +189,17 @@ class DSSM(BaseMatchModel):
     
     def item_tower(self, item_input: dict) -> torch.Tensor:
         """
-        Item tower: 将item特征编码为embedding
+        Item tower encodes item features into embeddings.
         
         Args:
-            item_input: item特征字典
+            item_input: item feature dict
         
         Returns:
-            item_emb: [batch_size, embedding_dim] 或 [batch_size, num_items, embedding_dim]
+            item_emb: [batch_size, embedding_dim] or [batch_size, num_items, embedding_dim]
         """
-        # 获取item特征的embedding
         all_item_features = self.item_dense_features + self.item_sparse_features + self.item_sequence_features
         item_emb = self.item_embedding(item_input, all_item_features, squeeze_dim=True)
         
-        # 通过item DNN
         item_emb = self.item_dnn(item_emb)
         
         # L2 normalize for cosine similarity