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

nextrec/loss/pointwise.py

@@ -0,0 +1,198 @@
+"""
+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,7 +49,13 @@ class DSSM(BaseMatchModel):
                  embedding_l2_reg: float = 0.0,
                  dense_l2_reg: float = 0.0,
                  early_stop_patience: int = 20,
-                 model_id: str = 'dssm'):
+                 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__(
             user_dense_features=user_dense_features,
@@ -67,7 +74,7 @@ class DSSM(BaseMatchModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=early_stop_patience,
-            model_id=model_id
+            **kwargs
         )
         
         self.embedding_dim = embedding_dim
@@ -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

nextrec/models/match/dssm_v2.py

@@ -44,7 +44,13 @@ class DSSM_v2(BaseMatchModel):
                  embedding_l2_reg: float = 0.0,
                  dense_l2_reg: float = 0.0,
                  early_stop_patience: int = 20,
-                 model_id: str = 'dssm_v2'):
+                 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_v2, self).__init__(
             user_dense_features=user_dense_features,
@@ -63,7 +69,7 @@ class DSSM_v2(BaseMatchModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=early_stop_patience,
-            model_id=model_id
+            **kwargs
         )
         
         self.embedding_dim = embedding_dim
@@ -137,6 +143,18 @@ class DSSM_v2(BaseMatchModel):
             include_modules=['item_dnn']
         )
         
+        if optimizer_params is None:
+            optimizer_params = {"lr": 1e-3, "weight_decay": 1e-5}
+
+        self.compile(
+            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:

nextrec/models/match/mind.py

@@ -41,7 +41,7 @@ class MIND(BaseMatchModel):
                  item_dnn_hidden_units: list[int] = [256, 128],
                  dnn_activation: str = 'relu',
                  dnn_dropout: float = 0.0,
-                 training_mode: Literal['pointwise', 'pairwise', 'listwise'] = 'listwise',
+                 training_mode: Literal['pointwise', 'pairwise', 'listwise'] = 'pointwise',
                  num_negative_samples: int = 100,
                  temperature: float = 1.0,
                  similarity_metric: Literal['dot', 'cosine', 'euclidean'] = 'dot',
@@ -51,7 +51,13 @@ class MIND(BaseMatchModel):
                  embedding_l2_reg: float = 0.0,
                  dense_l2_reg: float = 0.0,
                  early_stop_patience: int = 20,
-                 model_id: str = 'mind'):
+                 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(MIND, self).__init__(
             user_dense_features=user_dense_features,
@@ -70,7 +76,7 @@ class MIND(BaseMatchModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=early_stop_patience,
-            model_id=model_id
+            **kwargs
         )
         
         self.embedding_dim = embedding_dim
@@ -152,6 +158,15 @@ class MIND(BaseMatchModel):
             include_modules=['item_dnn'] if self.item_dnn else []
         )
         
+        self.compile(
+            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:

nextrec/models/match/sdm.py

@@ -52,7 +52,13 @@ class SDM(BaseMatchModel):
                  embedding_l2_reg: float = 0.0,
                  dense_l2_reg: float = 0.0,
                  early_stop_patience: int = 20,
-                 model_id: str = 'sdm'):
+                 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(SDM, self).__init__(
             user_dense_features=user_dense_features,
@@ -71,7 +77,7 @@ class SDM(BaseMatchModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=early_stop_patience,
-            model_id=model_id
+            **kwargs
         )
         
         self.embedding_dim = embedding_dim
@@ -179,6 +185,15 @@ class SDM(BaseMatchModel):
             include_modules=['item_dnn'] if self.item_dnn else []
         )
         
+        self.compile(
+            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:

nextrec/models/match/youtube_dnn.py

@@ -17,11 +17,10 @@ from nextrec.basic.layers import MLP, EmbeddingLayer, AveragePooling
 
 class YoutubeDNN(BaseMatchModel):
     """
-    YouTube Deep Neural Network for Recommendations
-    
-    用户塔：历史行为序列 + 用户特征 -> 用户embedding
-    物品塔：物品特征 -> 物品embedding
-    训练：sampled softmax loss (listwise)
+    YouTube Deep Neural Network for Recommendations.
+    User tower: behavior sequence + user features -> user embedding.
+    Item tower: item features -> item embedding.
+    Training usually uses listwise / sampled softmax style objectives.
     """
     
     @property
@@ -50,7 +49,13 @@ class YoutubeDNN(BaseMatchModel):
                  embedding_l2_reg: float = 0.0,
                  dense_l2_reg: float = 0.0,
                  early_stop_patience: int = 20,
-                 model_id: str = 'youtube_dnn'):
+                 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(YoutubeDNN, self).__init__(
             user_dense_features=user_dense_features,
@@ -69,7 +74,7 @@ class YoutubeDNN(BaseMatchModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=early_stop_patience,
-            model_id=model_id
+            **kwargs
         )
         
         self.embedding_dim = embedding_dim
@@ -94,7 +99,7 @@ class YoutubeDNN(BaseMatchModel):
             for feat in user_sparse_features or []:
                 user_input_dim += feat.embedding_dim
             for feat in user_sequence_features or []:
-                # 序列特征通过平均池化聚合
+                # Sequence features are pooled before entering the DNN
                 user_input_dim += feat.embedding_dim
             
             user_dnn_units = user_dnn_hidden_units + [embedding_dim]
@@ -144,12 +149,20 @@ class YoutubeDNN(BaseMatchModel):
             include_modules=['item_dnn']
         )
         
+        self.compile(
+            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 tower to encode historical behavior sequences and user features.
         """
         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)

nextrec/models/multi_task/esmm.py

@@ -1,9 +1,7 @@
 """
 Date: create on 09/11/2025
-Author:
-    Yang Zhou,zyaztec@gmail.com
-Reference:
-    [1] Ma X, Zhao L, Huang G, et al. Entire space multi-task model: An effective approach for estimating post-click conversion rate[C]//SIGIR. 2018: 1137-1140.
+Author: Yang Zhou,zyaztec@gmail.com
+Reference: [1] Ma X, Zhao L, Huang G, et al. Entire space multi-task model: An effective approach for estimating post-click conversion rate[C]//SIGIR. 2018: 1137-1140.
 """
 
 import torch
@@ -46,12 +44,13 @@ class ESMM(BaseModel):
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | list[str | nn.Module] | None = "bce",
+                 loss_params: dict | list[dict] | None = None,
                  device: str = 'cpu',
-                 model_id: str = "baseline",
                  embedding_l1_reg=1e-6,
                  dense_l1_reg=1e-5,
                  embedding_l2_reg=1e-5,
-                 dense_l2_reg=1e-4):
+                 dense_l2_reg=1e-4,
+                 **kwargs):
         
         # ESMM requires exactly 2 targets: ctr and ctcvr
         if len(target) != 2:
@@ -69,7 +68,7 @@ class ESMM(BaseModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=20,
-            model_id=model_id
+            **kwargs
         )
 
         self.loss = loss
@@ -106,7 +105,8 @@ class ESMM(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/multi_task/mmoe.py

@@ -1,9 +1,7 @@
 """
 Date: create on 09/11/2025
-Author:
-    Yang Zhou,zyaztec@gmail.com
-Reference:
-    [1] Ma J, Zhao Z, Yi X, et al. Modeling task relationships in multi-task learning with multi-gate mixture-of-experts[C]//KDD. 2018: 1930-1939.
+Author: Yang Zhou,zyaztec@gmail.com
+Reference: [1] Ma J, Zhao Z, Yi X, et al. Modeling task relationships in multi-task learning with multi-gate mixture-of-experts[C]//KDD. 2018: 1930-1939.
 """
 
 import torch
@@ -44,12 +42,13 @@ class MMOE(BaseModel):
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | list[str | nn.Module] | None = "bce",
+                 loss_params: dict | list[dict] | None = None,
                  device: str = 'cpu',
-                 model_id: str = "baseline",
                  embedding_l1_reg=1e-6,
                  dense_l1_reg=1e-5,
                  embedding_l2_reg=1e-5,
-                 dense_l2_reg=1e-4):
+                 dense_l2_reg=1e-4,
+                 **kwargs):
         
         super(MMOE, self).__init__(
             dense_features=dense_features,
@@ -63,7 +62,7 @@ class MMOE(BaseModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=20,
-            model_id=model_id
+            **kwargs
         )
 
         self.loss = loss
@@ -128,7 +127,8 @@ class MMOE(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):

nextrec/models/multi_task/ple.py

@@ -1,9 +1,7 @@
 """
 Date: create on 09/11/2025
-Author:
-    Yang Zhou,zyaztec@gmail.com
-Reference:
-    [1] Tang H, Liu J, Zhao M, et al. Progressive layered extraction (ple): A novel multi-task learning (mtl) model for personalized recommendations[C]//RecSys. 2020: 269-278.
+Author: Yang Zhou,zyaztec@gmail.com
+Reference: [1] Tang H, Liu J, Zhao M, et al. Progressive layered extraction (ple): A novel multi-task learning (mtl) model for personalized recommendations[C]//RecSys. 2020: 269-278.
 """
 
 import torch
@@ -47,12 +45,13 @@ class PLE(BaseModel):
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | list[str | nn.Module] | None = "bce",
+                 loss_params: dict | list[dict] | None = None,
                  device: str = 'cpu',
-                 model_id: str = "baseline",
                  embedding_l1_reg=1e-6,
                  dense_l1_reg=1e-5,
                  embedding_l2_reg=1e-5,
-                 dense_l2_reg=1e-4):
+                 dense_l2_reg=1e-4,
+                 **kwargs):
         
         super(PLE, self).__init__(
             dense_features=dense_features,
@@ -66,7 +65,7 @@ class PLE(BaseModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=20,
-            model_id=model_id
+            **kwargs
         )
 
         self.loss = loss
@@ -166,7 +165,8 @@ class PLE(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss
+            loss=loss,
+            loss_params=loss_params,
         )
 
     def forward(self, x):