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

nextrec/models/ranking/din.py

@@ -12,7 +12,12 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.model import BaseModel
-from nextrec.basic.layers import EmbeddingLayer, MLP, AttentionPoolingLayer, PredictionLayer
+from nextrec.basic.layers import (
+    EmbeddingLayer,
+    MLP,
+    AttentionPoolingLayer,
+    PredictionLayer,
+)
 from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
 
 
@@ -24,28 +29,30 @@ class DIN(BaseModel):
     @property
     def default_task(self):
         return "binary"
-    
-    def __init__(self,
-                 dense_features: list[DenseFeature],
-                 sparse_features: list[SparseFeature],
-                 sequence_features: list[SequenceFeature],
-                 mlp_params: dict,
-                 attention_hidden_units: list[int] = [80, 40],
-                 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",
-                 loss_params: dict | list[dict] | None = None,
-                 device: str = 'cpu',
-                 embedding_l1_reg=1e-6,
-                 dense_l1_reg=1e-5,
-                 embedding_l2_reg=1e-5,
-                 dense_l2_reg=1e-4,
-                 **kwargs):
-        
+
+    def __init__(
+        self,
+        dense_features: list[DenseFeature],
+        sparse_features: list[SparseFeature],
+        sequence_features: list[SequenceFeature],
+        mlp_params: dict,
+        attention_hidden_units: list[int] = [80, 40],
+        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",
+        loss_params: dict | list[dict] | None = None,
+        device: str = "cpu",
+        embedding_l1_reg=1e-6,
+        dense_l1_reg=1e-5,
+        embedding_l2_reg=1e-5,
+        dense_l2_reg=1e-4,
+        **kwargs,
+    ):
+
         super(DIN, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -57,43 +64,54 @@ class DIN(BaseModel):
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            **kwargs
+            **kwargs,
         )
 
         self.loss = loss
         if self.loss is None:
             self.loss = "bce"
-        
+
         # Features classification
         # DIN requires: candidate item + user behavior sequence + other features
         if len(sequence_features) == 0:
-            raise ValueError("DIN requires at least one sequence feature for user behavior history")
-        
+            raise ValueError(
+                "DIN requires at least one sequence feature for user behavior history"
+            )
+
         self.behavior_feature = sequence_features[0]  # User behavior sequence
-        self.candidate_feature = sparse_features[-1] if sparse_features else None  # Candidate item
-        
+        self.candidate_feature = (
+            sparse_features[-1] if sparse_features else None
+        )  # Candidate item
+
         # Other features (excluding behavior sequence in final concatenation)
-        self.other_sparse_features = sparse_features[:-1] if self.candidate_feature else sparse_features
+        self.other_sparse_features = (
+            sparse_features[:-1] if self.candidate_feature else sparse_features
+        )
         self.dense_features_list = dense_features
-        
+
         # All features for embedding
         self.all_features = dense_features + sparse_features + sequence_features
 
         # Embedding layer
         self.embedding = EmbeddingLayer(features=self.all_features)
-        
+
         # Attention layer for behavior sequence
         behavior_emb_dim = self.behavior_feature.embedding_dim
         self.candidate_attention_proj = None
-        if self.candidate_feature is not None and self.candidate_feature.embedding_dim != behavior_emb_dim:
-            self.candidate_attention_proj = nn.Linear(self.candidate_feature.embedding_dim, behavior_emb_dim)
+        if (
+            self.candidate_feature is not None
+            and self.candidate_feature.embedding_dim != behavior_emb_dim
+        ):
+            self.candidate_attention_proj = nn.Linear(
+                self.candidate_feature.embedding_dim, behavior_emb_dim
+            )
         self.attention = AttentionPoolingLayer(
             embedding_dim=behavior_emb_dim,
             hidden_units=attention_hidden_units,
             activation=attention_activation,
-            use_softmax=attention_use_softmax
+            use_softmax=attention_use_softmax,
         )
-        
+
         # Calculate MLP input dimension
         # candidate + attention_pooled_behavior + other_sparse + dense
         mlp_input_dim = 0
@@ -101,16 +119,18 @@ class DIN(BaseModel):
             mlp_input_dim += self.candidate_feature.embedding_dim
         mlp_input_dim += behavior_emb_dim  # attention pooled
         mlp_input_dim += sum([f.embedding_dim for f in self.other_sparse_features])
-        mlp_input_dim += sum([getattr(f, "embedding_dim", 1) or 1 for f in dense_features])
-        
+        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)
 
         # Register regularization weights
         self.register_regularization_weights(
-            embedding_attr='embedding',
-            include_modules=['attention', 'mlp', 'candidate_attention_proj']
+            embedding_attr="embedding",
+            include_modules=["attention", "mlp", "candidate_attention_proj"],
         )
 
         self.compile(
@@ -123,61 +143,69 @@ class DIN(BaseModel):
     def forward(self, x):
         # Get candidate item embedding
         if self.candidate_feature:
-            candidate_emb = self.embedding.embed_dict[self.candidate_feature.embedding_name](
+            candidate_emb = self.embedding.embed_dict[
+                self.candidate_feature.embedding_name
+            ](
                 x[self.candidate_feature.name].long()
             )  # [B, emb_dim]
         else:
             candidate_emb = None
-        
+
         # Get behavior sequence embedding
         behavior_seq = x[self.behavior_feature.name].long()  # [B, seq_len]
         behavior_emb = self.embedding.embed_dict[self.behavior_feature.embedding_name](
             behavior_seq
         )  # [B, seq_len, emb_dim]
-        
+
         # Create mask for padding
         if self.behavior_feature.padding_idx is not None:
-            mask = (behavior_seq != self.behavior_feature.padding_idx).unsqueeze(-1).float()
+            mask = (
+                (behavior_seq != self.behavior_feature.padding_idx)
+                .unsqueeze(-1)
+                .float()
+            )
         else:
             mask = (behavior_seq != 0).unsqueeze(-1).float()
-        
+
         # Apply attention pooling
         if candidate_emb is not None:
             candidate_query = candidate_emb
             if self.candidate_attention_proj is not None:
                 candidate_query = self.candidate_attention_proj(candidate_query)
             pooled_behavior = self.attention(
-                query=candidate_query,
-                keys=behavior_emb,
-                mask=mask
+                query=candidate_query, keys=behavior_emb, mask=mask
             )  # [B, emb_dim]
         else:
             # If no candidate, use mean pooling
-            pooled_behavior = torch.sum(behavior_emb * mask, dim=1) / (mask.sum(dim=1) + 1e-9)
-        
+            pooled_behavior = torch.sum(behavior_emb * mask, dim=1) / (
+                mask.sum(dim=1) + 1e-9
+            )
+
         # Get other features
         other_embeddings = []
-        
+
         if candidate_emb is not None:
             other_embeddings.append(candidate_emb)
-        
+
         other_embeddings.append(pooled_behavior)
-        
+
         # Other sparse features
         for feat in self.other_sparse_features:
-            feat_emb = self.embedding.embed_dict[feat.embedding_name](x[feat.name].long())
+            feat_emb = self.embedding.embed_dict[feat.embedding_name](
+                x[feat.name].long()
+            )
             other_embeddings.append(feat_emb)
-        
+
         # Dense features
         for feat in self.dense_features_list:
             val = x[feat.name].float()
             if val.dim() == 1:
                 val = val.unsqueeze(1)
             other_embeddings.append(val)
-        
+
         # Concatenate all features
         concat_input = torch.cat(other_embeddings, dim=-1)  # [B, total_dim]
-        
+
         # MLP prediction
         y = self.mlp(concat_input)  # [B, 1]
         return self.prediction_layer(y)

nextrec/models/ranking/fibinet.py

@@ -30,27 +30,29 @@ class FiBiNET(BaseModel):
     @property
     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 = {},
-                 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",
-                 loss_params: dict | list[dict] | None = None,
-                 device: str = 'cpu',
-                 embedding_l1_reg=1e-6,
-                 dense_l1_reg=1e-5,
-                 embedding_l2_reg=1e-5,
-                 dense_l2_reg=1e-4,
-                 **kwargs):
-        
+
+    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 = [],
+        task: str | list[str] | None = None,
+        optimizer: str = "adam",
+        optimizer_params: dict = {},
+        loss: str | nn.Module | None = "bce",
+        loss_params: dict | list[dict] | None = None,
+        device: str = "cpu",
+        embedding_l1_reg=1e-6,
+        dense_l1_reg=1e-5,
+        embedding_l2_reg=1e-5,
+        dense_l2_reg=1e-4,
+        **kwargs,
+    ):
+
         super(FiBiNET, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -62,28 +64,36 @@ class FiBiNET(BaseModel):
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            **kwargs
+            **kwargs,
         )
 
         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,
@@ -105,8 +115,14 @@ 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(
@@ -117,10 +133,14 @@ class FiBiNET(BaseModel):
         )
 
     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,24 +26,26 @@ class FM(BaseModel):
     @property
     def default_task(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 = [],
-                 task: str | list[str] | None = None,
-                 optimizer: str = "adam",
-                 optimizer_params: dict = {},
-                 loss: str | nn.Module | None = "bce",
-                 loss_params: dict | list[dict] | None = None,
-                 device: str = 'cpu',
-                 embedding_l1_reg=1e-6,
-                 dense_l1_reg=1e-5,
-                 embedding_l2_reg=1e-5,
-                 dense_l2_reg=1e-4,
-                 **kwargs):
-        
+
+    def __init__(
+        self,
+        dense_features: list[DenseFeature] | list = [],
+        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",
+        loss_params: dict | list[dict] | None = None,
+        device: str = "cpu",
+        embedding_l1_reg=1e-6,
+        dense_l1_reg=1e-5,
+        embedding_l2_reg=1e-5,
+        dense_l2_reg=1e-4,
+        **kwargs,
+    ):
+
         super(FM, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -50,13 +57,13 @@ class FM(BaseModel):
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
-            **kwargs
+            **kwargs,
         )
 
         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.")
@@ -70,8 +77,7 @@ class FM(BaseModel):
 
         # Register regularization weights
         self.register_regularization_weights(
-            embedding_attr='embedding',
-            include_modules=['linear']
+            embedding_attr="embedding", include_modules=["linear"]
         )
 
         self.compile(

nextrec/models/ranking/masknet.py

@@ -69,12 +69,13 @@ class InstanceGuidedMask(nn.Module):
         self.fc2 = nn.Linear(hidden_dim, output_dim)
 
     def forward(self, v_emb_flat: torch.Tensor) -> torch.Tensor:
-        # v_emb_flat: [batch, features count * embedding_dim] 
+        # v_emb_flat: [batch, features count * embedding_dim]
         x = self.fc1(v_emb_flat)
         x = F.relu(x)
         v_mask = self.fc2(x)
         return v_mask
 
+
 class MaskBlockOnEmbedding(nn.Module):
     def __init__(
         self,
@@ -86,20 +87,28 @@ class MaskBlockOnEmbedding(nn.Module):
         super().__init__()
         self.num_fields = num_fields
         self.embedding_dim = embedding_dim
-        self.input_dim = num_fields * embedding_dim  # input_dim = features count * embedding_dim
+        self.input_dim = (
+            num_fields * embedding_dim
+        )  # input_dim = features count * embedding_dim
         self.ln_emb = nn.LayerNorm(embedding_dim)
-        self.mask_gen = InstanceGuidedMask(input_dim=self.input_dim, hidden_dim=mask_hidden_dim, output_dim=self.input_dim,)
+        self.mask_gen = InstanceGuidedMask(
+            input_dim=self.input_dim,
+            hidden_dim=mask_hidden_dim,
+            output_dim=self.input_dim,
+        )
         self.ffn = nn.Linear(self.input_dim, hidden_dim)
         self.ln_hid = nn.LayerNorm(hidden_dim)
 
     # different from MaskBlockOnHidden: input is field embeddings
-    def forward(self, field_emb: torch.Tensor, v_emb_flat: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, field_emb: torch.Tensor, v_emb_flat: torch.Tensor
+    ) -> torch.Tensor:
         B = field_emb.size(0)
-        norm_emb = self.ln_emb(field_emb)        # [B, features count, embedding_dim]
-        norm_emb_flat = norm_emb.view(B, -1)     # [B, features count * embedding_dim]
-        v_mask = self.mask_gen(v_emb_flat)       # [B, features count * embedding_dim]
-        v_masked_emb = v_mask * norm_emb_flat    # [B, features count * embedding_dim]
-        hidden = self.ffn(v_masked_emb)          # [B, hidden_dim]
+        norm_emb = self.ln_emb(field_emb)  # [B, features count, embedding_dim]
+        norm_emb_flat = norm_emb.view(B, -1)  # [B, features count * embedding_dim]
+        v_mask = self.mask_gen(v_emb_flat)  # [B, features count * embedding_dim]
+        v_masked_emb = v_mask * norm_emb_flat  # [B, features count * embedding_dim]
+        hidden = self.ffn(v_masked_emb)  # [B, hidden_dim]
         hidden = self.ln_hid(hidden)
         hidden = F.relu(hidden)
 
@@ -123,15 +132,21 @@ class MaskBlockOnHidden(nn.Module):
         self.ln_input = nn.LayerNorm(hidden_dim)
         self.ln_output = nn.LayerNorm(hidden_dim)
 
-        self.mask_gen = InstanceGuidedMask(input_dim=self.v_emb_dim, hidden_dim=mask_hidden_dim, output_dim=hidden_dim,)
+        self.mask_gen = InstanceGuidedMask(
+            input_dim=self.v_emb_dim,
+            hidden_dim=mask_hidden_dim,
+            output_dim=hidden_dim,
+        )
         self.ffn = nn.Linear(hidden_dim, hidden_dim)
 
     # different from MaskBlockOnEmbedding: input is hidden representation
-    def forward(self, hidden_in: torch.Tensor, v_emb_flat: torch.Tensor) -> torch.Tensor:
-        norm_hidden = self.ln_input(hidden_in)  
+    def forward(
+        self, hidden_in: torch.Tensor, v_emb_flat: torch.Tensor
+    ) -> torch.Tensor:
+        norm_hidden = self.ln_input(hidden_in)
         v_mask = self.mask_gen(v_emb_flat)
-        v_masked_hid = v_mask * norm_hidden     
-        out = self.ffn(v_masked_hid)  
+        v_masked_hid = v_mask * norm_hidden
+        out = self.ffn(v_masked_hid)
         out = self.ln_output(out)
         out = F.relu(out)
         return out
@@ -151,7 +166,7 @@ class MaskNet(BaseModel):
         dense_features: list[DenseFeature] | None = None,
         sparse_features: list[SparseFeature] | None = None,
         sequence_features: list[SequenceFeature] | None = None,
-        model_type: str = "parallel",         # "serial" or "parallel"
+        model_type: str = "parallel",  # "serial" or "parallel"
         num_blocks: int = 3,
         mask_hidden_dim: int = 64,
         block_hidden_dim: int = 256,
@@ -199,50 +214,97 @@ class MaskNet(BaseModel):
         self.sparse_features = sparse_features
         self.sequence_features = sequence_features
         self.mask_features = self.all_features  # use all features for masking
-        assert len(self.mask_features) > 0, "MaskNet requires at least one feature for masking."
+        assert (
+            len(self.mask_features) > 0
+        ), "MaskNet requires at least one feature for masking."
         self.embedding = EmbeddingLayer(features=self.mask_features)
         self.num_fields = len(self.mask_features)
         self.embedding_dim = getattr(self.mask_features[0], "embedding_dim", None)
-        assert self.embedding_dim is not None, "MaskNet requires mask_features to have 'embedding_dim' defined."
+        assert (
+            self.embedding_dim is not None
+        ), "MaskNet requires mask_features to have 'embedding_dim' defined."
 
         for f in self.mask_features:
             edim = getattr(f, "embedding_dim", None)
             if edim is None or edim != self.embedding_dim:
-                raise ValueError(f"MaskNet expects identical embedding_dim across all mask_features, but got {edim} for feature {getattr(f, 'name', type(f))}.")
+                raise ValueError(
+                    f"MaskNet expects identical embedding_dim across all mask_features, but got {edim} for feature {getattr(f, 'name', type(f))}."
+                )
 
         self.v_emb_dim = self.num_fields * self.embedding_dim
         self.model_type = model_type.lower()
-        assert self.model_type in ("serial", "parallel"), "model_type must be either 'serial' or 'parallel'."
+        assert self.model_type in (
+            "serial",
+            "parallel",
+        ), "model_type must be either 'serial' or 'parallel'."
 
         self.num_blocks = max(1, num_blocks)
         self.block_hidden_dim = block_hidden_dim
-        self.block_dropout = nn.Dropout(block_dropout) if block_dropout > 0 else nn.Identity()
+        self.block_dropout = (
+            nn.Dropout(block_dropout) if block_dropout > 0 else nn.Identity()
+        )
 
         if self.model_type == "serial":
-            self.first_block = MaskBlockOnEmbedding(num_fields=self.num_fields, embedding_dim=self.embedding_dim, mask_hidden_dim=mask_hidden_dim, hidden_dim=block_hidden_dim,)
+            self.first_block = MaskBlockOnEmbedding(
+                num_fields=self.num_fields,
+                embedding_dim=self.embedding_dim,
+                mask_hidden_dim=mask_hidden_dim,
+                hidden_dim=block_hidden_dim,
+            )
             self.hidden_blocks = nn.ModuleList(
-                [MaskBlockOnHidden(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 - 1)])
+                [
+                    MaskBlockOnHidden(
+                        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 - 1)
+                ]
+            )
             self.mask_blocks = nn.ModuleList([self.first_block, *self.hidden_blocks])
             self.output_layer = nn.Linear(block_hidden_dim, 1)
             self.final_mlp = None
 
         else:  # parallel
-            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.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)
 
         if self.model_type == "serial":
-            self.register_regularization_weights(embedding_attr="embedding", include_modules=["mask_blocks", "output_layer"],)
+            self.register_regularization_weights(
+                embedding_attr="embedding",
+                include_modules=["mask_blocks", "output_layer"],
+            )
         # serial
         else:
-            self.register_regularization_weights(embedding_attr="embedding", include_modules=["mask_blocks", "final_mlp"])
-        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss, loss_params=loss_params)
+            self.register_regularization_weights(
+                embedding_attr="embedding", include_modules=["mask_blocks", "final_mlp"]
+            )
+        self.compile(
+            optimizer=optimizer,
+            optimizer_params=optimizer_params,
+            loss=loss,
+            loss_params=loss_params,
+        )
 
     def forward(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
         field_emb = self.embedding(x=x, features=self.mask_features, squeeze_dim=False)
         B = field_emb.size(0)
-        v_emb_flat = field_emb.view(B, -1)  # flattened embeddings 
+        v_emb_flat = field_emb.view(B, -1)  # flattened embeddings
 
         if self.model_type == "parallel":
             block_outputs = []
@@ -253,7 +315,7 @@ class MaskNet(BaseModel):
             concat_hidden = torch.cat(block_outputs, dim=-1)
             logit = self.final_mlp(concat_hidden)  # [B, 1]
         # serial
-        else: 
+        else:
             hidden = self.first_block(field_emb, v_emb_flat)
             hidden = self.block_dropout(hidden)
             for block in self.hidden_blocks: