nextrec 0.1.4__py3-none-any.whl → 0.1.7__py3-none-any.whl

nextrec/models/ranking/autoint.py

@@ -3,8 +3,8 @@ Date: create on 09/11/2025
 Author:
     Yang Zhou,zyaztec@gmail.com
 Reference:
-    [1] Song W, Shi C, Xiao Z, et al. Autoint: Automatic feature interaction learning via
-        self-attentive neural networks[C]//Proceedings of the 28th ACM international conference
+    [1] Song W, Shi C, Xiao Z, et al. Autoint: Automatic feature interaction learning via 
+        self-attentive neural networks[C]//Proceedings of the 28th ACM international conference 
         on information and knowledge management. 2019: 1161-1170.
         (https://arxiv.org/abs/1810.11921)
 """
@@ -25,29 +25,27 @@ class AutoInt(BaseModel):
     @property
     def task_type(self):
         return "binary"
-
-    def __init__(
-        self,
-        dense_features: list[DenseFeature],
-        sparse_features: list[SparseFeature],
-        sequence_features: list[SequenceFeature],
-        att_layer_num: int = 3,
-        att_embedding_dim: int = 8,
-        att_head_num: int = 2,
-        att_dropout: float = 0.0,
-        att_use_residual: 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],
+                 att_layer_num: int = 3,
+                 att_embedding_dim: int = 8,
+                 att_head_num: int = 2,
+                 att_dropout: float = 0.0,
+                 att_use_residual: 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(AutoInt, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -60,94 +58,83 @@ class AutoInt(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.att_layer_num = att_layer_num
         self.att_embedding_dim = att_embedding_dim
-
+        
         # Use sparse and sequence features for interaction
         self.interaction_features = sparse_features + sequence_features
-
+        
         # All features for embedding
         self.all_features = dense_features + sparse_features + sequence_features
 
         # Embedding layer
         self.embedding = EmbeddingLayer(features=self.all_features)
-
+        
         # Project embeddings to attention embedding dimension
         num_fields = len(self.interaction_features)
         total_embedding_dim = sum([f.embedding_dim for f in self.interaction_features])
-
+        
         # If embeddings have different dimensions, project them to att_embedding_dim
-        self.need_projection = not all(
-            f.embedding_dim == att_embedding_dim for f in self.interaction_features
-        )
+        self.need_projection = not all(f.embedding_dim == att_embedding_dim for f in self.interaction_features)
         self.projection_layers = None
         if self.need_projection:
-            self.projection_layers = nn.ModuleList(
-                [
-                    nn.Linear(f.embedding_dim, att_embedding_dim, bias=False)
-                    for f in self.interaction_features
-                ]
-            )
-
+            self.projection_layers = nn.ModuleList([
+                nn.Linear(f.embedding_dim, att_embedding_dim, bias=False) 
+                for f in self.interaction_features
+            ])
+        
         # Multi-head self-attention layers
-        self.attention_layers = nn.ModuleList(
-            [
-                MultiHeadSelfAttention(
-                    embedding_dim=att_embedding_dim,
-                    num_heads=att_head_num,
-                    dropout=att_dropout,
-                    use_residual=att_use_residual,
-                )
-                for _ in range(att_layer_num)
-            ]
-        )
-
+        self.attention_layers = nn.ModuleList([
+            MultiHeadSelfAttention(
+                embedding_dim=att_embedding_dim,
+                num_heads=att_head_num,
+                dropout=att_dropout,
+                use_residual=att_use_residual
+            ) for _ in range(att_layer_num)
+        ])
+        
         # Final prediction layer
         self.fc = nn.Linear(num_fields * att_embedding_dim, 1)
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
         # Register regularization weights
         self._register_regularization_weights(
-            embedding_attr="embedding",
-            include_modules=["projection_layers", "attention_layers", "fc"],
+            embedding_attr='embedding',
+            include_modules=['projection_layers', 'attention_layers', 'fc']
         )
 
-        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 field-by-field so mixed dimensions can be projected safely
         field_embeddings = []
         if len(self.interaction_features) == 0:
-            raise ValueError(
-                "AutoInt requires at least one sparse or sequence feature for interactions."
-            )
+            raise ValueError("AutoInt requires at least one sparse or sequence feature for interactions.")
         for idx, feature in enumerate(self.interaction_features):
             feature_emb = self.embedding(x=x, features=[feature], squeeze_dim=False)
             feature_emb = feature_emb.squeeze(1)  # [B, embedding_dim]
             if self.need_projection and self.projection_layers is not None:
                 feature_emb = self.projection_layers[idx](feature_emb)
-            field_embeddings.append(
-                feature_emb.unsqueeze(1)
-            )  # [B, 1, att_embedding_dim or original_dim]
+            field_embeddings.append(feature_emb.unsqueeze(1))  # [B, 1, att_embedding_dim or original_dim]
         embeddings = torch.cat(field_embeddings, dim=1)
-
+        
         # Apply multi-head self-attention layers
         attention_output = embeddings
         for att_layer in self.attention_layers:
-            attention_output = att_layer(
-                attention_output
-            )  # [B, num_fields, att_embedding_dim]
-
+            attention_output = att_layer(attention_output)  # [B, num_fields, att_embedding_dim]
+        
         # Flatten and predict
-        attention_output_flat = attention_output.flatten(
-            start_dim=1
-        )  # [B, num_fields * att_embedding_dim]
+        attention_output_flat = attention_output.flatten(start_dim=1)  # [B, num_fields * att_embedding_dim]
         y = self.fc(attention_output_flat)  # [B, 1]
         return self.prediction_layer(y)

nextrec/models/ranking/dcn.py

@@ -24,26 +24,24 @@ class DCN(BaseModel):
     @property
     def task_type(self):
         return "binary"
-
-    def __init__(
-        self,
-        dense_features: list[DenseFeature],
-        sparse_features: list[SparseFeature],
-        sequence_features: list[SequenceFeature],
-        cross_num: int = 3,
-        mlp_params: dict | None = None,
-        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],
+                 cross_num: int = 3,
+                 mlp_params: dict | None = None,
+                 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(DCN, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -56,13 +54,13 @@ class DCN(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"
-
+            
         # All features
         self.all_features = dense_features + sparse_features + sequence_features
 
@@ -70,21 +68,13 @@ class DCN(BaseModel):
         self.embedding = EmbeddingLayer(features=self.all_features)
 
         # Calculate input dimension
-        emb_dim_total = sum(
-            [
-                f.embedding_dim
-                for f in self.all_features
-                if not isinstance(f, DenseFeature)
-            ]
-        )
-        dense_input_dim = sum(
-            [getattr(f, "embedding_dim", 1) or 1 for f in dense_features]
-        )
+        emb_dim_total = sum([f.embedding_dim for f in self.all_features if not isinstance(f, DenseFeature)])
+        dense_input_dim = sum([getattr(f, "embedding_dim", 1) or 1 for f in dense_features])
         input_dim = emb_dim_total + dense_input_dim
-
+        
         # Cross Network
         self.cross_network = CrossNetwork(input_dim=input_dim, num_layers=cross_num)
-
+        
         # Deep Network (optional)
         if mlp_params is not None:
             self.use_dnn = True
@@ -100,29 +90,31 @@ class DCN(BaseModel):
 
         # Register regularization weights
         self._register_regularization_weights(
-            embedding_attr="embedding",
-            include_modules=["cross_network", "mlp", "final_layer"],
+            embedding_attr='embedding',
+            include_modules=['cross_network', 'mlp', 'final_layer']
         )
 
-        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 all embeddings and flatten
         input_flat = self.embedding(x=x, features=self.all_features, squeeze_dim=True)
-
+        
         # Cross Network
         cross_output = self.cross_network(input_flat)  # [B, input_dim]
-
+        
         if self.use_dnn:
             # Deep Network
             deep_output = self.mlp(input_flat)  # [B, 1]
             # Concatenate cross and deep
-            combined = torch.cat(
-                [cross_output, deep_output], dim=-1
-            )  # [B, input_dim + 1]
+            combined = torch.cat([cross_output, deep_output], dim=-1)  # [B, input_dim + 1]
         else:
             combined = cross_output
-
+        
         # Final prediction
         y = self.final_layer(combined)
         return self.prediction_layer(y)

nextrec/models/ranking/deepfm.py

@@ -13,7 +13,6 @@ from nextrec.basic.model import BaseModel
 from nextrec.basic.layers import FM, LR, EmbeddingLayer, MLP, PredictionLayer
 from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
 
-
 class DeepFM(BaseModel):
     @property
     def model_name(self):
@@ -22,25 +21,23 @@ class DeepFM(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 = {},
-        target: list[str] | 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]|list = [],
+                 sparse_features: list[SparseFeature]|list = [],
+                 sequence_features: list[SequenceFeature]|list = [],
+                 mlp_params: dict = {},
+                 target: list[str]|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(DeepFM, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -53,29 +50,21 @@ class DeepFM(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
         self.deep_features = dense_features + sparse_features + sequence_features
 
         self.embedding = EmbeddingLayer(features=self.deep_features)
 
         fm_emb_dim_total = sum([f.embedding_dim for f in self.fm_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]
-        )
+        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.linear = LR(fm_emb_dim_total)
         self.fm = FM(reduce_sum=True)
@@ -84,10 +73,15 @@ class DeepFM(BaseModel):
 
         # 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):
         input_deep = self.embedding(x=x, features=self.deep_features, squeeze_dim=True)