nextrec 0.4.16__py3-none-any.whl → 0.4.18__py3-none-any.whl

nextrec/loss/grad_norm.py

@@ -2,7 +2,7 @@
 GradNorm loss weighting for multi-task learning.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 20/12/2025
+Checkpoint: edit on 24/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 
 Reference:
@@ -45,7 +45,7 @@ class GradNormLossWeighting:
     Adaptive multi-task loss weighting with GradNorm.
 
     Args:
-        num_tasks: Number of tasks.
+        nums_task: Number of tasks.
         alpha: GradNorm balancing strength.
         lr: Learning rate for the weight optimizer.
         init_weights: Optional initial weights per task.
@@ -58,7 +58,7 @@ class GradNormLossWeighting:
 
     def __init__(
         self,
-        num_tasks: int,
+        nums_task: int,
         alpha: float = 1.5,
         lr: float = 0.025,
         init_weights: Iterable[float] | None = None,
@@ -68,9 +68,9 @@ class GradNormLossWeighting:
         init_ema_decay: float = 0.9,
         eps: float = 1e-8,
     ) -> None:
-        if num_tasks <= 1:
-            raise ValueError("GradNorm requires num_tasks > 1.")
-        self.num_tasks = num_tasks
+        if nums_task <= 1:
+            raise ValueError("GradNorm requires nums_task > 1.")
+        self.nums_task = nums_task
         self.alpha = alpha
         self.eps = eps
         if ema_decay is not None:
@@ -87,12 +87,12 @@ class GradNormLossWeighting:
         self.init_ema_count = 0
 
         if init_weights is None:
-            weights = torch.ones(self.num_tasks, dtype=torch.float32)
+            weights = torch.ones(self.nums_task, dtype=torch.float32)
         else:
             weights = torch.tensor(list(init_weights), dtype=torch.float32)
-            if weights.numel() != self.num_tasks:
+            if weights.numel() != self.nums_task:
                 raise ValueError(
-                    "init_weights length must match num_tasks for GradNorm."
+                    "init_weights length must match nums_task for GradNorm."
                 )
         if device is not None:
             weights = weights.to(device)
@@ -123,9 +123,9 @@ class GradNormLossWeighting:
         """
         Return weighted total loss and update task weights with GradNorm.
         """
-        if len(task_losses) != self.num_tasks:
+        if len(task_losses) != self.nums_task:
             raise ValueError(
-                f"Expected {self.num_tasks} task losses, got {len(task_losses)}."
+                f"Expected {self.nums_task} task losses, got {len(task_losses)}."
             )
         shared_params = [p for p in shared_params if p.requires_grad]
         if not shared_params:
@@ -152,7 +152,7 @@ class GradNormLossWeighting:
 
         weights_detached = self.weights.detach()
         weighted_losses = [
-            weights_detached[i] * task_losses[i] for i in range(self.num_tasks)
+            weights_detached[i] * task_losses[i] for i in range(self.nums_task)
         ]
         total_loss = torch.stack(weighted_losses).sum()
 
@@ -226,7 +226,7 @@ class GradNormLossWeighting:
 
         with torch.no_grad():
             w = self.weights.clamp(min=self.eps)
-            w = w * self.num_tasks / (w.sum() + self.eps)
+            w = w * self.nums_task / (w.sum() + self.eps)
             self.weights.copy_(w)
 
         self.pending_grad = None

nextrec/models/multi_task/esmm.py

@@ -1,6 +1,6 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 29/11/2025
+Checkpoint: edit on 23/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 Reference:
 [1] Ma X, Zhao L, Huang G, et al. Entire space multi-task model: An effective approach
@@ -45,7 +45,8 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -100,17 +101,17 @@ class ESMM(BaseModel):
                 f"ESMM requires exactly 2 targets (ctr and ctcvr), got {len(target)}"
             )
 
-        self.num_tasks = len(target)
+        self.nums_task = len(target)
         resolved_task = task
         if resolved_task is None:
             resolved_task = self.default_task
         elif isinstance(resolved_task, str):
-            resolved_task = [resolved_task] * self.num_tasks
-        elif len(resolved_task) == 1 and self.num_tasks > 1:
-            resolved_task = resolved_task * self.num_tasks
-        elif len(resolved_task) != self.num_tasks:
+            resolved_task = [resolved_task] * self.nums_task
+        elif len(resolved_task) == 1 and self.nums_task > 1:
+            resolved_task = resolved_task * self.nums_task
+        elif len(resolved_task) != self.nums_task:
             raise ValueError(
-                f"Length of task ({len(resolved_task)}) must match number of targets ({self.num_tasks})."
+                f"Length of task ({len(resolved_task)}) must match number of targets ({self.nums_task})."
             )
         # resolved_task is now guaranteed to be a list[str]
 
@@ -139,9 +140,7 @@ class ESMM(BaseModel):
         # CVR tower
         self.cvr_tower = MLP(input_dim=input_dim, output_layer=True, **cvr_params)
         self.grad_norm_shared_modules = ["embedding"]
-        self.prediction_layer = PredictionLayer(
-            task_type=self.default_task, task_dims=[1, 1]
-        )
+        self.prediction_layer = TaskHead(task_type=self.default_task, task_dims=[1, 1])
         # Register regularization weights
         self.register_regularization_weights(
             embedding_attr="embedding", include_modules=["ctr_tower", "cvr_tower"]

nextrec/models/multi_task/mmoe.py

@@ -1,6 +1,6 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 29/11/2025
+Checkpoint: edit on 23/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 Reference:
 [1] Ma J, Zhao Z, Yi X, et al. Modeling task relationships in multi-task learning with
@@ -46,7 +46,8 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -66,9 +67,9 @@ class MMOE(BaseModel):
 
     @property
     def default_task(self):
-        num_tasks = getattr(self, "num_tasks", None)
-        if num_tasks is not None and num_tasks > 0:
-            return ["binary"] * num_tasks
+        nums_task = getattr(self, "nums_task", None)
+        if nums_task is not None and nums_task > 0:
+            return ["binary"] * nums_task
         return ["binary"]
 
     def __init__(
@@ -106,18 +107,18 @@ class MMOE(BaseModel):
         elif isinstance(target, str):
             target = [target]
 
-        self.num_tasks = len(target) if target else 1
+        self.nums_task = len(target) if target else 1
 
         resolved_task = task
         if resolved_task is None:
             resolved_task = self.default_task
         elif isinstance(resolved_task, str):
-            resolved_task = [resolved_task] * self.num_tasks
-        elif len(resolved_task) == 1 and self.num_tasks > 1:
-            resolved_task = resolved_task * self.num_tasks
-        elif len(resolved_task) != self.num_tasks:
+            resolved_task = [resolved_task] * self.nums_task
+        elif len(resolved_task) == 1 and self.nums_task > 1:
+            resolved_task = resolved_task * self.nums_task
+        elif len(resolved_task) != self.nums_task:
             raise ValueError(
-                f"Length of task ({len(resolved_task)}) must match number of targets ({self.num_tasks})."
+                f"Length of task ({len(resolved_task)}) must match number of targets ({self.nums_task})."
             )
 
         super(MMOE, self).__init__(
@@ -137,12 +138,12 @@ class MMOE(BaseModel):
         self.loss = loss
 
         # Number of tasks and experts
-        self.num_tasks = len(target)
+        self.nums_task = len(target)
         self.num_experts = num_experts
 
-        if len(tower_params_list) != self.num_tasks:
+        if len(tower_params_list) != self.nums_task:
             raise ValueError(
-                f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.num_tasks})"
+                f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.nums_task})"
             )
 
         self.embedding = EmbeddingLayer(features=self.all_features)
@@ -162,7 +163,7 @@ class MMOE(BaseModel):
 
         # Task-specific gates
         self.gates = nn.ModuleList()
-        for _ in range(self.num_tasks):
+        for _ in range(self.nums_task):
             gate = nn.Sequential(nn.Linear(input_dim, num_experts), nn.Softmax(dim=1))
             self.gates.append(gate)
         self.grad_norm_shared_modules = ["embedding", "experts", "gates"]
@@ -172,8 +173,8 @@ class MMOE(BaseModel):
         for tower_params in tower_params_list:
             tower = MLP(input_dim=expert_output_dim, output_layer=True, **tower_params)
             self.towers.append(tower)
-        self.prediction_layer = PredictionLayer(
-            task_type=self.default_task, task_dims=[1] * self.num_tasks
+        self.prediction_layer = TaskHead(
+            task_type=self.default_task, task_dims=[1] * self.nums_task
         )
         # Register regularization weights
         self.register_regularization_weights(
@@ -198,7 +199,7 @@ class MMOE(BaseModel):
 
         # Task-specific processing
         task_outputs = []
-        for task_idx in range(self.num_tasks):
+        for task_idx in range(self.nums_task):
             # Gate weights for this task: [B, num_experts]
             gate_weights = self.gates[task_idx](input_flat)  # [B, num_experts]
 
@@ -217,6 +218,6 @@ class MMOE(BaseModel):
             tower_output = self.towers[task_idx](gated_output)  # [B, 1]
             task_outputs.append(tower_output)
 
-        # Stack outputs: [B, num_tasks]
+        # Stack outputs: [B, nums_task]
         y = torch.cat(task_outputs, dim=1)
         return self.prediction_layer(y)

nextrec/models/multi_task/ple.py

@@ -1,6 +1,6 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 29/11/2025
+Checkpoint: edit on 23/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 Reference:
 [1] Tang H, Liu J, Zhao M, et al. Progressive layered extraction (PLE): A novel
@@ -49,7 +49,8 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 from nextrec.utils.model import get_mlp_output_dim
 
@@ -63,18 +64,18 @@ class CGCLayer(nn.Module):
     def __init__(
         self,
         input_dim: int,
-        num_tasks: int,
+        nums_task: int,
         num_shared_experts: int,
         num_specific_experts: int,
         shared_expert_params: dict,
         specific_expert_params: dict | list[dict],
     ):
         super().__init__()
-        if num_tasks < 1:
-            raise ValueError("num_tasks must be >= 1")
+        if nums_task < 1:
+            raise ValueError("nums_task must be >= 1")
 
         specific_params_list = self.normalize_specific_params(
-            specific_expert_params, num_tasks
+            specific_expert_params, nums_task
         )
 
         self.output_dim = get_mlp_output_dim(shared_expert_params, input_dim)
@@ -120,23 +121,23 @@ class CGCLayer(nn.Module):
                     nn.Linear(input_dim, task_gate_expert_num),
                     nn.Softmax(dim=1),
                 )
-                for _ in range(num_tasks)
+                for _ in range(nums_task)
             ]
         )
-        shared_gate_expert_num = num_shared_experts + num_specific_experts * num_tasks
+        shared_gate_expert_num = num_shared_experts + num_specific_experts * nums_task
         self.shared_gate = nn.Sequential(
             nn.Linear(input_dim, shared_gate_expert_num),
             nn.Softmax(dim=1),
         )
 
-        self.num_tasks = num_tasks
+        self.nums_task = nums_task
 
     def forward(
         self, task_inputs: list[torch.Tensor], shared_input: torch.Tensor
     ) -> tuple[list[torch.Tensor], torch.Tensor]:
-        if len(task_inputs) != self.num_tasks:
+        if len(task_inputs) != self.nums_task:
             raise ValueError(
-                f"Expected {self.num_tasks} task inputs, got {len(task_inputs)}"
+                f"Expected {self.nums_task} task inputs, got {len(task_inputs)}"
             )
 
         shared_outputs = [expert(shared_input) for expert in self.shared_experts]
@@ -145,7 +146,7 @@ class CGCLayer(nn.Module):
         new_task_fea: list[torch.Tensor] = []
         all_specific_for_shared: list[torch.Tensor] = []
 
-        for task_idx in range(self.num_tasks):
+        for task_idx in range(self.nums_task):
             task_input = task_inputs[task_idx]
             task_specific_outputs = [expert(task_input) for expert in self.specific_experts[task_idx]]  # type: ignore
             all_specific_for_shared.extend(task_specific_outputs)
@@ -167,15 +168,15 @@ class CGCLayer(nn.Module):
 
     @staticmethod
     def normalize_specific_params(
-        params: dict | list[dict], num_tasks: int
+        params: dict | list[dict], nums_task: int
     ) -> list[dict]:
         if isinstance(params, list):
-            if len(params) != num_tasks:
+            if len(params) != nums_task:
                 raise ValueError(
-                    f"Length of specific_expert_params ({len(params)}) must match num_tasks ({num_tasks})."
+                    f"Length of specific_expert_params ({len(params)}) must match nums_task ({nums_task})."
                 )
             return [p.copy() for p in params]
-        return [params.copy() for _ in range(num_tasks)]
+        return [params.copy() for _ in range(nums_task)]
 
 
 class PLE(BaseModel):
@@ -194,9 +195,9 @@ class PLE(BaseModel):
 
     @property
     def default_task(self):
-        num_tasks = getattr(self, "num_tasks", None)
-        if num_tasks is not None and num_tasks > 0:
-            return ["binary"] * num_tasks
+        nums_task = getattr(self, "nums_task", None)
+        if nums_task is not None and nums_task > 0:
+            return ["binary"] * nums_task
         return ["binary"]
 
     def __init__(
@@ -224,18 +225,18 @@ class PLE(BaseModel):
         **kwargs,
     ):
 
-        self.num_tasks = len(target)
+        self.nums_task = len(target)
 
         resolved_task = task
         if resolved_task is None:
             resolved_task = self.default_task
         elif isinstance(resolved_task, str):
-            resolved_task = [resolved_task] * self.num_tasks
-        elif len(resolved_task) == 1 and self.num_tasks > 1:
-            resolved_task = resolved_task * self.num_tasks
-        elif len(resolved_task) != self.num_tasks:
+            resolved_task = [resolved_task] * self.nums_task
+        elif len(resolved_task) == 1 and self.nums_task > 1:
+            resolved_task = resolved_task * self.nums_task
+        elif len(resolved_task) != self.nums_task:
             raise ValueError(
-                f"Length of task ({len(resolved_task)}) must match number of targets ({self.num_tasks})."
+                f"Length of task ({len(resolved_task)}) must match number of targets ({self.nums_task})."
             )
 
         super(PLE, self).__init__(
@@ -256,15 +257,15 @@ class PLE(BaseModel):
         if self.loss is None:
             self.loss = "bce"
         # Number of tasks, experts, and levels
-        self.num_tasks = len(target)
+        self.nums_task = len(target)
         self.num_shared_experts = num_shared_experts
         self.num_specific_experts = num_specific_experts
         self.num_levels = num_levels
         if optimizer_params is None:
             optimizer_params = {}
-        if len(tower_params_list) != self.num_tasks:
+        if len(tower_params_list) != self.nums_task:
             raise ValueError(
-                f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.num_tasks})"
+                f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.nums_task})"
             )
         # Embedding layer
         self.embedding = EmbeddingLayer(features=self.all_features)
@@ -287,7 +288,7 @@ class PLE(BaseModel):
             level_input_dim = input_dim if level == 0 else expert_output_dim
             cgc_layer = CGCLayer(
                 input_dim=level_input_dim,
-                num_tasks=self.num_tasks,
+                nums_task=self.nums_task,
                 num_shared_experts=num_shared_experts,
                 num_specific_experts=num_specific_experts,
                 shared_expert_params=shared_expert_params,
@@ -302,8 +303,8 @@ class PLE(BaseModel):
         for tower_params in tower_params_list:
             tower = MLP(input_dim=expert_output_dim, output_layer=True, **tower_params)
             self.towers.append(tower)
-        self.prediction_layer = PredictionLayer(
-            task_type=self.default_task, task_dims=[1] * self.num_tasks
+        self.prediction_layer = TaskHead(
+            task_type=self.default_task, task_dims=[1] * self.nums_task
         )
         # Register regularization weights
         self.register_regularization_weights(
@@ -321,7 +322,7 @@ class PLE(BaseModel):
         input_flat = self.embedding(x=x, features=self.all_features, squeeze_dim=True)
 
         # Initial features for each task and shared
-        task_fea = [input_flat for _ in range(self.num_tasks)]
+        task_fea = [input_flat for _ in range(self.nums_task)]
         shared_fea = input_flat
 
         # Progressive Layered Extraction: CGC
@@ -330,10 +331,10 @@ class PLE(BaseModel):
 
         # task tower
         task_outputs = []
-        for task_idx in range(self.num_tasks):
+        for task_idx in range(self.nums_task):
             tower_output = self.towers[task_idx](task_fea[task_idx])  # [B, 1]
             task_outputs.append(tower_output)
 
-        # [B, num_tasks]
+        # [B, nums_task]
         y = torch.cat(task_outputs, dim=1)
         return self.prediction_layer(y)

nextrec/models/multi_task/poso.py

@@ -1,5 +1,6 @@
 """
 Date: create on 28/11/2025
+Checkpoint: edit on 23/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 Reference:
 [1] Wang et al. "POSO: Personalized Cold Start Modules for Large-scale Recommender Systems", 2021.
@@ -44,7 +45,8 @@ import torch.nn.functional as F
 
 from nextrec.basic.activation import activation_layer
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 from nextrec.utils.model import select_features
 
@@ -195,7 +197,7 @@ class POSOMMoE(nn.Module):
         pc_dim: int,  # for poso feature dimension
         num_experts: int,
         expert_hidden_dims: list[int],
-        num_tasks: int,
+        nums_task: int,
         activation: str = "relu",
         expert_dropout: float = 0.0,
         gate_hidden_dim: int = 32,  # for poso gate hidden dimension
@@ -204,7 +206,7 @@ class POSOMMoE(nn.Module):
     ) -> None:
         super().__init__()
         self.num_experts = num_experts
-        self.num_tasks = num_tasks
+        self.nums_task = nums_task
 
         # Experts built with framework MLP, same as standard MMoE
         self.experts = nn.ModuleList(
@@ -225,7 +227,7 @@ class POSOMMoE(nn.Module):
 
         # Task-specific gates: gate_t(x) over experts
         self.gates = nn.ModuleList(
-            [nn.Linear(input_dim, num_experts) for _ in range(num_tasks)]
+            [nn.Linear(input_dim, num_experts) for _ in range(nums_task)]
         )
         self.gate_use_softmax = gate_use_softmax
 
@@ -247,7 +249,7 @@ class POSOMMoE(nn.Module):
         """
         x:  (B, input_dim)
         pc: (B, pc_dim)
-        return: list of task outputs z_t with length num_tasks, each (B, D)
+        return: list of task outputs z_t with length nums_task, each (B, D)
         """
         # 1) Expert outputs with POSO PC gate
         masked_expert_outputs = []
@@ -261,7 +263,7 @@ class POSOMMoE(nn.Module):
 
         # 2) Task gates depend on x as in standard MMoE
         task_outputs: list[torch.Tensor] = []
-        for t in range(self.num_tasks):
+        for t in range(self.nums_task):
             logits = self.gates[t](x)  # (B, E)
             if self.gate_use_softmax:
                 gate = F.softmax(logits, dim=1)
@@ -288,9 +290,9 @@ class POSO(BaseModel):
 
     @property
     def default_task(self) -> list[str]:
-        num_tasks = getattr(self, "num_tasks", None)
-        if num_tasks is not None and num_tasks > 0:
-            return ["binary"] * num_tasks
+        nums_task = getattr(self, "nums_task", None)
+        if nums_task is not None and nums_task > 0:
+            return ["binary"] * nums_task
         return ["binary"]
 
     def __init__(
@@ -332,24 +334,24 @@ class POSO(BaseModel):
         dense_l2_reg: float = 1e-4,
         **kwargs,
     ):
-        self.num_tasks = len(target)
+        self.nums_task = len(target)
 
-        # Normalize task to match num_tasks
+        # Normalize task to match nums_task
         resolved_task = task
         if resolved_task is None:
             resolved_task = self.default_task
         elif isinstance(resolved_task, str):
-            resolved_task = [resolved_task] * self.num_tasks
-        elif len(resolved_task) == 1 and self.num_tasks > 1:
-            resolved_task = resolved_task * self.num_tasks
-        elif len(resolved_task) != self.num_tasks:
+            resolved_task = [resolved_task] * self.nums_task
+        elif len(resolved_task) == 1 and self.nums_task > 1:
+            resolved_task = resolved_task * self.nums_task
+        elif len(resolved_task) != self.nums_task:
             raise ValueError(
-                f"Length of task ({len(resolved_task)}) must match number of targets ({self.num_tasks})."
+                f"Length of task ({len(resolved_task)}) must match number of targets ({self.nums_task})."
             )
 
-        if len(tower_params_list) != self.num_tasks:
+        if len(tower_params_list) != self.nums_task:
             raise ValueError(
-                f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.num_tasks})"
+                f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.nums_task})"
             )
 
         super().__init__(
@@ -465,7 +467,7 @@ class POSO(BaseModel):
                 pc_dim=self.pc_input_dim,
                 num_experts=num_experts,
                 expert_hidden_dims=expert_hidden_dims,
-                num_tasks=self.num_tasks,
+                nums_task=self.nums_task,
                 activation=expert_activation,
                 expert_dropout=expert_dropout,
                 gate_hidden_dim=expert_gate_hidden_dim,
@@ -487,9 +489,9 @@ class POSO(BaseModel):
             self.grad_norm_shared_modules = ["embedding"]
         else:
             self.grad_norm_shared_modules = ["embedding", "mmoe"]
-        self.prediction_layer = PredictionLayer(
+        self.prediction_layer = TaskHead(
             task_type=self.default_task,
-            task_dims=[1] * self.num_tasks,
+            task_dims=[1] * self.nums_task,
         )
         include_modules = (
             ["towers", "tower_heads"]

nextrec/models/multi_task/share_bottom.py

@@ -1,6 +1,6 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 24/11/2025
+Checkpoint: edit on 23/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 Reference:
 [1] Caruana R. Multitask learning[J]. Machine Learning, 1997, 28: 41-75.
@@ -43,7 +43,8 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -54,9 +55,9 @@ class ShareBottom(BaseModel):
 
     @property
     def default_task(self):
-        num_tasks = getattr(self, "num_tasks", None)
-        if num_tasks is not None and num_tasks > 0:
-            return ["binary"] * num_tasks
+        nums_task = getattr(self, "nums_task", None)
+        if nums_task is not None and nums_task > 0:
+            return ["binary"] * nums_task
         return ["binary"]
 
     def __init__(
@@ -82,18 +83,18 @@ class ShareBottom(BaseModel):
 
         optimizer_params = optimizer_params or {}
 
-        self.num_tasks = len(target)
+        self.nums_task = len(target)
 
         resolved_task = task
         if resolved_task is None:
             resolved_task = self.default_task
         elif isinstance(resolved_task, str):
-            resolved_task = [resolved_task] * self.num_tasks
-        elif len(resolved_task) == 1 and self.num_tasks > 1:
-            resolved_task = resolved_task * self.num_tasks
-        elif len(resolved_task) != self.num_tasks:
+            resolved_task = [resolved_task] * self.nums_task
+        elif len(resolved_task) == 1 and self.nums_task > 1:
+            resolved_task = resolved_task * self.nums_task
+        elif len(resolved_task) != self.nums_task:
             raise ValueError(
-                f"Length of task ({len(resolved_task)}) must match number of targets ({self.num_tasks})."
+                f"Length of task ({len(resolved_task)}) must match number of targets ({self.nums_task})."
             )
 
         super(ShareBottom, self).__init__(
@@ -114,10 +115,10 @@ class ShareBottom(BaseModel):
         if self.loss is None:
             self.loss = "bce"
         # Number of tasks
-        self.num_tasks = len(target)
-        if len(tower_params_list) != self.num_tasks:
+        self.nums_task = len(target)
+        if len(tower_params_list) != self.nums_task:
             raise ValueError(
-                f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.num_tasks})"
+                f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.nums_task})"
             )
         # Embedding layer
         self.embedding = EmbeddingLayer(features=self.all_features)
@@ -142,8 +143,8 @@ class ShareBottom(BaseModel):
         for tower_params in tower_params_list:
             tower = MLP(input_dim=bottom_output_dim, output_layer=True, **tower_params)
             self.towers.append(tower)
-        self.prediction_layer = PredictionLayer(
-            task_type=self.default_task, task_dims=[1] * self.num_tasks
+        self.prediction_layer = TaskHead(
+            task_type=self.default_task, task_dims=[1] * self.nums_task
         )
         # Register regularization weights
         self.register_regularization_weights(
@@ -169,6 +170,6 @@ class ShareBottom(BaseModel):
             tower_output = tower(bottom_output)  # [B, 1]
             task_outputs.append(tower_output)
 
-        # Stack outputs: [B, num_tasks]
+        # Stack outputs: [B, nums_task]
         y = torch.cat(task_outputs, dim=1)
         return self.prediction_layer(y)