nextrec 0.3.5__py3-none-any.whl → 0.4.1__py3-none-any.whl

nextrec/models/ranking/__init__.py

@@ -1,27 +0,0 @@
-from .fm import FM
-from .afm import AFM
-from .masknet import MaskNet
-from .pnn import PNN
-from .deepfm import DeepFM
-from .autoint import AutoInt
-from .widedeep import WideDeep
-from .xdeepfm import xDeepFM
-from .dcn import DCN
-from .fibinet import FiBiNET
-from .din import DIN
-from .dien import DIEN
-
-__all__ = [
-    'DeepFM',
-    'AutoInt', 
-    'WideDeep',
-    'xDeepFM',
-    'DCN',
-    'DIN',
-    'DIEN',
-    'FM',
-    'AFM',
-    'MaskNet',
-    'PNN',
-    'FiBiNET',
-]

nextrec/models/ranking/afm.py

@@ -1,17 +1,46 @@
 """
 Date: create on 09/11/2025
-Author:
-    Yang Zhou,zyaztec@gmail.com
+Checkpoint: edit on 06/12/2025
+Author: Yang Zhou,zyaztec@gmail.com
 Reference:
-    [1] Xiao J, Ye H, He X, et al. Attentional factorization machines: Learning the weight of
-        feature interactions via attention networks[C]//IJCAI. 2017: 3119-3125.
+[1] Xiao J, Ye H, He X, et al. Attentional factorization machines: Learning the weight of
+feature interactions via attention networks[C]//IJCAI. 2017: 3119-3125.
+
+Attentional Factorization Machine (AFM) builds on FM by learning an importance
+weight for every second-order interaction instead of treating all pairs equally.
+It retains FM’s linear (first-order) component for sparsity-friendly modeling,
+while using an attention network to reweight the element-wise product of field
+embeddings before aggregation.
+
+In each forward pass:
+  (1) Embed each field and compute pairwise element-wise products v_i ⊙ v_j
+  (2) Pass interactions through an attention MLP (ReLU + projection) to score them
+  (3) Softmax-normalize scores to obtain interaction weights
+  (4) Weighted sum of interactions -> linear projection -> add FM first-order term
+
+Key Advantages:
+- Learns which feature pairs contribute most via attention weights
+- Keeps FM efficiency and interpretability by preserving first-order terms
+- Softmax-normalized reweighting reduces noise from uninformative interactions
+
+AFM 在 FM 的二阶交互上引入注意力，为每个特征对学习重要性权重；同时保留 FM 的一阶项，
+保持对稀疏特征的友好与可解释性。具体流程：
+  (1) 对各字段做 embedding，并计算所有特征对的元素积 v_i ⊙ v_j
+  (2) 经由注意力 MLP（ReLU + 线性映射）得到交互得分
+  (3) 通过 softmax 归一化交互得分，得到权重
+  (4) 将加权交互求和、线性映射，再与一阶项相加得到最终预测
+
+主要优点：
+- 注意力显式告诉哪些特征对更重要
+- 保留 FM 的效率和可解释性
+- softmax 归一化减弱噪声交互的影响
 """
 
 import torch
 import torch.nn as nn
 
 from nextrec.basic.model import BaseModel
-from nextrec.basic.layers import EmbeddingLayer, LR, PredictionLayer
+from nextrec.basic.layers import EmbeddingLayer, LR, PredictionLayer, InputMask
 from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
 
 
@@ -21,7 +50,7 @@ class AFM(BaseModel):
         return "AFM"
 
     @property
-    def task_type(self):
+    def default_task(self):
         return "binary"
     
     def __init__(self,
@@ -31,6 +60,7 @@ class AFM(BaseModel):
                  attention_dim: int = 32,
                  attention_dropout: float = 0.0,
                  target: list[str] | list = [],
+                 task: str | list[str] | None = None,
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | None = "bce",
@@ -46,45 +76,64 @@ class AFM(BaseModel):
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=self.task_type,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs
         )
 
-        self.loss = loss
-        if self.loss is None:
-            self.loss = "bce"
+        if target is None:
+            target = []
+        if optimizer_params is None:
+            optimizer_params = {}
+        if loss is None:
+            loss = "bce"
             
         self.fm_features = sparse_features + sequence_features
         if len(self.fm_features) < 2:
             raise ValueError("AFM requires at least two sparse/sequence features to build pairwise interactions.")
 
-        # Assume uniform embedding dimension across FM fields
+        # make sure all embedding dimension are the same for FM features
         self.embedding_dim = self.fm_features[0].embedding_dim
         if any(f.embedding_dim != self.embedding_dim for f in self.fm_features):
             raise ValueError("All FM features must share the same embedding_dim for AFM.")
 
-        self.embedding = EmbeddingLayer(features=self.fm_features)
-
-        fm_input_dim = sum([f.embedding_dim for f in self.fm_features])
-        self.linear = LR(fm_input_dim)
+        self.embedding = EmbeddingLayer(features=self.fm_features) # [Batch, Field, Dim ]
+
+        # First-order terms: dense linear + one hot embeddings
+        self.dense_features = list(dense_features)
+        dense_input_dim = sum([f.input_dim for f in self.dense_features])
+        self.linear_dense = nn.Linear(dense_input_dim, 1, bias=True) if dense_input_dim > 0 else None
+
+        # First-order term: sparse/sequence features one-hot
+        # **INFO**: source paper does not contain sequence features in experiments,
+        # but we implement it here for completeness. if you want follow the paper strictly,
+        # remove sequence features from fm_features.
+        self.first_order_embeddings = nn.ModuleDict()
+        for feature in self.fm_features:
+            if feature.embedding_name in self.first_order_embeddings: # shared embedding
+                continue
+            emb = nn.Embedding(num_embeddings=feature.vocab_size, embedding_dim=1, padding_idx=feature.padding_idx) # equal to one-hot encoding weight
+            # nn.init.zeros_(emb.weight)
+            self.first_order_embeddings[feature.embedding_name] = emb
 
         self.attention_linear = nn.Linear(self.embedding_dim, attention_dim)
         self.attention_p = nn.Linear(attention_dim, 1, bias=False)
         self.attention_dropout = nn.Dropout(attention_dropout)
         self.output_projection = nn.Linear(self.embedding_dim, 1, bias=False)
-        self.prediction_layer = PredictionLayer(task_type=self.task_type)
+        self.prediction_layer = PredictionLayer(task_type=self.default_task)
+        self.input_mask = InputMask()
 
         # Register regularization weights
         self.register_regularization_weights(
             embedding_attr='embedding',
-            include_modules=['linear', 'attention_linear', 'attention_p', 'output_projection']
+            include_modules=['linear_dense', 'attention_linear', 'attention_p', 'output_projection']
         )
+        # add first-order embeddings to embedding regularization list
+        self.embedding_params.extend(emb.weight for emb in self.first_order_embeddings.values())
 
         self.compile(
             optimizer=optimizer,
@@ -95,10 +144,53 @@ class AFM(BaseModel):
 
     def forward(self, x):
         field_emb = self.embedding(x=x, features=self.fm_features, squeeze_dim=False)  # [B, F, D]
-        input_linear = field_emb.flatten(start_dim=1)
-        y_linear = self.linear(input_linear)
+        batch_size = field_emb.size(0)
+        y_linear = torch.zeros(batch_size, 1, device=field_emb.device)
+
+        # First-order dense part
+        if self.linear_dense is not None:
+            dense_inputs = [x[f.name].float().view(batch_size, -1) for f in self.dense_features]
+            dense_stack = torch.cat(dense_inputs, dim=1) if dense_inputs else None
+            if dense_stack is not None:
+                y_linear = y_linear + self.linear_dense(dense_stack)
+
+        # First-order sparse/sequence part
+        first_order_terms = []
+        for feature in self.fm_features:
+            emb = self.first_order_embeddings[feature.embedding_name]
+            if isinstance(feature, SparseFeature):
+                term = emb(x[feature.name].long())  # [B, 1]
+            else:  # SequenceFeature
+                seq_input = x[feature.name].long() # [B, 1]
+                if feature.max_len is not None and seq_input.size(1) > feature.max_len:
+                    seq_input = seq_input[:, -feature.max_len :]
+                mask = self.input_mask(x, feature, seq_input).squeeze(1)  # [B, 1]
+                seq_weight = emb(seq_input).squeeze(-1)  # [B, L]
+                term = (seq_weight * mask).sum(dim=1, keepdim=True)  # [B, 1]
+            first_order_terms.append(term)
+        if first_order_terms:
+            y_linear = y_linear + torch.sum(torch.cat(first_order_terms, dim=1), dim=1, keepdim=True)
 
         interactions = []
+        feature_values = []
+        for feature in self.fm_features:
+            value = x.get(f"{feature.name}_value")
+            if value is not None:
+                value = value.float()
+                if value.dim() == 1:
+                    value = value.unsqueeze(-1)
+            else:
+                if isinstance(feature, SequenceFeature):
+                    seq_input = x[feature.name].long()
+                    if feature.max_len is not None and seq_input.size(1) > feature.max_len:
+                        seq_input = seq_input[:, -feature.max_len :]
+                    value = self.input_mask(x, feature, seq_input).sum(dim=2)  # [B, 1]
+                else:
+                    value = torch.ones(batch_size, 1, device=field_emb.device)
+            feature_values.append(value)
+        feature_values_tensor = torch.cat(feature_values, dim=1).unsqueeze(-1)  # [B, F, 1]
+        field_emb = field_emb * feature_values_tensor
+
         num_fields = field_emb.shape[1]
         for i in range(num_fields - 1):
             vi = field_emb[:, i, :]
@@ -107,7 +199,7 @@ class AFM(BaseModel):
                 interactions.append(vi * vj)
 
         pair_tensor = torch.stack(interactions, dim=1)  # [B, num_pairs, D]
-        attention_scores = torch.tanh(self.attention_linear(pair_tensor))
+        attention_scores = torch.relu(self.attention_linear(pair_tensor))
         attention_scores = self.attention_p(attention_scores)  # [B, num_pairs, 1]
         attention_weights = torch.softmax(attention_scores, dim=1)
 

nextrec/models/ranking/autoint.py

@@ -68,7 +68,7 @@ class AutoInt(BaseModel):
         return "AutoInt"
 
     @property
-    def task_type(self):
+    def default_task(self):
         return "binary"
     
     def __init__(self,
@@ -80,9 +80,10 @@ class AutoInt(BaseModel):
                  att_head_num: int = 2,
                  att_dropout: float = 0.0,
                  att_use_residual: bool = True,
-                 target: list[str] = [],
+                 target: list[str] | None = None,
+                 task: str | list[str] | None = None,
                  optimizer: str = "adam",
-                 optimizer_params: dict = {},
+                 optimizer_params: dict | None = None,
                  loss: str | nn.Module | None = "bce",
                  loss_params: dict | list[dict] | None = None,
                  device: str = 'cpu',
@@ -97,24 +98,29 @@ class AutoInt(BaseModel):
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=self.task_type,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs
         )
 
-        self.loss = loss
-        if self.loss is None:
-            self.loss = "bce"
+        if target is None:
+            target = []
+        if optimizer_params is None:
+            optimizer_params = {}
+        if loss is None:
+            loss = "bce"
             
         self.att_layer_num = att_layer_num
         self.att_embedding_dim = att_embedding_dim
         
         # Use sparse and sequence features for interaction
+        # **INFO**: this is different from the original paper, we also include dense features
+        # if you want to follow the paper strictly, set dense_features=[]
+        # or modify the code accordingly
         self.interaction_features = dense_features + sparse_features + sequence_features 
         
         # All features for embedding
@@ -147,7 +153,7 @@ class AutoInt(BaseModel):
         
         # Final prediction layer
         self.fc = nn.Linear(num_fields * att_embedding_dim, 1)
-        self.prediction_layer = PredictionLayer(task_type=self.task_type)
+        self.prediction_layer = PredictionLayer(task_type=self.default_task)
 
         # Register regularization weights
         self.register_regularization_weights(

nextrec/models/ranking/dcn.py

@@ -25,15 +25,11 @@ class CrossNetwork(nn.Module):
         self.b = torch.nn.ParameterList([torch.nn.Parameter(torch.zeros((input_dim,))) for _ in range(num_layers)])
 
     def forward(self, x):
-        """
-        :param x: Float tensor of size ``(batch_size, num_fields, embed_dim)``
-        """
         x0 = x
         for i in range(self.num_layers):
             xw = self.w[i](x)
             x = x0 * xw + self.b[i] + x
-        return x
-
+        return x # [batch_size, input_dim]
 
 class DCN(BaseModel):
     @property
@@ -41,9 +37,9 @@ class DCN(BaseModel):
         return "DCN"
 
     @property
-    def task_type(self):
+    def default_task(self):
         return "binary"
-    
+
     def __init__(self,
                  dense_features: list[DenseFeature],
                  sparse_features: list[SparseFeature],
@@ -51,6 +47,7 @@ class DCN(BaseModel):
                  cross_num: int = 3,
                  mlp_params: dict | None = None,
                  target: list[str] = [],
+                 task: str | list[str] | None = None,
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | None = "bce",
@@ -67,13 +64,12 @@ class DCN(BaseModel):
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=self.task_type,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs
         )
 
@@ -99,14 +95,15 @@ class DCN(BaseModel):
         if mlp_params is not None:
             self.use_dnn = True
             self.mlp = MLP(input_dim=input_dim, **mlp_params)
+            deep_dim = self.mlp.output_dim
             # Final layer combines cross and deep
-            self.final_layer = nn.Linear(input_dim + 1, 1)  # +1 for MLP output
+            self.final_layer = nn.Linear(input_dim + deep_dim, 1)  # + deep_dim for MLP output
         else:
             self.use_dnn = False
             # Final layer only uses cross network output
             self.final_layer = nn.Linear(input_dim, 1)
 
-        self.prediction_layer = PredictionLayer(task_type=self.task_type)
+        self.prediction_layer = PredictionLayer(task_type=self.task)
 
         # Register regularization weights
         self.register_regularization_weights(

nextrec/models/ranking/deepfm.py

@@ -56,15 +56,16 @@ class DeepFM(BaseModel):
         return "DeepFM"
 
     @property
-    def task_type(self):
+    def default_task(self):
         return "binary"
-    
+
     def __init__(self,
                  dense_features: list[DenseFeature]|list = [],
                  sparse_features: list[SparseFeature]|list = [],
                  sequence_features: list[SequenceFeature]|list = [],
                  mlp_params: dict = {},
                  target: list[str]|str = [],
+                 task: str | list[str] | None = None,
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | None = "bce",
@@ -80,13 +81,12 @@ class DeepFM(BaseModel):
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=self.task_type,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs
         )
 
@@ -104,7 +104,7 @@ class DeepFM(BaseModel):
         self.linear = LR(fm_emb_dim_total)
         self.fm = FM(reduce_sum=True)
         self.mlp = MLP(input_dim=mlp_input_dim, **mlp_params)
-        self.prediction_layer = PredictionLayer(task_type=self.task_type)
+        self.prediction_layer = PredictionLayer(task_type=self.default_task)
 
         # Register regularization weights
         self.register_regularization_weights(embedding_attr='embedding', include_modules=['linear', 'mlp'])

nextrec/models/ranking/dien.py

@@ -146,7 +146,7 @@ class DIEN(BaseModel):
         return "DIEN"
 
     @property
-    def task_type(self):
+    def default_task(self):
         return "binary"
     
     def __init__(self,
@@ -159,6 +159,7 @@ class DIEN(BaseModel):
                  attention_activation: str = 'sigmoid',
                  use_negsampling: bool = False,
                  target: list[str] = [],
+                 task: str | list[str] | None = None,
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | None = "bce",
@@ -175,13 +176,12 @@ class DIEN(BaseModel):
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=self.task_type,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs
         )
 
@@ -235,7 +235,7 @@ class DIEN(BaseModel):
         mlp_input_dim += sum([getattr(f, "embedding_dim", 1) or 1 for f in dense_features])
         # MLP for final prediction
         self.mlp = MLP(input_dim=mlp_input_dim, **mlp_params)
-        self.prediction_layer = PredictionLayer(task_type=self.task_type)
+        self.prediction_layer = PredictionLayer(task_type=self.task)
         # Register regularization weights
         self.register_regularization_weights(embedding_attr='embedding', include_modules=['interest_extractor', 'interest_evolution', 'attention_layer', 'mlp', 'candidate_proj'])
         self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params)

nextrec/models/ranking/din.py

@@ -22,7 +22,7 @@ class DIN(BaseModel):
         return "DIN"
 
     @property
-    def task_type(self):
+    def default_task(self):
         return "binary"
     
     def __init__(self,
@@ -34,6 +34,7 @@ class DIN(BaseModel):
                  attention_activation: str = 'sigmoid',
                  attention_use_softmax: bool = True,
                  target: list[str] = [],
+                 task: str | list[str] | None = None,
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | None = "bce",
@@ -50,13 +51,12 @@ class DIN(BaseModel):
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=self.task_type,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs
         )
 
@@ -105,7 +105,7 @@ class DIN(BaseModel):
         
         # MLP for final prediction
         self.mlp = MLP(input_dim=mlp_input_dim, **mlp_params)
-        self.prediction_layer = PredictionLayer(task_type=self.task_type)
+        self.prediction_layer = PredictionLayer(task_type=self.task)
 
         # Register regularization weights
         self.register_regularization_weights(

nextrec/models/ranking/fibinet.py

@@ -28,7 +28,7 @@ class FiBiNET(BaseModel):
         return "FiBiNET"
 
     @property
-    def task_type(self):
+    def default_task(self):
         return "binary"
     
     def __init__(self,
@@ -39,6 +39,7 @@ class FiBiNET(BaseModel):
                  bilinear_type: str = "field_interaction",
                  senet_reduction: int = 3,
                  target: list[str] | list = [],
+                 task: str | list[str] | None = None,
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | None = "bce",
@@ -55,13 +56,12 @@ class FiBiNET(BaseModel):
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=self.task_type,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs
         )
 
@@ -101,7 +101,7 @@ class FiBiNET(BaseModel):
         num_pairs = self.num_fields * (self.num_fields - 1) // 2
         interaction_dim = num_pairs * self.embedding_dim * 2
         self.mlp = MLP(input_dim=interaction_dim, **mlp_params)
-        self.prediction_layer = PredictionLayer(task_type=self.task_type)
+        self.prediction_layer = PredictionLayer(task_type=self.default_task)
 
         # Register regularization weights
         self.register_regularization_weights(

nextrec/models/ranking/fm.py

@@ -19,7 +19,7 @@ class FM(BaseModel):
         return "FM"
 
     @property
-    def task_type(self):
+    def default_task(self):
         return "binary"
     
     def __init__(self,
@@ -27,6 +27,7 @@ class FM(BaseModel):
                  sparse_features: list[SparseFeature] | list = [],
                  sequence_features: list[SequenceFeature] | list = [],
                  target: list[str] | list = [],
+                 task: str | list[str] | None = None,
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | None = "bce",
@@ -43,13 +44,12 @@ class FM(BaseModel):
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=self.task_type,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs
         )
 
@@ -66,7 +66,7 @@ class FM(BaseModel):
         fm_input_dim = sum([f.embedding_dim for f in self.fm_features])
         self.linear = LR(fm_input_dim)
         self.fm = FMInteraction(reduce_sum=True)
-        self.prediction_layer = PredictionLayer(task_type=self.task_type)
+        self.prediction_layer = PredictionLayer(task_type=self.task)
 
         # Register regularization weights
         self.register_regularization_weights(

nextrec/models/ranking/masknet.py

@@ -143,8 +143,7 @@ class MaskNet(BaseModel):
         return "MaskNet"
 
     @property
-    def task_type(self):
-        # Align with PredictionLayer supported task types
+    def default_task(self):
         return "binary"
 
     def __init__(
@@ -159,6 +158,7 @@ class MaskNet(BaseModel):
         block_dropout: float = 0.0,
         mlp_params: dict | None = None,
         target: list[str] | None = None,
+        task: str | list[str] | None = None,
         optimizer: str = "adam",
         optimizer_params: dict | None = None,
         loss: str | nn.Module | None = "bce",
@@ -182,13 +182,12 @@ class MaskNet(BaseModel):
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=self.task_type,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs,
         )
 
@@ -231,7 +230,7 @@ class MaskNet(BaseModel):
             self.mask_blocks = nn.ModuleList([MaskBlockOnEmbedding(num_fields=self.num_fields, embedding_dim=self.embedding_dim, mask_hidden_dim=mask_hidden_dim, hidden_dim=block_hidden_dim) for _ in range(self.num_blocks)])
             self.final_mlp = MLP(input_dim=self.num_blocks * block_hidden_dim, **mlp_params)
             self.output_layer = None
-        self.prediction_layer = PredictionLayer(task_type=self.task_type)
+        self.prediction_layer = PredictionLayer(task_type=self.task)
 
         if self.model_type == "serial":
             self.register_regularization_weights(embedding_attr="embedding", include_modules=["mask_blocks", "output_layer"],)

nextrec/models/ranking/pnn.py

@@ -20,7 +20,7 @@ class PNN(BaseModel):
         return "PNN"
 
     @property
-    def task_type(self):
+    def default_task(self):
         return "binary"
     
     def __init__(self,
@@ -31,6 +31,7 @@ class PNN(BaseModel):
                  product_type: str = "inner",
                  outer_product_dim: int | None = None,
                  target: list[str] | list = [],
+                 task: str | list[str] | None = None,
                  optimizer: str = "adam",
                  optimizer_params: dict = {},
                  loss: str | nn.Module | None = "bce",
@@ -47,13 +48,12 @@ class PNN(BaseModel):
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=self.task_type,
+            task=task or self.default_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            early_stop_patience=20,
             **kwargs
         )
 
@@ -86,7 +86,7 @@ class PNN(BaseModel):
 
         linear_dim = self.num_fields * self.embedding_dim
         self.mlp = MLP(input_dim=linear_dim + product_dim, **mlp_params)
-        self.prediction_layer = PredictionLayer(task_type=self.task_type)
+        self.prediction_layer = PredictionLayer(task_type=self.task)
 
         modules = ['mlp']
         if self.product_type == "outer":