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

nextrec/__version__.py

@@ -1 +1 @@
-__version__ = "0.4.2"
+__version__ = "0.4.4"

nextrec/basic/layers.py

@@ -51,11 +51,11 @@ class PredictionLayer(nn.Module):
 
         # slice offsets per task
         start = 0
-        self._task_slices: list[tuple[int, int]] = []
+        self.task_slices: list[tuple[int, int]] = []
         for dim in self.task_dims:
             if dim < 1:
                 raise ValueError("Each task dimension must be >= 1.")
-            self._task_slices.append((start, start + dim))
+            self.task_slices.append((start, start + dim))
             start += dim
         if use_bias:
             self.bias = nn.Parameter(torch.zeros(self.total_dim))
@@ -71,7 +71,7 @@ class PredictionLayer(nn.Module):
             )
         logits = x if self.bias is None else x + self.bias
         outputs = []
-        for task_type, (start, end) in zip(self.task_types, self._task_slices):
+        for task_type, (start, end) in zip(self.task_types, self.task_slices):
             task_logits = logits[..., start:end]  # logits for the current task
             if self.return_logits:
                 outputs.append(task_logits)
@@ -367,20 +367,29 @@ class MLP(nn.Module):
         dims: list[int] | None = None,
         dropout: float = 0.0,
         activation: str = "relu",
+        use_norm: bool = True,
+        norm_type: str = "layer_norm",
     ):
         super().__init__()
         if dims is None:
             dims = []
         layers = []
         current_dim = input_dim
-
         for i_dim in dims:
             layers.append(nn.Linear(current_dim, i_dim))
-            layers.append(nn.BatchNorm1d(i_dim))
+            if use_norm:
+                if norm_type == "batch_norm":
+                    # **IMPORTANT** be careful when using BatchNorm1d in distributed training, nextrec does not support sync batch norm now
+                    layers.append(nn.BatchNorm1d(i_dim))
+                elif norm_type == "layer_norm":
+                    layers.append(nn.LayerNorm(i_dim))
+                else:
+                    raise ValueError(f"Unsupported norm_type: {norm_type}")
+
             layers.append(activation_layer(activation))
             layers.append(nn.Dropout(p=dropout))
             current_dim = i_dim
-
+        # output layer
         if output_layer:
             layers.append(nn.Linear(current_dim, 1))
             self.output_dim = 1
@@ -471,6 +480,21 @@ class BiLinearInteractionLayer(nn.Module):
         return torch.cat(bilinear_list, dim=1)
 
 
+class HadamardInteractionLayer(nn.Module):
+    """Hadamard interaction layer for Deep-FiBiNET (0 case in 01/11)."""
+
+    def __init__(self, num_fields: int):
+        super().__init__()
+        self.num_fields = num_fields
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B, F, D]
+        feature_emb = torch.split(x, 1, dim=1)  # list of F tensors [B,1,D]
+
+        hadamard_list = [v_i * v_j for (v_i, v_j) in combinations(feature_emb, 2)]
+        return torch.cat(hadamard_list, dim=1)  # [B, num_pairs, D]
+
+
 class MultiHeadSelfAttention(nn.Module):
     def __init__(
         self,
@@ -542,7 +566,7 @@ class AttentionPoolingLayer(nn.Module):
         embedding_dim: int,
         hidden_units: list = [80, 40],
         activation: str = "sigmoid",
-        use_softmax: bool = True,
+        use_softmax: bool = False,
     ):
         super().__init__()
         self.embedding_dim = embedding_dim
@@ -553,7 +577,7 @@ class AttentionPoolingLayer(nn.Module):
         layers = []
         for hidden_unit in hidden_units:
             layers.append(nn.Linear(input_dim, hidden_unit))
-            layers.append(activation_layer(activation))
+            layers.append(activation_layer(activation, emb_size=hidden_unit))
             input_dim = hidden_unit
         layers.append(nn.Linear(input_dim, 1))
         self.attention_net = nn.Sequential(*layers)

nextrec/basic/loggers.py

@@ -103,7 +103,7 @@ def setup_logger(session_id: str | os.PathLike | None = None):
     session = create_session(str(session_id) if session_id is not None else None)
     log_dir = session.logs_dir
     log_dir.mkdir(parents=True, exist_ok=True)
-    log_file = log_dir / f"{session.log_basename}.log"
+    log_file = log_dir / "runs.log"
 
     console_format = "%(message)s"
     file_format = "%(asctime)s - %(levelname)s - %(message)s"

nextrec/basic/metrics.py

@@ -260,7 +260,7 @@ def compute_mrr_at_k(
         order = np.argsort(scores)[::-1]
         k_user = min(k, idx.size)
         topk = order[:k_user]
-        ranked_labels = labels[order]
+        ranked_labels = labels[topk]
         rr = 0.0
         for rank, lab in enumerate(ranked_labels[:k_user], start=1):
             if lab > 0:
@@ -612,6 +612,7 @@ def evaluate_metrics(
                     if task_type in ["binary", "multilabel"]:
                         should_compute = metric_lower in {
                             "auc",
+                            "gauc",
                             "ks",
                             "logloss",
                             "accuracy",

nextrec/basic/model.py

@@ -455,7 +455,7 @@ class BaseModel(FeatureSet, nn.Module):
         if hasattr(
             self, "prediction_layer"
         ):  # we need to use registered task_slices for multi-task and multi-class
-            slices = self.prediction_layer._task_slices  # type: ignore
+            slices = self.prediction_layer.task_slices  # type: ignore
         else:
             slices = [(i, i + 1) for i in range(self.nums_task)]
         task_losses = []
@@ -1369,7 +1369,7 @@ class BaseModel(FeatureSet, nn.Module):
         pred_columns: list[str] = []
         if self.target_columns:
             for name in self.target_columns[:num_outputs]:
-                pred_columns.append(f"{name}_pred")
+                pred_columns.append(f"{name}")
         while len(pred_columns) < num_outputs:
             pred_columns.append(f"pred_{len(pred_columns)}")
         if include_ids and predict_id_columns:
@@ -1496,7 +1496,7 @@ class BaseModel(FeatureSet, nn.Module):
                     pred_columns = []
                     if self.target_columns:
                         for name in self.target_columns[:num_outputs]:
-                            pred_columns.append(f"{name}_pred")
+                            pred_columns.append(f"{name}")
                     while len(pred_columns) < num_outputs:
                         pred_columns.append(f"pred_{len(pred_columns)}")
 

nextrec/cli.py

@@ -8,10 +8,10 @@ following script to execute the desired operations.
 
 Examples:
     # Train a model
-    nextrec --mode=train --train_config=tutorials/iflytek/scripts/masknet/train_config.yaml
+    nextrec --mode=train --train_config=nextrec_cli_preset/train_config.yaml
 
     # Run prediction
-    nextrec --mode=predict --predict_config=tutorials/iflytek/scripts/masknet/predict_config.yaml
+    nextrec --mode=predict --predict_config=nextrec_cli_preset/predict_config.yaml
 
 Date: create on 06/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
@@ -33,7 +33,6 @@ from nextrec.data.preprocessor import DataProcessor
 from nextrec.utils.config import (
     build_feature_objects,
     build_model_instance,
-    extract_feature_groups,
     register_processor_features,
     resolve_path,
     select_features,
@@ -115,16 +114,13 @@ def train_model(train_config_path: str) -> None:
         df = read_table(data_path, data_cfg.get("format"))
         df_columns = list(df.columns)
 
-    # for some models have independent feature groups, we need to extract them here
-    feature_groups, grouped_columns = extract_feature_groups(feature_cfg, df_columns)
-    if feature_groups:
-        model_cfg.setdefault("params", {})
-        model_cfg["params"].setdefault("feature_groups", feature_groups)
-
     dense_names, sparse_names, sequence_names = select_features(feature_cfg, df_columns)
-    used_columns = (
-        dense_names + sparse_names + sequence_names + grouped_columns + target
-    )
+
+    # Extract id_column from data config for GAUC metrics
+    id_column = data_cfg.get("id_column") or data_cfg.get("user_id_column")
+    id_columns = [id_column] if id_column else []
+
+    used_columns = dense_names + sparse_names + sequence_names + target + id_columns
 
     # keep order but drop duplicates
     seen = set()
@@ -183,17 +179,16 @@ def train_model(train_config_path: str) -> None:
             streaming_valid_files = file_paths[-val_count:]
             streaming_train_files = file_paths[:-val_count]
             logger.info(
-                "使用 valid_ratio=%.3f 切分文件: 训练 %d 个文件, 验证 %d 个文件",
-                ratio,
-                len(streaming_train_files),
-                len(streaming_valid_files),
+                f"Split files for streaming training and validation using valid_ratio={ratio:.3f}: training {len(streaming_train_files)} files, validation {len(streaming_valid_files)} files"
             )
     train_data: Dict[str, Any]
     valid_data: Dict[str, Any] | None
 
     if val_data_path and not streaming:
         # Use specified validation dataset path
-        logger.info("使用指定的验证集路径: %s", val_data_path)
+        logger.info(
+            f"Validation using specified validation dataset path: {val_data_path}"
+        )
         val_data_resolved = resolve_path(val_data_path, config_dir)
         val_df = read_table(val_data_resolved, data_cfg.get("format"))
         val_df = val_df[unique_used_columns]
@@ -206,17 +201,21 @@ def train_model(train_config_path: str) -> None:
         valid_data = valid_data_result
         train_size = len(list(train_data.values())[0])
         valid_size = len(list(valid_data.values())[0])
-        logger.info("训练集样本数: %s, 验证集样本数: %s", train_size, valid_size)
+        logger.info(
+            f"Sample count - Training set: {train_size}, Validation set: {valid_size}"
+        )
     elif streaming:
         train_data = None  # type: ignore[assignment]
         valid_data = None
         if not val_data_path and not streaming_valid_files:
             logger.info(
-                "流式训练模式，未指定验证集路径且未配置 valid_ratio，跳过验证集创建"
+                "Streaming training mode: No validation dataset path specified and valid_ratio not configured, skipping validation dataset creation"
             )
     else:
         # Split data using valid_ratio
-        logger.info("使用 valid_ratio 切分数据: %s", data_cfg.get("valid_ratio", 0.2))
+        logger.info(
+            f"Splitting data using valid_ratio: {data_cfg.get('valid_ratio', 0.2)}"
+        )
         if not isinstance(processed, dict):
             raise TypeError("Processed data must be a dictionary for splitting")
         train_data, valid_data = split_dict_random(
@@ -230,6 +229,7 @@ def train_model(train_config_path: str) -> None:
         sparse_features=sparse_features,
         sequence_features=sequence_features,
         target=target,
+        id_columns=id_columns,
         processor=processor if streaming else None,
     )
     if streaming:
@@ -240,6 +240,7 @@ def train_model(train_config_path: str) -> None:
             shuffle=dataloader_cfg.get("train_shuffle", True),
             load_full=False,
             chunk_size=dataloader_chunk_size,
+            num_workers=dataloader_cfg.get("num_workers", 0),
         )
         valid_loader = None
         if val_data_path:
@@ -250,6 +251,7 @@ def train_model(train_config_path: str) -> None:
                 shuffle=dataloader_cfg.get("valid_shuffle", False),
                 load_full=False,
                 chunk_size=dataloader_chunk_size,
+                num_workers=dataloader_cfg.get("num_workers", 0),
             )
         elif streaming_valid_files:
             valid_loader = dataloader.create_dataloader(
@@ -258,17 +260,20 @@ def train_model(train_config_path: str) -> None:
                 shuffle=dataloader_cfg.get("valid_shuffle", False),
                 load_full=False,
                 chunk_size=dataloader_chunk_size,
+                num_workers=dataloader_cfg.get("num_workers", 0),
             )
     else:
         train_loader = dataloader.create_dataloader(
             data=train_data,
             batch_size=dataloader_cfg.get("train_batch_size", 512),
             shuffle=dataloader_cfg.get("train_shuffle", True),
+            num_workers=dataloader_cfg.get("num_workers", 0),
         )
         valid_loader = dataloader.create_dataloader(
             data=valid_data,
             batch_size=dataloader_cfg.get("valid_batch_size", 512),
             shuffle=dataloader_cfg.get("valid_shuffle", False),
+            num_workers=dataloader_cfg.get("num_workers", 0),
         )
 
     model_cfg.setdefault("session_id", session_id)
@@ -300,6 +305,9 @@ def train_model(train_config_path: str) -> None:
             "batch_size", dataloader_cfg.get("train_batch_size", 512)
         ),
         shuffle=train_cfg.get("shuffle", True),
+        num_workers=dataloader_cfg.get("num_workers", 0),
+        user_id_column=id_column,
+        tensorboard=False,
     )
 
 
@@ -325,19 +333,15 @@ def predict_model(predict_config_path: str) -> None:
     model_cfg_path = resolve_path(
         cfg.get("model_config", "model_config.yaml"), config_dir
     )
-    feature_cfg_path = resolve_path(
-        cfg.get("feature_config", "feature_config.yaml"), config_dir
-    )
+    # feature_cfg_path = resolve_path(
+    #     cfg.get("feature_config", "feature_config.yaml"), config_dir
+    # )
 
     model_cfg = read_yaml(model_cfg_path)
-    feature_cfg = read_yaml(feature_cfg_path)
+    # feature_cfg = read_yaml(feature_cfg_path)
     model_cfg.setdefault("session_id", session_id)
-    feature_groups_raw = feature_cfg.get("feature_groups") or {}
     model_cfg.setdefault("params", {})
 
-    # attach feature_groups in predict phase to avoid missing bindings
-    model_cfg["params"]["feature_groups"] = feature_groups_raw
-
     processor = DataProcessor.load(processor_path)
 
     # Load checkpoint and ensure required parameters are passed
@@ -383,13 +387,6 @@ def predict_model(predict_config_path: str) -> None:
     if target_override:
         target_cols = normalize_to_list(target_override)
 
-    # Recompute feature_groups with available feature names to drive bindings
-    feature_group_names = [f.name for f in all_features if hasattr(f, "name")]
-    parsed_feature_groups, _ = extract_feature_groups(feature_cfg, feature_group_names)
-    if parsed_feature_groups:
-        model_cfg.setdefault("params", {})
-        model_cfg["params"]["feature_groups"] = parsed_feature_groups
-
     model = build_model_instance(
         model_cfg=model_cfg,
         model_cfg_path=model_cfg_path,
@@ -440,6 +437,7 @@ def predict_model(predict_config_path: str) -> None:
         return_dataframe=False,
         save_path=output_path,
         save_format=predict_cfg.get("save_format", "csv"),
+        num_workers=predict_cfg.get("num_workers", 0),
     )
     duration = time.time() - start
     logger.info(f"Prediction completed, results saved to: {output_path}")
@@ -448,7 +446,7 @@ def predict_model(predict_config_path: str) -> None:
     preview_rows = predict_cfg.get("preview_rows", 0)
     if preview_rows > 0:
         try:
-            preview = pd.read_csv(output_path, nrows=preview_rows)
+            preview = pd.read_csv(output_path, nrows=preview_rows, low_memory=False)
             logger.info(f"Output preview:\n{preview}")
         except Exception as exc:  # pragma: no cover
             logger.warning(f"Failed to read output preview: {exc}")
@@ -472,25 +470,21 @@ Examples:
         "--mode",
         choices=["train", "predict"],
         required=True,
-        help="运行模式：train 或 predict",
-    )
-    parser.add_argument("--train_config", help="训练配置文件路径")
-    parser.add_argument("--predict_config", help="预测配置文件路径")
-    parser.add_argument(
-        "--config",
-        help="通用配置文件路径（已废弃，建议使用 --train_config 或 --predict_config）",
+        help="Running mode: train or predict",
     )
+    parser.add_argument("--train_config", help="Training configuration file path")
+    parser.add_argument("--predict_config", help="Prediction configuration file path")
     args = parser.parse_args()
 
     if args.mode == "train":
-        config_path = args.train_config or args.config
+        config_path = args.train_config
         if not config_path:
-            parser.error("train 模式需要提供 --train_config")
+            parser.error("[NextRec CLI Error] train mode requires --train_config")
         train_model(config_path)
     else:
-        config_path = args.predict_config or args.config
+        config_path = args.predict_config
         if not config_path:
-            parser.error("predict 模式需要提供 --predict_config")
+            parser.error("[NextRec CLI Error] predict mode requires --predict_config")
         predict_model(config_path)
 
 

nextrec/data/dataloader.py

@@ -322,7 +322,7 @@ class RecDataLoader(FeatureSet):
                 except OSError:
                     pass
                 try:
-                    df = read_table(file_path, file_type=file_type)
+                    df = read_table(file_path, data_format=file_type)
                     dfs.append(df)
                 except MemoryError as exc:
                     raise MemoryError(

nextrec/models/multi_task/esmm.py

@@ -76,10 +76,10 @@ class ESMM(BaseModel):
         sequence_features: list[SequenceFeature],
         ctr_params: dict,
         cvr_params: dict,
-        target: list[str] = ["ctr", "ctcvr"],  # Note: ctcvr = ctr * cvr
+        target: list[str] | None = None,  # Note: ctcvr = ctr * cvr
         task: list[str] | None = None,
         optimizer: str = "adam",
-        optimizer_params: dict = {},
+        optimizer_params: dict | None = None,
         loss: str | nn.Module | list[str | nn.Module] | None = "bce",
         loss_params: dict | list[dict] | None = None,
         device: str = "cpu",
@@ -90,19 +90,36 @@ class ESMM(BaseModel):
         **kwargs,
     ):
 
-        # ESMM requires exactly 2 targets: ctr and ctcvr
+        target = target or ["ctr", "ctcvr"]
+        optimizer_params = optimizer_params or {}
+        if loss is None:
+            loss = "bce"
+
         if len(target) != 2:
             raise ValueError(
                 f"ESMM requires exactly 2 targets (ctr and ctcvr), got {len(target)}"
             )
 
+        self.num_tasks = 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:
+            raise ValueError(
+                f"Length of task ({len(resolved_task)}) must match number of targets ({self.num_tasks})."
+            )
+        # resolved_task is now guaranteed to be a list[str]
+
         super(ESMM, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=task
-            or self.default_task,  # Both CTR and CTCVR are binary classification
+            task=resolved_task,  # Both CTR and CTCVR are binary classification
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
@@ -112,19 +129,9 @@ class ESMM(BaseModel):
         )
 
         self.loss = loss
-        if self.loss is None:
-            self.loss = "bce"
 
-        # All features
-        self.all_features = dense_features + sparse_features + sequence_features
-        # Shared embedding layer
         self.embedding = EmbeddingLayer(features=self.all_features)
-        input_dim = (
-            self.embedding.input_dim
-        )  # Calculate input dimension, better way than below
-        # 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
+        input_dim = self.embedding.input_dim
 
         # CTR tower
         self.ctr_tower = MLP(input_dim=input_dim, output_layer=True, **ctr_params)

nextrec/models/multi_task/mmoe.py

@@ -73,16 +73,16 @@ class MMOE(BaseModel):
 
     def __init__(
         self,
-        dense_features: list[DenseFeature] = [],
-        sparse_features: list[SparseFeature] = [],
-        sequence_features: list[SequenceFeature] = [],
-        expert_params: dict = {},
+        dense_features: list[DenseFeature] | None = None,
+        sparse_features: list[SparseFeature] | None = None,
+        sequence_features: list[SequenceFeature] | None = None,
+        expert_params: dict | None = None,
         num_experts: int = 3,
-        tower_params_list: list[dict] = [],
-        target: list[str] = [],
-        task: str | list[str] | None = None,
+        tower_params_list: list[dict] | None = None,
+        target: list[str] | str | None = None,
+        task: str | list[str] = "binary",
         optimizer: str = "adam",
-        optimizer_params: dict = {},
+        optimizer_params: dict | None = None,
         loss: str | nn.Module | list[str | nn.Module] | None = "bce",
         loss_params: dict | list[dict] | None = None,
         device: str = "cpu",
@@ -93,14 +93,39 @@ class MMOE(BaseModel):
         **kwargs,
     ):
 
-        self.num_tasks = len(target)
+        dense_features = dense_features or []
+        sparse_features = sparse_features or []
+        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):
+            target = [target]
+
+        self.num_tasks = 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:
+            raise ValueError(
+                f"Length of task ({len(resolved_task)}) must match number of targets ({self.num_tasks})."
+            )
 
         super(MMOE, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=task or self.default_task,
+            task=resolved_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
@@ -110,8 +135,6 @@ class MMOE(BaseModel):
         )
 
         self.loss = loss
-        if self.loss is None:
-            self.loss = "bce"
 
         # Number of tasks and experts
         self.num_tasks = len(target)
@@ -122,12 +145,8 @@ class MMOE(BaseModel):
                 f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.num_tasks})"
             )
 
-        self.all_features = dense_features + sparse_features + sequence_features
         self.embedding = EmbeddingLayer(features=self.all_features)
         input_dim = self.embedding.input_dim
-        # 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
 
         # Expert networks (shared by all tasks)
         self.experts = nn.ModuleList()
@@ -162,7 +181,7 @@ class MMOE(BaseModel):
         self.compile(
             optimizer=optimizer,
             optimizer_params=optimizer_params,
-            loss=loss,
+            loss=self.loss,
             loss_params=loss_params,
         )
 

nextrec/models/multi_task/ple.py

@@ -51,6 +51,7 @@ import torch.nn as nn
 from nextrec.basic.model import BaseModel
 from nextrec.basic.layers import EmbeddingLayer, MLP, PredictionLayer
 from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
+from nextrec.utils.model import get_mlp_output_dim
 
 
 class CGCLayer(nn.Module):
@@ -72,13 +73,13 @@ class CGCLayer(nn.Module):
         if num_tasks < 1:
             raise ValueError("num_tasks must be >= 1")
 
-        specific_params_list = self._normalize_specific_params(
+        specific_params_list = self.normalize_specific_params(
             specific_expert_params, num_tasks
         )
 
-        self.output_dim = self._get_output_dim(shared_expert_params, input_dim)
+        self.output_dim = get_mlp_output_dim(shared_expert_params, input_dim)
         specific_dims = [
-            self._get_output_dim(params, input_dim) for params in specific_params_list
+            get_mlp_output_dim(params, input_dim) for params in specific_params_list
         ]
         dims_set = set(specific_dims + [self.output_dim])
         if len(dims_set) != 1:
@@ -165,14 +166,7 @@ class CGCLayer(nn.Module):
         return new_task_fea, new_shared
 
     @staticmethod
-    def _get_output_dim(params: dict, fallback: int) -> int:
-        dims = params.get("dims")
-        if dims:
-            return dims[-1]
-        return fallback
-
-    @staticmethod
-    def _normalize_specific_params(
+    def normalize_specific_params(
         params: dict | list[dict], num_tasks: int
     ) -> list[dict]:
         if isinstance(params, list):
@@ -232,12 +226,24 @@ class PLE(BaseModel):
 
         self.num_tasks = 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:
+            raise ValueError(
+                f"Length of task ({len(resolved_task)}) must match number of targets ({self.num_tasks})."
+            )
+
         super(PLE, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
             sequence_features=sequence_features,
             target=target,
-            task=task or self.default_task,
+            task=resolved_task,
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,