nextrec 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

nextrec/models/ranking/fibinet.py

@@ -30,25 +30,27 @@ class FiBiNET(BaseModel):
     @property
     def task_type(self):
         return "binary"
-    
-    def __init__(self,
-                 dense_features: list[DenseFeature] | list = [],
-                 sparse_features: list[SparseFeature] | list = [],
-                 sequence_features: list[SequenceFeature] | list = [],
-                 mlp_params: dict = {},
-                 bilinear_type: str = "field_interaction",
-                 senet_reduction: int = 3,
-                 target: list[str] | list = [],
-                 optimizer: str = "adam",
-                 optimizer_params: dict = {},
-                 loss: str | nn.Module | None = "bce",
-                 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):
-        
+
+    def __init__(
+        self,
+        dense_features: list[DenseFeature] | list = [],
+        sparse_features: list[SparseFeature] | list = [],
+        sequence_features: list[SequenceFeature] | list = [],
+        mlp_params: dict = {},
+        bilinear_type: str = "field_interaction",
+        senet_reduction: int = 3,
+        target: list[str] | list = [],
+        optimizer: str = "adam",
+        optimizer_params: dict = {},
+        loss: str | nn.Module | None = "bce",
+        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,
+    ):
+
         super(FiBiNET, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -61,28 +63,36 @@ class FiBiNET(BaseModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=20,
-            model_id=model_id
+            model_id=model_id,
         )
 
         self.loss = loss
         if self.loss is None:
             self.loss = "bce"
-            
+
         self.linear_features = sparse_features + sequence_features
         self.deep_features = dense_features + sparse_features + sequence_features
         self.interaction_features = sparse_features + sequence_features
 
         if len(self.interaction_features) < 2:
-            raise ValueError("FiBiNET requires at least two sparse/sequence features for interactions.")
+            raise ValueError(
+                "FiBiNET requires at least two sparse/sequence features for interactions."
+            )
 
         self.embedding = EmbeddingLayer(features=self.deep_features)
 
         self.num_fields = len(self.interaction_features)
         self.embedding_dim = self.interaction_features[0].embedding_dim
-        if any(f.embedding_dim != self.embedding_dim for f in self.interaction_features):
-            raise ValueError("All interaction features must share the same embedding_dim in FiBiNET.")
-
-        self.senet = SENETLayer(num_fields=self.num_fields, reduction_ratio=senet_reduction)
+        if any(
+            f.embedding_dim != self.embedding_dim for f in self.interaction_features
+        ):
+            raise ValueError(
+                "All interaction features must share the same embedding_dim in FiBiNET."
+            )
+
+        self.senet = SENETLayer(
+            num_fields=self.num_fields, reduction_ratio=senet_reduction
+        )
         self.bilinear_standard = BiLinearInteractionLayer(
             input_dim=self.embedding_dim,
             num_fields=self.num_fields,
@@ -104,21 +114,27 @@ class FiBiNET(BaseModel):
 
         # Register regularization weights
         self._register_regularization_weights(
-            embedding_attr='embedding',
-            include_modules=['linear', 'senet', 'bilinear_standard', 'bilinear_senet', 'mlp']
+            embedding_attr="embedding",
+            include_modules=[
+                "linear",
+                "senet",
+                "bilinear_standard",
+                "bilinear_senet",
+                "mlp",
+            ],
         )
 
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=loss
-        )
+        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss)
 
     def forward(self, x):
-        input_linear = self.embedding(x=x, features=self.linear_features, squeeze_dim=True)
+        input_linear = self.embedding(
+            x=x, features=self.linear_features, squeeze_dim=True
+        )
         y_linear = self.linear(input_linear)
 
-        field_emb = self.embedding(x=x, features=self.interaction_features, squeeze_dim=False)
+        field_emb = self.embedding(
+            x=x, features=self.interaction_features, squeeze_dim=False
+        )
         senet_emb = self.senet(field_emb)
 
         bilinear_standard = self.bilinear_standard(field_emb).flatten(start_dim=1)

nextrec/models/ranking/fm.py

@@ -9,7 +9,12 @@ Reference:
 import torch.nn as nn
 
 from nextrec.basic.model import BaseModel
-from nextrec.basic.layers import EmbeddingLayer, FM as FMInteraction, LR, PredictionLayer
+from nextrec.basic.layers import (
+    EmbeddingLayer,
+    FM as FMInteraction,
+    LR,
+    PredictionLayer,
+)
 from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
 
 
@@ -21,22 +26,24 @@ class FM(BaseModel):
     @property
     def task_type(self):
         return "binary"
-    
-    def __init__(self,
-                 dense_features: list[DenseFeature] | list = [],
-                 sparse_features: list[SparseFeature] | list = [],
-                 sequence_features: list[SequenceFeature] | list = [],
-                 target: list[str] | list = [],
-                 optimizer: str = "adam",
-                 optimizer_params: dict = {},
-                 loss: str | nn.Module | None = "bce",
-                 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):
-        
+
+    def __init__(
+        self,
+        dense_features: list[DenseFeature] | list = [],
+        sparse_features: list[SparseFeature] | list = [],
+        sequence_features: list[SequenceFeature] | list = [],
+        target: list[str] | list = [],
+        optimizer: str = "adam",
+        optimizer_params: dict = {},
+        loss: str | nn.Module | None = "bce",
+        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,
+    ):
+
         super(FM, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -49,13 +56,13 @@ class FM(BaseModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=20,
-            model_id=model_id
+            model_id=model_id,
         )
 
         self.loss = loss
         if self.loss is None:
             self.loss = "bce"
-            
+
         self.fm_features = sparse_features + sequence_features
         if len(self.fm_features) == 0:
             raise ValueError("FM requires at least one sparse or sequence feature.")
@@ -69,15 +76,10 @@ class FM(BaseModel):
 
         # Register regularization weights
         self._register_regularization_weights(
-            embedding_attr='embedding',
-            include_modules=['linear']
+            embedding_attr="embedding", include_modules=["linear"]
         )
 
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=loss
-        )
+        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss)
 
     def forward(self, x):
         input_fm = self.embedding(x=x, features=self.fm_features, squeeze_dim=False)

nextrec/models/ranking/masknet.py

@@ -23,26 +23,28 @@ class MaskNet(BaseModel):
     @property
     def task_type(self):
         return "binary"
-    
-    def __init__(self,
-                 dense_features: list[DenseFeature] | list = [],
-                 sparse_features: list[SparseFeature] | list = [],
-                 sequence_features: list[SequenceFeature] | list = [],
-                 num_blocks: int = 3,
-                 mask_hidden_dim: int = 64,
-                 block_dropout: float = 0.1,
-                 mlp_params: dict = {},
-                 target: list[str] | list = [],
-                 optimizer: str = "adam",
-                 optimizer_params: dict = {},
-                 loss: str | nn.Module | None = "bce",
-                 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):
-        
+
+    def __init__(
+        self,
+        dense_features: list[DenseFeature] | list = [],
+        sparse_features: list[SparseFeature] | list = [],
+        sequence_features: list[SequenceFeature] | list = [],
+        num_blocks: int = 3,
+        mask_hidden_dim: int = 64,
+        block_dropout: float = 0.1,
+        mlp_params: dict = {},
+        target: list[str] | list = [],
+        optimizer: str = "adam",
+        optimizer_params: dict = {},
+        loss: str | nn.Module | None = "bce",
+        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,
+    ):
+
         super(MaskNet, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -55,13 +57,13 @@ class MaskNet(BaseModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=20,
-            model_id=model_id
+            model_id=model_id,
         )
 
         self.loss = loss
         if self.loss is None:
             self.loss = "bce"
-            
+
         self.mask_features = sparse_features + sequence_features
         if len(self.mask_features) == 0:
             raise ValueError("MaskNet requires at least one sparse/sequence feature.")
@@ -70,7 +72,9 @@ class MaskNet(BaseModel):
         self.num_fields = len(self.mask_features)
         self.embedding_dim = self.mask_features[0].embedding_dim
         if any(f.embedding_dim != self.embedding_dim for f in self.mask_features):
-            raise ValueError("MaskNet expects identical embedding_dim across mask_features.")
+            raise ValueError(
+                "MaskNet expects identical embedding_dim across mask_features."
+            )
 
         self.num_blocks = max(1, num_blocks)
         self.field_dim = self.num_fields * self.embedding_dim
@@ -82,7 +86,7 @@ class MaskNet(BaseModel):
                 nn.Sequential(
                     nn.Linear(self.field_dim, mask_hidden_dim),
                     nn.ReLU(),
-                    nn.Linear(mask_hidden_dim, self.num_fields)
+                    nn.Linear(mask_hidden_dim, self.num_fields),
                 )
             )
 
@@ -91,15 +95,11 @@ class MaskNet(BaseModel):
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
         self._register_regularization_weights(
-            embedding_attr='embedding',
-            include_modules=['linear', 'mask_generators', 'final_mlp']
+            embedding_attr="embedding",
+            include_modules=["linear", "mask_generators", "final_mlp"],
         )
 
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=loss
-        )
+        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss)
 
     def forward(self, x):
         field_emb = self.embedding(x=x, features=self.mask_features, squeeze_dim=False)

nextrec/models/ranking/pnn.py

@@ -22,25 +22,27 @@ class PNN(BaseModel):
     @property
     def task_type(self):
         return "binary"
-    
-    def __init__(self,
-                 dense_features: list[DenseFeature] | list = [],
-                 sparse_features: list[SparseFeature] | list = [],
-                 sequence_features: list[SequenceFeature] | list = [],
-                 mlp_params: dict = {},
-                 product_type: str = "inner",
-                 outer_product_dim: int | None = None,
-                 target: list[str] | list = [],
-                 optimizer: str = "adam",
-                 optimizer_params: dict = {},
-                 loss: str | nn.Module | None = "bce",
-                 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):
-        
+
+    def __init__(
+        self,
+        dense_features: list[DenseFeature] | list = [],
+        sparse_features: list[SparseFeature] | list = [],
+        sequence_features: list[SequenceFeature] | list = [],
+        mlp_params: dict = {},
+        product_type: str = "inner",
+        outer_product_dim: int | None = None,
+        target: list[str] | list = [],
+        optimizer: str = "adam",
+        optimizer_params: dict = {},
+        loss: str | nn.Module | None = "bce",
+        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,
+    ):
+
         super(PNN, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -53,13 +55,13 @@ class PNN(BaseModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=20,
-            model_id=model_id
+            model_id=model_id,
         )
 
         self.loss = loss
         if self.loss is None:
             self.loss = "bce"
-            
+
         self.field_features = sparse_features + sequence_features
         if len(self.field_features) < 2:
             raise ValueError("PNN requires at least two sparse/sequence features.")
@@ -68,7 +70,9 @@ class PNN(BaseModel):
         self.num_fields = len(self.field_features)
         self.embedding_dim = self.field_features[0].embedding_dim
         if any(f.embedding_dim != self.embedding_dim for f in self.field_features):
-            raise ValueError("All field features must share the same embedding_dim for PNN.")
+            raise ValueError(
+                "All field features must share the same embedding_dim for PNN."
+            )
 
         self.product_type = product_type.lower()
         if self.product_type not in {"inner", "outer"}:
@@ -87,19 +91,14 @@ class PNN(BaseModel):
         self.mlp = MLP(input_dim=linear_dim + product_dim, **mlp_params)
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
-        modules = ['mlp']
+        modules = ["mlp"]
         if self.product_type == "outer":
-            modules.append('kernel')
+            modules.append("kernel")
         self._register_regularization_weights(
-            embedding_attr='embedding',
-            include_modules=modules
+            embedding_attr="embedding", include_modules=modules
         )
 
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=loss
-        )
+        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss)
 
     def forward(self, x):
         field_emb = self.embedding(x=x, features=self.field_features, squeeze_dim=False)

nextrec/models/ranking/widedeep.py

@@ -24,23 +24,25 @@ class WideDeep(BaseModel):
     @property
     def task_type(self):
         return "binary"
-    
-    def __init__(self,
-                 dense_features: list[DenseFeature],
-                 sparse_features: list[SparseFeature],
-                 sequence_features: list[SequenceFeature],
-                 mlp_params: dict,
-                 target: list[str] = [],
-                 optimizer: str = "adam",
-                 optimizer_params: dict = {},
-                 loss: str | nn.Module | None = "bce",
-                 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):
-        
+
+    def __init__(
+        self,
+        dense_features: list[DenseFeature],
+        sparse_features: list[SparseFeature],
+        sequence_features: list[SequenceFeature],
+        mlp_params: dict,
+        target: list[str] = [],
+        optimizer: str = "adam",
+        optimizer_params: dict = {},
+        loss: str | nn.Module | None = "bce",
+        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,
+    ):
+
         super(WideDeep, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -53,16 +55,16 @@ class WideDeep(BaseModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=20,
-            model_id=model_id
+            model_id=model_id,
         )
 
         self.loss = loss
         if self.loss is None:
             self.loss = "bce"
-            
+
         # Wide part: use all features for linear model
         self.wide_features = sparse_features + sequence_features
-        
+
         # Deep part: use all features
         self.deep_features = dense_features + sparse_features + sequence_features
 
@@ -72,30 +74,33 @@ class WideDeep(BaseModel):
         # Wide part: Linear layer
         wide_dim = sum([f.embedding_dim for f in self.wide_features])
         self.linear = LR(wide_dim)
-        
+
         # Deep part: MLP
-        deep_emb_dim_total = sum([f.embedding_dim for f in self.deep_features if not isinstance(f, DenseFeature)])
-        dense_input_dim = sum([getattr(f, "embedding_dim", 1) or 1 for f in dense_features])
+        deep_emb_dim_total = sum(
+            [
+                f.embedding_dim
+                for f in self.deep_features
+                if not isinstance(f, DenseFeature)
+            ]
+        )
+        dense_input_dim = sum(
+            [getattr(f, "embedding_dim", 1) or 1 for f in dense_features]
+        )
         self.mlp = MLP(input_dim=deep_emb_dim_total + dense_input_dim, **mlp_params)
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
         # Register regularization weights
         self._register_regularization_weights(
-            embedding_attr='embedding',
-            include_modules=['linear', 'mlp']
+            embedding_attr="embedding", include_modules=["linear", "mlp"]
         )
 
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=loss
-        )
+        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss)
 
     def forward(self, x):
         # Deep part
         input_deep = self.embedding(x=x, features=self.deep_features, squeeze_dim=True)
         y_deep = self.mlp(input_deep)  # [B, 1]
-        
+
         # Wide part
         input_wide = self.embedding(x=x, features=self.wide_features, squeeze_dim=True)
         y_wide = self.linear(input_wide)

nextrec/models/ranking/xdeepfm.py

@@ -3,8 +3,8 @@ Date: create on 09/11/2025
 Author:
     Yang Zhou,zyaztec@gmail.com
 Reference:
-    [1] Lian J, Zhou X, Zhang F, et al. xdeepfm: Combining explicit and implicit feature interactions 
-        for recommender systems[C]//Proceedings of the 24th ACM SIGKDD international conference on 
+    [1] Lian J, Zhou X, Zhang F, et al. xdeepfm: Combining explicit and implicit feature interactions
+        for recommender systems[C]//Proceedings of the 24th ACM SIGKDD international conference on
         knowledge discovery & data mining. 2018: 1754-1763.
         (https://arxiv.org/abs/1803.05170)
 """
@@ -25,25 +25,27 @@ class xDeepFM(BaseModel):
     @property
     def task_type(self):
         return "binary"
-    
-    def __init__(self,
-                 dense_features: list[DenseFeature],
-                 sparse_features: list[SparseFeature],
-                 sequence_features: list[SequenceFeature],
-                 mlp_params: dict,
-                 cin_size: list[int] = [128, 128],
-                 split_half: bool = True,
-                 target: list[str] = [],
-                 optimizer: str = "adam",
-                 optimizer_params: dict = {},
-                 loss: str | nn.Module | None = "bce",
-                 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):
-        
+
+    def __init__(
+        self,
+        dense_features: list[DenseFeature],
+        sparse_features: list[SparseFeature],
+        sequence_features: list[SequenceFeature],
+        mlp_params: dict,
+        cin_size: list[int] = [128, 128],
+        split_half: bool = True,
+        target: list[str] = [],
+        optimizer: str = "adam",
+        optimizer_params: dict = {},
+        loss: str | nn.Module | None = "bce",
+        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,
+    ):
+
         super(xDeepFM, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -56,16 +58,16 @@ class xDeepFM(BaseModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=20,
-            model_id=model_id
+            model_id=model_id,
         )
 
         self.loss = loss
         if self.loss is None:
             self.loss = "bce"
-            
+
         # Linear part and CIN part: use sparse and sequence features
         self.linear_features = sparse_features + sequence_features
-        
+
         # Deep part: use all features
         self.deep_features = dense_features + sparse_features + sequence_features
 
@@ -75,39 +77,44 @@ class xDeepFM(BaseModel):
         # Linear part
         linear_dim = sum([f.embedding_dim for f in self.linear_features])
         self.linear = LR(linear_dim)
-        
+
         # CIN part: Compressed Interaction Network
         num_fields = len(self.linear_features)
         self.cin = CIN(input_dim=num_fields, cin_size=cin_size, split_half=split_half)
-        
+
         # Deep part: DNN
-        deep_emb_dim_total = sum([f.embedding_dim for f in self.deep_features if not isinstance(f, DenseFeature)])
-        dense_input_dim = sum([getattr(f, "embedding_dim", 1) or 1 for f in dense_features])
+        deep_emb_dim_total = sum(
+            [
+                f.embedding_dim
+                for f in self.deep_features
+                if not isinstance(f, DenseFeature)
+            ]
+        )
+        dense_input_dim = sum(
+            [getattr(f, "embedding_dim", 1) or 1 for f in dense_features]
+        )
         self.mlp = MLP(input_dim=deep_emb_dim_total + dense_input_dim, **mlp_params)
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
         # Register regularization weights
         self._register_regularization_weights(
-            embedding_attr='embedding',
-            include_modules=['linear', 'cin', 'mlp']
+            embedding_attr="embedding", include_modules=["linear", "cin", "mlp"]
         )
 
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=loss
-        )
+        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss)
 
     def forward(self, x):
         # Get embeddings for linear and CIN (sparse features only)
-        input_linear = self.embedding(x=x, features=self.linear_features, squeeze_dim=False)
-        
+        input_linear = self.embedding(
+            x=x, features=self.linear_features, squeeze_dim=False
+        )
+
         # Linear part
         y_linear = self.linear(input_linear.flatten(start_dim=1))
-        
+
         # CIN part
         y_cin = self.cin(input_linear)  # [B, 1]
-        
+
         # Deep part
         input_deep = self.embedding(x=x, features=self.deep_features, squeeze_dim=True)
         y_deep = self.mlp(input_deep)  # [B, 1]