nextrec 0.4.21__py3-none-any.whl → 0.4.23__py3-none-any.whl

nextrec/data/preprocessor.py

@@ -2,7 +2,7 @@
 DataProcessor for data preprocessing including numeric, sparse, sequence features and target processing.
 
 Date: create on 13/11/2025
-Checkpoint: edit on 24/12/2025
+Checkpoint: edit on 29/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
@@ -79,6 +79,14 @@ class DataProcessor(FeatureSet):
         ] = "standard",
         fill_na: Optional[float] = None,
     ):
+        """Add a numeric feature configuration.
+
+        Args:
+            name (str): Feature name.
+            scaler (Optional[Literal["standard", "minmax", "robust", "maxabs", "log", "none"]], optional): Scaler type. Defaults to "standard".
+            fill_na (Optional[float], optional): Fill value for missing entries. Defaults to None.
+        """
+
         self.numeric_features[name] = {"scaler": scaler, "fill_na": fill_na}
 
     def add_sparse_feature(
@@ -88,6 +96,14 @@ class DataProcessor(FeatureSet):
         hash_size: Optional[int] = None,
         fill_na: str = "<UNK>",
     ):
+        """Add a sparse feature configuration.
+
+        Args:
+            name (str): Feature name.
+            encode_method (Literal["hash", "label"], optional): Encoding method, including "hash encoding" and "label encoding". Defaults to "label".
+            hash_size (Optional[int], optional): Hash size for hash encoding. Required if encode_method is "hash".
+            fill_na (str, optional): Fill value for missing entries. Defaults to "<UNK>".
+        """
         if encode_method == "hash" and hash_size is None:
             raise ValueError(
                 "[Data Processor Error] hash_size must be specified when encode_method='hash'"
@@ -101,7 +117,7 @@ class DataProcessor(FeatureSet):
     def add_sequence_feature(
         self,
         name: str,
-        encode_method: Literal["hash", "label"] = "label",
+        encode_method: Literal["hash", "label"] = "hash",
         hash_size: Optional[int] = None,
         max_len: Optional[int] = 50,
         pad_value: int = 0,
@@ -110,6 +126,17 @@ class DataProcessor(FeatureSet):
         ] = "pre",  # pre: keep last max_len items, post: keep first max_len items
         separator: str = ",",
     ):
+        """Add a sequence feature configuration.
+
+        Args:
+            name (str): Feature name.
+            encode_method (Literal["hash", "label"], optional): Encoding method, including "hash encoding" and "label encoding". Defaults to "hash".
+            hash_size (Optional[int], optional): Hash size for hash encoding. Required if encode_method is "hash".
+            max_len (Optional[int], optional): Maximum sequence length. Defaults to 50.
+            pad_value (int, optional): Padding value for sequences shorter than max_len. Defaults to 0.
+            truncate (Literal["pre", "post"], optional): Truncation strategy for sequences longer than max_len, including "pre" (keep last max_len items) and "post" (keep first max_len items). Defaults to "pre".
+            separator (str, optional): Separator for string sequences. Defaults to ",".
+        """
         if encode_method == "hash" and hash_size is None:
             raise ValueError(
                 "[Data Processor Error] hash_size must be specified when encode_method='hash'"
@@ -131,6 +158,14 @@ class DataProcessor(FeatureSet):
             Dict[str, int]
         ] = None,  # example: {'click': 1, 'no_click': 0}
     ):
+        """Add a target configuration.
+
+        Args:
+            name (str): Target name.
+            target_type (Literal["binary", "regression"], optional): Target type. Defaults to "binary".
+            label_map (Optional[Dict[str, int]], optional): Label mapping for binary targets. Defaults to None.
+        """
+
         self.target_features[name] = {
             "target_type": target_type,
             "label_map": label_map,
@@ -392,7 +427,15 @@ class DataProcessor(FeatureSet):
         )
 
     def load_dataframe_from_path(self, path: str) -> pd.DataFrame:
-        """Load all data from a file or directory path into a single DataFrame."""
+        """
+        Load all data from a file or directory path into a single DataFrame.
+
+        Args:
+            path (str): File or directory path.
+
+        Returns:
+            pd.DataFrame: Loaded DataFrame.
+        """
         file_paths, file_type = resolve_file_paths(path)
         frames = load_dataframes(file_paths, file_type)
         return pd.concat(frames, ignore_index=True) if len(frames) > 1 else frames[0]
@@ -411,7 +454,16 @@ class DataProcessor(FeatureSet):
         return [str(value)]
 
     def fit_from_path(self, path: str, chunk_size: int) -> "DataProcessor":
-        """Fit processor statistics by streaming files to reduce memory usage."""
+        """
+        Fit processor statistics by streaming files to reduce memory usage.
+
+        Args:
+            path (str): File or directory path.
+            chunk_size (int): Number of rows per chunk.
+
+        Returns:
+            DataProcessor: Fitted DataProcessor instance.
+        """
         logger = logging.getLogger()
         logger.info(
             colorize(
@@ -428,7 +480,7 @@ class DataProcessor(FeatureSet):
                 "Use fit(dataframe) with in-memory data or convert the data format."
             )
 
-        numeric_acc: Dict[str, Dict[str, float]] = {}
+        numeric_acc = {}
         for name in self.numeric_features.keys():
             numeric_acc[name] = {
                 "sum": 0.0,
@@ -609,6 +661,21 @@ class DataProcessor(FeatureSet):
         output_path: Optional[str],
         warn_missing: bool = True,
     ):
+        """
+        Transform in-memory data and optionally persist the transformed data.
+
+        Args:
+            data (Union[pd.DataFrame, Dict[str, Any]]): Input data.
+            return_dict (bool): Whether to return a dictionary of numpy arrays.
+            persist (bool): Whether to persist the transformed data to disk.
+            save_format (Optional[str]): Format to save the data if persisting.
+            output_path (Optional[str]): Output path to save the data if persisting.
+            warn_missing (bool): Whether to warn about missing features in the data.
+
+        Returns:
+            Union[pd.DataFrame, Dict[str, np.ndarray]]: Transformed data.
+        """
+
         logger = logging.getLogger()
         data_dict = data if isinstance(data, dict) else None
 
@@ -719,6 +786,12 @@ class DataProcessor(FeatureSet):
         """Transform data from files under a path and save them to a new location.
 
         Uses chunked reading/writing to keep peak memory bounded for large files.
+
+        Args:
+            input_path (str): Input file or directory path.
+            output_path (Optional[str]): Output directory path. If None, defaults to input_path/transformed_data.
+            save_format (Optional[str]): Format to save transformed files. If None, uses input file format.
+            chunk_size (int): Number of rows per chunk.
         """
         logger = logging.getLogger()
         file_paths, file_type = resolve_file_paths(input_path)
@@ -876,6 +949,17 @@ class DataProcessor(FeatureSet):
         data: Union[pd.DataFrame, Dict[str, Any], str, os.PathLike],
         chunk_size: int = 200000,
     ):
+        """
+        Fit the DataProcessor to the provided data.
+
+        Args:
+            data (Union[pd.DataFrame, Dict[str, Any], str, os.PathLike]): Input data for fitting.
+            chunk_size (int): Number of rows per chunk when streaming from path.
+
+        Returns:
+            DataProcessor: Fitted DataProcessor instance.
+        """
+
         logger = logging.getLogger()
         if isinstance(data, (str, os.PathLike)):
             path_str = str(data)
@@ -915,6 +999,19 @@ class DataProcessor(FeatureSet):
         output_path: Optional[str] = None,
         chunk_size: int = 200000,
     ):
+        """
+        Transform the provided data using the fitted DataProcessor.
+
+        Args:
+            data (Union[pd.DataFrame, Dict[str, Any], str, os.PathLike]): Input data to transform.
+            return_dict (bool): Whether to return a dictionary of numpy arrays.
+            save_format (Optional[str]): Format to save the data if output_path is provided.
+            output_path (Optional[str]): Output path to save the transformed data.
+            chunk_size (int): Number of rows per chunk when streaming from path.
+        Returns:
+            Union[pd.DataFrame, Dict[str, np.ndarray], List[str]]: Transformed data or list of saved file paths.
+        """
+
         if not self.is_fitted:
             raise ValueError(
                 "[Data Processor Error] DataProcessor must be fitted before transform"
@@ -943,6 +1040,19 @@ class DataProcessor(FeatureSet):
         output_path: Optional[str] = None,
         chunk_size: int = 200000,
     ):
+        """
+        Fit the DataProcessor to the provided data and then transform it.
+
+        Args:
+            data (Union[pd.DataFrame, Dict[str, Any], str, os.PathLike]): Input data for fitting and transforming.
+            return_dict (bool): Whether to return a dictionary of numpy arrays.
+            save_format (Optional[str]): Format to save the data if output_path is provided.
+            output_path (Optional[str]): Output path to save the transformed data.
+            chunk_size (int): Number of rows per chunk when streaming from path.
+        Returns:
+            Union[pd.DataFrame, Dict[str, np.ndarray], List[str]]: Transformed data or list of saved file paths.
+        """
+
         self.fit(data, chunk_size=chunk_size)
         return self.transform(
             data,
@@ -952,6 +1062,12 @@ class DataProcessor(FeatureSet):
         )
 
     def save(self, save_path: str | Path):
+        """
+        Save the fitted DataProcessor to a file.
+
+        Args:
+            save_path (str | Path): Path to save the DataProcessor.
+        """
         logger = logging.getLogger()
         assert isinstance(save_path, (str, Path)), "save_path must be a string or Path"
         save_path = Path(save_path)
@@ -983,6 +1099,16 @@ class DataProcessor(FeatureSet):
 
     @classmethod
     def load(cls, load_path: str | Path) -> "DataProcessor":
+        """
+        Load a fitted DataProcessor from a file.
+
+        Args:
+            load_path (str | Path): Path to load the DataProcessor from.
+
+        Returns:
+            DataProcessor: Loaded DataProcessor instance.
+        """
+
         logger = logging.getLogger()
         load_path = Path(load_path)
         with open(load_path, "rb") as f:
@@ -1003,6 +1129,12 @@ class DataProcessor(FeatureSet):
         return processor
 
     def get_vocab_sizes(self) -> Dict[str, int]:
+        """
+        Get vocabulary sizes for all sparse and sequence features.
+
+        Returns:
+            Dict[str, int]: Mapping of feature names to vocabulary sizes.
+        """
         vocab_sizes = {}
         for name, config in self.sparse_features.items():
             vocab_sizes[name] = config.get("vocab_size", 0)

nextrec/loss/__init__.py

@@ -6,7 +6,6 @@ from nextrec.loss.listwise import (
     SampledSoftmaxLoss,
 )
 from nextrec.loss.grad_norm import GradNormLossWeighting
-from nextrec.loss.loss_utils import VALID_TASK_TYPES, get_loss_fn, get_loss_kwargs
 from nextrec.loss.pairwise import BPRLoss, HingeLoss, TripletLoss
 from nextrec.loss.pointwise import (
     ClassBalancedFocalLoss,
@@ -34,7 +33,4 @@ __all__ = [
     # Multi-task weighting
     "GradNormLossWeighting",
     # Utilities
-    "get_loss_fn",
-    "get_loss_kwargs",
-    "VALID_TASK_TYPES",
 ]

nextrec/loss/grad_norm.py

@@ -20,9 +20,9 @@ import torch.nn.functional as F
 
 
 def get_grad_norm_shared_params(
-    model: torch.nn.Module,
-    shared_modules: Iterable[str] | None = None,
-) -> list[torch.nn.Parameter]:
+    model,
+    shared_modules=None,
+):
     if not shared_modules:
         return [p for p in model.parameters() if p.requires_grad]
     shared_params = []

nextrec/loss/listwise.py

@@ -2,10 +2,11 @@
 Listwise loss functions for ranking and contrastive training.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 29/11/2025
+Checkpoint: edit on 29/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
+from typing import Literal
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -16,7 +17,7 @@ class SampledSoftmaxLoss(nn.Module):
     Softmax over one positive and multiple sampled negatives.
     """
 
-    def __init__(self, reduction: str = "mean"):
+    def __init__(self, reduction: Literal["mean", "sum", "none"] = "mean"):
         super().__init__()
         self.reduction = reduction
 
@@ -37,7 +38,11 @@ class InfoNCELoss(nn.Module):
     InfoNCE loss for contrastive learning with one positive and many negatives.
     """
 
-    def __init__(self, temperature: float = 0.07, reduction: str = "mean"):
+    def __init__(
+        self,
+        temperature: float = 0.07,
+        reduction: Literal["mean", "sum", "none"] = "mean",
+    ):
         super().__init__()
         self.temperature = temperature
         self.reduction = reduction
@@ -61,7 +66,11 @@ class ListNetLoss(nn.Module):
     Reference: Cao et al. (ICML 2007)
     """
 
-    def __init__(self, temperature: float = 1.0, reduction: str = "mean"):
+    def __init__(
+        self,
+        temperature: float = 1.0,
+        reduction: Literal["mean", "sum", "none"] = "mean",
+    ):
         super().__init__()
         self.temperature = temperature
         self.reduction = reduction
@@ -84,7 +93,7 @@ class ListMLELoss(nn.Module):
     Reference: Xia et al. (ICML 2008)
     """
 
-    def __init__(self, reduction: str = "mean"):
+    def __init__(self, reduction: Literal["mean", "sum", "none"] = "mean"):
         super().__init__()
         self.reduction = reduction
 
@@ -117,7 +126,11 @@ class ApproxNDCGLoss(nn.Module):
     Reference: Qin et al. (2010)
     """
 
-    def __init__(self, temperature: float = 1.0, reduction: str = "mean"):
+    def __init__(
+        self,
+        temperature: float = 1.0,
+        reduction: Literal["mean", "sum", "none"] = "mean",
+    ):
         super().__init__()
         self.temperature = temperature
         self.reduction = reduction

nextrec/loss/pairwise.py

@@ -2,7 +2,7 @@
 Pairwise loss functions for learning-to-rank and matching tasks.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 29/11/2025
+Checkpoint: edit on 29/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
@@ -18,7 +18,7 @@ class BPRLoss(nn.Module):
     Bayesian Personalized Ranking loss with support for multiple negatives.
     """
 
-    def __init__(self, reduction: str = "mean"):
+    def __init__(self, reduction: Literal["mean", "sum", "none"] = "mean"):
         super().__init__()
         self.reduction = reduction
 
@@ -42,7 +42,9 @@ class HingeLoss(nn.Module):
     Hinge loss for pairwise ranking.
     """
 
-    def __init__(self, margin: float = 1.0, reduction: str = "mean"):
+    def __init__(
+        self, margin: float = 1.0, reduction: Literal["mean", "sum", "none"] = "mean"
+    ):
         super().__init__()
         self.margin = margin
         self.reduction = reduction
@@ -69,7 +71,7 @@ class TripletLoss(nn.Module):
     def __init__(
         self,
         margin: float = 1.0,
-        reduction: str = "mean",
+        reduction: Literal["mean", "sum", "none"] = "mean",
         distance: Literal["euclidean", "cosine"] = "euclidean",
     ):
         super().__init__()

nextrec/loss/pointwise.py

@@ -2,11 +2,11 @@
 Pointwise loss functions, including imbalance-aware variants.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 29/11/2025
+Checkpoint: edit on 29/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
-from typing import Optional, Sequence
+from typing import Optional, Sequence, Literal
 
 import torch
 import torch.nn as nn
@@ -18,7 +18,9 @@ class CosineContrastiveLoss(nn.Module):
     Contrastive loss using cosine similarity for positive/negative pairs.
     """
 
-    def __init__(self, margin: float = 0.5, reduction: str = "mean"):
+    def __init__(
+        self, margin: float = 0.5, reduction: Literal["mean", "sum", "none"] = "mean"
+    ):
         super().__init__()
         self.margin = margin
         self.reduction = reduction
@@ -50,7 +52,7 @@ class WeightedBCELoss(nn.Module):
     def __init__(
         self,
         pos_weight: float | torch.Tensor | None = None,
-        reduction: str = "mean",
+        reduction: Literal["mean", "sum", "none"] = "mean",
         logits: bool = False,
         auto_balance: bool = False,
     ):
@@ -110,7 +112,7 @@ class FocalLoss(nn.Module):
         self,
         gamma: float = 2.0,
         alpha: Optional[float | Sequence[float] | torch.Tensor] = None,
-        reduction: str = "mean",
+        reduction: Literal["mean", "sum", "none"] = "mean",
         logits: bool = False,
     ):
         super().__init__()
@@ -187,7 +189,7 @@ class ClassBalancedFocalLoss(nn.Module):
         class_counts: Sequence[int] | torch.Tensor,
         beta: float = 0.9999,
         gamma: float = 2.0,
-        reduction: str = "mean",
+        reduction: Literal["mean", "sum", "none"] = "mean",
     ):
         super().__init__()
         self.gamma = gamma

nextrec/models/multi_task/esmm.py

@@ -42,12 +42,12 @@ CVR 预测 P(conversion|click)，二者相乘得到 CTCVR 并在曝光标签上
 """
 
 import torch
-import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
 from nextrec.basic.layers import MLP, EmbeddingLayer
 from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
+from nextrec.utils.types import TaskTypeName
 
 
 class ESMM(BaseModel):
@@ -77,23 +77,12 @@ class ESMM(BaseModel):
         sequence_features: list[SequenceFeature],
         ctr_params: dict,
         cvr_params: dict,
+        task: TaskTypeName | list[TaskTypeName] | None = None,
         target: list[str] | None = None,  # Note: ctcvr = ctr * cvr
-        task: list[str] | None = None,
-        optimizer: str = "adam",
-        optimizer_params: dict | None = None,
-        loss: str | nn.Module | list[str | nn.Module] | None = "bce",
-        loss_params: dict | list[dict] | None = None,
-        embedding_l1_reg=0.0,
-        dense_l1_reg=0.0,
-        embedding_l2_reg=0.0,
-        dense_l2_reg=0.0,
         **kwargs,
     ):
 
         target = target or ["ctr", "ctcvr"]
-        optimizer_params = optimizer_params or {}
-        if loss is None:
-            loss = "bce"
 
         if len(target) != 2:
             raise ValueError(
@@ -120,15 +109,9 @@ class ESMM(BaseModel):
             sequence_features=sequence_features,
             target=target,
             task=resolved_task,  # Both CTR and CTCVR are binary classification
-            embedding_l1_reg=embedding_l1_reg,
-            dense_l1_reg=dense_l1_reg,
-            embedding_l2_reg=embedding_l2_reg,
-            dense_l2_reg=dense_l2_reg,
             **kwargs,
         )
 
-        self.loss = loss
-
         self.embedding = EmbeddingLayer(features=self.all_features)
         input_dim = self.embedding.input_dim
 
@@ -138,17 +121,11 @@ 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 = TaskHead(task_type=self.default_task, task_dims=[1, 1])
+        self.prediction_layer = TaskHead(task_type=self.task, task_dims=[1, 1])
         # Register regularization weights
         self.register_regularization_weights(
             embedding_attr="embedding", include_modules=["ctr_tower", "cvr_tower"]
         )
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=loss,
-            loss_params=loss_params,
-        )
 
     def forward(self, x):
         # Get all embeddings and flatten

nextrec/models/multi_task/mmoe.py

@@ -82,14 +82,6 @@ class MMOE(BaseModel):
         tower_params_list: list[dict] | None = None,
         target: list[str] | str | None = None,
         task: str | list[str] = "binary",
-        optimizer: str = "adam",
-        optimizer_params: dict | None = None,
-        loss: str | nn.Module | list[str | nn.Module] | None = "bce",
-        loss_params: dict | list[dict] | None = None,
-        embedding_l1_reg=0.0,
-        dense_l1_reg=0.0,
-        embedding_l2_reg=0.0,
-        dense_l2_reg=0.0,
         **kwargs,
     ):
 
@@ -98,9 +90,7 @@ class MMOE(BaseModel):
         sequence_features = sequence_features or []
         expert_params = expert_params or {}
         tower_params_list = tower_params_list or []
-        optimizer_params = optimizer_params or {}
-        if loss is None:
-            loss = "bce"
+
         if target is None:
             target = []
         elif isinstance(target, str):
@@ -126,15 +116,9 @@ class MMOE(BaseModel):
             sequence_features=sequence_features,
             target=target,
             task=resolved_task,
-            embedding_l1_reg=embedding_l1_reg,
-            dense_l1_reg=dense_l1_reg,
-            embedding_l2_reg=embedding_l2_reg,
-            dense_l2_reg=dense_l2_reg,
             **kwargs,
         )
 
-        self.loss = loss
-
         # Number of tasks and experts
         self.nums_task = len(target)
         self.num_experts = num_experts
@@ -172,18 +156,12 @@ class MMOE(BaseModel):
             tower = MLP(input_dim=expert_output_dim, output_layer=True, **tower_params)
             self.towers.append(tower)
         self.prediction_layer = TaskHead(
-            task_type=self.default_task, task_dims=[1] * self.nums_task
+            task_type=self.task, task_dims=[1] * self.nums_task
         )
         # Register regularization weights
         self.register_regularization_weights(
             embedding_attr="embedding", include_modules=["experts", "gates", "towers"]
         )
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=self.loss,
-            loss_params=loss_params,
-        )
 
     def forward(self, x):
         # Get all embeddings and flatten

nextrec/models/multi_task/ple.py

@@ -202,29 +202,21 @@ class PLE(BaseModel):
 
     def __init__(
         self,
-        dense_features: list[DenseFeature],
-        sparse_features: list[SparseFeature],
-        sequence_features: list[SequenceFeature],
-        shared_expert_params: dict,
-        specific_expert_params: dict | list[dict],
-        num_shared_experts: int,
-        num_specific_experts: int,
-        num_levels: int,
-        tower_params_list: list[dict],
-        target: list[str],
+        dense_features: list[DenseFeature] | None = None,
+        sparse_features: list[SparseFeature] | None = None,
+        sequence_features: list[SequenceFeature] | None = None,
+        shared_expert_params: dict | None = None,
+        specific_expert_params: dict | list[dict] | None = None,
+        num_shared_experts: int = 2,
+        num_specific_experts: int = 2,
+        num_levels: int = 2,
+        tower_params_list: list[dict] | None = None,
+        target: list[str] | None = None,
         task: str | list[str] | None = None,
-        optimizer: str = "adam",
-        optimizer_params: dict | None = None,
-        loss: str | nn.Module | list[str | nn.Module] | None = "bce",
-        loss_params: dict | list[dict] | None = None,
-        embedding_l1_reg=0.0,
-        dense_l1_reg=0.0,
-        embedding_l2_reg=0.0,
-        dense_l2_reg=0.0,
         **kwargs,
     ):
 
-        self.nums_task = len(target)
+        self.nums_task = len(target) if target is not None else 1
 
         resolved_task = task
         if resolved_task is None:
@@ -244,23 +236,15 @@ class PLE(BaseModel):
             sequence_features=sequence_features,
             target=target,
             task=resolved_task,
-            embedding_l1_reg=embedding_l1_reg,
-            dense_l1_reg=dense_l1_reg,
-            embedding_l2_reg=embedding_l2_reg,
-            dense_l2_reg=dense_l2_reg,
             **kwargs,
         )
 
-        self.loss = loss
-        if self.loss is None:
-            self.loss = "bce"
         # Number of tasks, experts, and levels
         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.nums_task:
             raise ValueError(
                 f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.nums_task})"
@@ -302,18 +286,12 @@ class PLE(BaseModel):
             tower = MLP(input_dim=expert_output_dim, output_layer=True, **tower_params)
             self.towers.append(tower)
         self.prediction_layer = TaskHead(
-            task_type=self.default_task, task_dims=[1] * self.nums_task
+            task_type=self.task, task_dims=[1] * self.nums_task
         )
         # Register regularization weights
         self.register_regularization_weights(
             embedding_attr="embedding", include_modules=["cgc_layers", "towers"]
         )
-        self.compile(
-            optimizer=optimizer,
-            optimizer_params=optimizer_params,
-            loss=self.loss,
-            loss_params=loss_params,
-        )
 
     def forward(self, x):
         # Get all embeddings and flatten