nextrec 0.4.11__py3-none-any.whl → 0.4.13__py3-none-any.whl

nextrec/basic/model.py

@@ -2,7 +2,7 @@
 Base Model & Base Match Model Class
 
 Date: create on 27/10/2025
-Checkpoint: edit on 19/12/2025
+Checkpoint: edit on 20/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 """
 
@@ -50,12 +50,14 @@ from nextrec.data.dataloader import (
 )
 from nextrec.loss import (
     BPRLoss,
+    GradNormLossWeighting,
     HingeLoss,
     InfoNCELoss,
     SampledSoftmaxLoss,
     TripletLoss,
     get_loss_fn,
 )
+from nextrec.loss.grad_norm import get_grad_norm_shared_params
 from nextrec.utils.console import display_metrics_table, progress
 from nextrec.utils.torch_utils import (
     add_distributed_sampler,
@@ -169,6 +171,7 @@ class BaseModel(FeatureSet, nn.Module):
         self.loss_weight = None
 
         self.early_stop_patience = early_stop_patience
+        # max samples to keep for training metrics, in case of large training set
         self.max_metrics_samples = (
             None if max_metrics_samples is None else int(max_metrics_samples)
         )
@@ -176,6 +179,8 @@ class BaseModel(FeatureSet, nn.Module):
         self.logger_initialized = False
         self.training_logger = None
         self.callbacks = CallbackList(callbacks) if callbacks else CallbackList()
+        self.grad_norm: GradNormLossWeighting | None = None
+        self.grad_norm_shared_params: list[torch.nn.Parameter] | None = None
 
     def register_regularization_weights(
         self,
@@ -376,7 +381,7 @@ class BaseModel(FeatureSet, nn.Module):
         scheduler_params: dict | None = None,
         loss: str | nn.Module | list[str | nn.Module] | None = "bce",
         loss_params: dict | list[dict] | None = None,
-        loss_weights: int | float | list[int | float] | None = None,
+        loss_weights: int | float | list[int | float] | dict | str | None = None,
         callbacks: list[Callback] | None = None,
     ):
         """
@@ -389,6 +394,7 @@ class BaseModel(FeatureSet, nn.Module):
             loss: Loss function name, instance, or list for multi-task. e.g., 'bce', 'mse', or torch.nn.BCELoss(), you can also use custom loss functions.
             loss_params: Loss function parameters, or list for multi-task. e.g., {'weight': tensor([0.25, 0.75])}.
             loss_weights: Weights for each task loss, int/float for single-task or list for multi-task. e.g., 1.0, or [1.0, 0.5].
+                Use "grad_norm" or {"method": "grad_norm", ...} to enable GradNorm for multi-task loss balancing.
             callbacks: Additional callbacks to add to the existing callback list. e.g., [EarlyStopper(), CheckpointSaver()].
         """
         if loss_params is None:
@@ -442,7 +448,31 @@ class BaseModel(FeatureSet, nn.Module):
             for i in range(self.nums_task)
         ]
 
-        if loss_weights is None:
+        self.grad_norm = None
+        self.grad_norm_shared_params = None
+        if isinstance(loss_weights, str) and loss_weights.lower() == "grad_norm":
+            if self.nums_task == 1:
+                raise ValueError(
+                    "[BaseModel-compile Error] GradNorm requires multi-task setup."
+                )
+            self.grad_norm = GradNormLossWeighting(
+                num_tasks=self.nums_task, device=self.device
+            )
+            self.loss_weights = None
+        elif (
+            isinstance(loss_weights, dict) and loss_weights.get("method") == "grad_norm"
+        ):
+            if self.nums_task == 1:
+                raise ValueError(
+                    "[BaseModel-compile Error] GradNorm requires multi-task setup."
+                )
+            grad_norm_params = dict(loss_weights)
+            grad_norm_params.pop("method", None)
+            self.grad_norm = GradNormLossWeighting(
+                num_tasks=self.nums_task, device=self.device, **grad_norm_params
+            )
+            self.loss_weights = None
+        elif loss_weights is None:
             self.loss_weights = None
         elif self.nums_task == 1:
             if isinstance(loss_weights, (list, tuple)):
@@ -507,9 +537,20 @@ class BaseModel(FeatureSet, nn.Module):
             y_pred_i = y_pred[:, start:end]
             y_true_i = y_true[:, start:end]
             task_loss = self.loss_fn[i](y_pred_i, y_true_i)
-            if isinstance(self.loss_weights, (list, tuple)):
-                task_loss *= self.loss_weights[i]
             task_losses.append(task_loss)
+        if self.grad_norm is not None:
+            if self.grad_norm_shared_params is None:
+                self.grad_norm_shared_params = get_grad_norm_shared_params(
+                    self, getattr(self, "grad_norm_shared_modules", None)
+                )
+            return self.grad_norm.compute_weighted_loss(
+                task_losses, self.grad_norm_shared_params
+            )
+        if isinstance(self.loss_weights, (list, tuple)):
+            task_losses = [
+                task_loss * self.loss_weights[i]
+                for i, task_loss in enumerate(task_losses)
+            ]
         return torch.stack(task_losses).sum()
 
     def prepare_data_loader(
@@ -563,6 +604,7 @@ class BaseModel(FeatureSet, nn.Module):
         num_workers: int = 0,
         tensorboard: bool = True,
         auto_distributed_sampler: bool = True,
+        log_interval: int = 1,
     ):
         """
         Train the model.
@@ -579,6 +621,7 @@ class BaseModel(FeatureSet, nn.Module):
             num_workers: DataLoader worker count.
             tensorboard: Enable tensorboard logging.
             auto_distributed_sampler: Attach DistributedSampler automatically when distributed, set False to when data is already sharded per rank.
+            log_interval: Log validation metrics every N epochs (still computes metrics each epoch).
 
         Notes:
             - Distributed training uses DDP; init occurs via env vars (RANK/WORLD_SIZE/LOCAL_RANK).
@@ -630,6 +673,9 @@ class BaseModel(FeatureSet, nn.Module):
             )
         )  # ['auc', 'logloss'], {'target1': ['auc', 'logloss'], 'target2': ['mse']}, 'max'
 
+        if log_interval < 1:
+            raise ValueError("[BaseModel-fit Error] log_interval must be >= 1.")
+
         # Setup default callbacks if missing
         if self.nums_task == 1:
             monitor_metric = f"val_{self.metrics[0]}"
@@ -911,23 +957,27 @@ class BaseModel(FeatureSet, nn.Module):
                     user_ids=valid_user_ids if self.needs_user_ids else None,
                     num_workers=num_workers,
                 )
-                display_metrics_table(
-                    epoch=epoch + 1,
-                    epochs=epochs,
-                    split="Valid",
-                    loss=None,
-                    metrics=val_metrics,
-                    target_names=self.target_columns,
-                    base_metrics=(
-                        self.metrics
-                        if isinstance(getattr(self, "metrics", None), list)
-                        else None
-                    ),
-                    is_main_process=self.is_main_process,
-                    colorize=lambda s: colorize("  " + s, color="cyan"),
-                )
+                should_log_valid = (epoch + 1) % log_interval == 0 or (
+                    epoch + 1
+                ) == epochs
+                if should_log_valid:
+                    display_metrics_table(
+                        epoch=epoch + 1,
+                        epochs=epochs,
+                        split="Valid",
+                        loss=None,
+                        metrics=val_metrics,
+                        target_names=self.target_columns,
+                        base_metrics=(
+                            self.metrics
+                            if isinstance(getattr(self, "metrics", None), list)
+                            else None
+                        ),
+                        is_main_process=self.is_main_process,
+                        colorize=lambda s: colorize("  " + s, color="cyan"),
+                    )
                 self.callbacks.on_validation_end()
-                if val_metrics and self.training_logger:
+                if should_log_valid and val_metrics and self.training_logger:
                     self.training_logger.log_metrics(
                         val_metrics, step=epoch + 1, split="valid"
                     )
@@ -1043,6 +1093,8 @@ class BaseModel(FeatureSet, nn.Module):
             params = model.parameters() if self.ddp_model is not None else self.parameters()  # type: ignore # ddp model parameters or self parameters
             nn.utils.clip_grad_norm_(params, self.max_gradient_norm)
             self.optimizer_fn.step()
+            if self.grad_norm is not None:
+                self.grad_norm.step()
             accumulated_loss += loss.item()
 
             if (
@@ -1207,7 +1259,7 @@ class BaseModel(FeatureSet, nn.Module):
             user_id_column: Column name for user IDs if user_ids is not provided. e.g. 'user_id'
             num_workers: DataLoader worker count.
         """
-        model = self.ddp_model if getattr(self, "ddp_model", None) is not None else self
+        model = self.ddp_model if self.ddp_model is not None else self
         model.eval()
         eval_metrics = metrics if metrics is not None else self.metrics
         if eval_metrics is None:
@@ -1233,6 +1285,10 @@ class BaseModel(FeatureSet, nn.Module):
                 batch_count += 1
                 batch_dict = batch_to_dict(batch_data)
                 X_input, y_true = self.get_input(batch_dict, require_labels=True)
+                if X_input is None:
+                    raise ValueError(
+                        "[BaseModel-evaluate Error] No input features found in the evaluation data."
+                    )
                 y_pred = model(X_input)
                 if y_true is not None:
                     y_true_list.append(y_true.cpu().numpy())
@@ -1322,7 +1378,7 @@ class BaseModel(FeatureSet, nn.Module):
         return_dataframe: bool = True,
         streaming_chunk_size: int = 10000,
         num_workers: int = 0,
-    ) -> pd.DataFrame | np.ndarray:
+    ) -> pd.DataFrame | np.ndarray | Path | None:
         """
         Note: predict does not support distributed mode currently, consider it as a single-process operation.
         Make predictions on the given data.
@@ -1497,7 +1553,7 @@ class BaseModel(FeatureSet, nn.Module):
         streaming_chunk_size: int,
         return_dataframe: bool,
         id_columns: list[str] | None = None,
-    ) -> pd.DataFrame:
+    ) -> pd.DataFrame | Path:
         if isinstance(data, (str, os.PathLike)):
             rec_loader = RecDataLoader(
                 dense_features=self.dense_features,
@@ -1624,11 +1680,11 @@ class BaseModel(FeatureSet, nn.Module):
         )
         model_path = Path(target_path)
 
-        model_to_save = (
-            self.ddp_model.module
-            if getattr(self, "ddp_model", None) is not None
-            else self
-        )
+        ddp_model = getattr(self, "ddp_model", None)
+        if ddp_model is not None:
+            model_to_save = ddp_model.module
+        else:
+            model_to_save = self
         torch.save(model_to_save.state_dict(), model_path)
         # torch.save(self.state_dict(), model_path)
 
@@ -2025,33 +2081,18 @@ class BaseMatchModel(BaseModel):
         self.num_negative_samples = num_negative_samples
         self.temperature = temperature
         self.similarity_metric = similarity_metric
-
-        self.user_feature_names = [
-            f.name
-            for f in (
-                self.user_dense_features
-                + self.user_sparse_features
-                + self.user_sequence_features
-            )
-        ]
-        self.item_feature_names = [
-            f.name
-            for f in (
-                self.item_dense_features
-                + self.item_sparse_features
-                + self.item_sequence_features
-            )
-        ]
-
-    def get_user_features(self, X_input: dict) -> dict:
-        return {
-            name: X_input[name] for name in self.user_feature_names if name in X_input
-        }
-
-    def get_item_features(self, X_input: dict) -> dict:
-        return {
-            name: X_input[name] for name in self.item_feature_names if name in X_input
-        }
+        self.user_features_all = (
+            self.user_dense_features
+            + self.user_sparse_features
+            + self.user_sequence_features
+        )
+        self.item_features_all = (
+            self.item_dense_features
+            + self.item_sparse_features
+            + self.item_sequence_features
+        )
+        self.user_feature_names = {feature.name for feature in self.user_features_all}
+        self.item_feature_names = {feature.name for feature in self.item_features_all}
 
     def compile(
         self,
@@ -2068,13 +2109,11 @@ class BaseMatchModel(BaseModel):
         scheduler_params: dict | None = None,
         loss: str | nn.Module | list[str | nn.Module] | None = "bce",
         loss_params: dict | list[dict] | None = None,
-        loss_weights: int | float | list[int | float] | None = None,
+        loss_weights: int | float | list[int | float] | dict | str | None = None,
         callbacks: list[Callback] | None = None,
     ):
         """
         Configure the match model for training.
-
-        This mirrors `BaseModel.compile()` and additionally validates `training_mode`.
         """
         if self.training_mode not in self.support_training_modes:
             raise ValueError(
@@ -2090,7 +2129,7 @@ class BaseMatchModel(BaseModel):
         effective_loss: str | nn.Module | list[str | nn.Module] | None = loss
         if effective_loss is None:
             effective_loss = default_loss_by_mode[self.training_mode]
-        elif isinstance(effective_loss, (str,)):
+        elif isinstance(effective_loss, str):
             if self.training_mode in {"pairwise", "listwise"} and effective_loss in {
                 "bce",
                 "binary_crossentropy",
@@ -2124,6 +2163,7 @@ class BaseMatchModel(BaseModel):
     def inbatch_logits(
         self, user_emb: torch.Tensor, item_emb: torch.Tensor
     ) -> torch.Tensor:
+        """Compute in-batch logits matrix between user and item embeddings."""
         if self.similarity_metric == "dot":
             logits = torch.matmul(user_emb, item_emb.t())
         elif self.similarity_metric == "cosine":
@@ -2131,8 +2171,8 @@ class BaseMatchModel(BaseModel):
             item_norm = F.normalize(item_emb, p=2, dim=-1)
             logits = torch.matmul(user_norm, item_norm.t())
         elif self.similarity_metric == "euclidean":
-            user_sq = (user_emb**2).sum(dim=1, keepdim=True)  # [B, 1]
-            item_sq = (item_emb**2).sum(dim=1, keepdim=True).t()  # [1, B]
+            user_sq = torch.sum(user_emb**2, dim=1, keepdim=True)  # [B, 1]
+            item_sq = torch.sum(item_emb**2, dim=1, keepdim=True).t()  # [1, B]
             logits = -(user_sq + item_sq - 2.0 * torch.matmul(user_emb, item_emb.t()))
         else:
             raise ValueError(f"Unknown similarity metric: {self.similarity_metric}")
@@ -2141,56 +2181,43 @@ class BaseMatchModel(BaseModel):
     def compute_similarity(
         self, user_emb: torch.Tensor, item_emb: torch.Tensor
     ) -> torch.Tensor:
-        if self.similarity_metric == "dot":
-            if user_emb.dim() == 3 and item_emb.dim() == 3:
-                # [batch_size, num_items, emb_dim] @ [batch_size, num_items, emb_dim]
-                similarity = torch.sum(
-                    user_emb * item_emb, dim=-1
-                )  # [batch_size, num_items]
-            elif user_emb.dim() == 2 and item_emb.dim() == 3:
-                # [batch_size, emb_dim] @ [batch_size, num_items, emb_dim]
-                user_emb_expanded = user_emb.unsqueeze(1)  # [batch_size, 1, emb_dim]
-                similarity = torch.sum(
-                    user_emb_expanded * item_emb, dim=-1
-                )  # [batch_size, num_items]
-            else:
-                similarity = torch.sum(user_emb * item_emb, dim=-1)  # [batch_size]
+        """Compute similarity score between user and item embeddings."""
+        if user_emb.dim() == 2 and item_emb.dim() == 3:
+            user_emb = user_emb.unsqueeze(1)
 
+        if self.similarity_metric == "dot":
+            similarity = torch.sum(user_emb * item_emb, dim=-1)
         elif self.similarity_metric == "cosine":
-            if user_emb.dim() == 3 and item_emb.dim() == 3:
-                similarity = F.cosine_similarity(user_emb, item_emb, dim=-1)
-            elif user_emb.dim() == 2 and item_emb.dim() == 3:
-                user_emb_expanded = user_emb.unsqueeze(1)
-                similarity = F.cosine_similarity(user_emb_expanded, item_emb, dim=-1)
-            else:
-                similarity = F.cosine_similarity(user_emb, item_emb, dim=-1)
-
+            similarity = F.cosine_similarity(user_emb, item_emb, dim=-1)
         elif self.similarity_metric == "euclidean":
-            if user_emb.dim() == 3 and item_emb.dim() == 3:
-                distance = torch.sum((user_emb - item_emb) ** 2, dim=-1)
-            elif user_emb.dim() == 2 and item_emb.dim() == 3:
-                user_emb_expanded = user_emb.unsqueeze(1)
-                distance = torch.sum((user_emb_expanded - item_emb) ** 2, dim=-1)
-            else:
-                distance = torch.sum((user_emb - item_emb) ** 2, dim=-1)
-            similarity = -distance
-
+            similarity = -torch.sum((user_emb - item_emb) ** 2, dim=-1)
         else:
             raise ValueError(f"Unknown similarity metric: {self.similarity_metric}")
         similarity = similarity / self.temperature
         return similarity
 
     def user_tower(self, user_input: dict) -> torch.Tensor:
+        """User tower to encode user features into embeddings."""
         raise NotImplementedError
 
     def item_tower(self, item_input: dict) -> torch.Tensor:
+        """Item tower to encode item features into embeddings."""
         raise NotImplementedError
 
     def forward(
         self, X_input: dict
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        user_input = self.get_user_features(X_input)
-        item_input = self.get_item_features(X_input)
+        """Rewrite forward to handle user and item features separately."""
+        user_input = {
+            name: tensor
+            for name, tensor in X_input.items()
+            if name in self.user_feature_names
+        }
+        item_input = {
+            name: tensor
+            for name, tensor in X_input.items()
+            if name in self.item_feature_names
+        }
 
         user_emb = self.user_tower(user_input)  # [B, D]
         item_emb = self.item_tower(item_input)  # [B, D]
@@ -2254,11 +2281,35 @@ class BaseMatchModel(BaseModel):
             raise ValueError(f"Unknown training mode: {self.training_mode}")
 
     def prepare_feature_data(
-        self, data: dict | pd.DataFrame | DataLoader, features: list, batch_size: int
+        self,
+        data,
+        features: list,
+        batch_size: int,
+        num_workers: int = 0,
+        streaming_chunk_size: int = 10000,
     ) -> DataLoader:
         """Prepare data loader for specific features."""
         if isinstance(data, DataLoader):
             return data
+        if isinstance(data, (str, os.PathLike)):
+            dense_features = [f for f in features if isinstance(f, DenseFeature)]
+            sparse_features = [f for f in features if isinstance(f, SparseFeature)]
+            sequence_features = [f for f in features if isinstance(f, SequenceFeature)]
+            rec_loader = RecDataLoader(
+                dense_features=dense_features,
+                sparse_features=sparse_features,
+                sequence_features=sequence_features,
+                target=[],
+                id_columns=[],
+            )
+            return rec_loader.create_dataloader(
+                data=data,
+                batch_size=batch_size,
+                shuffle=False,
+                streaming=True,
+                chunk_size=streaming_chunk_size,
+                num_workers=num_workers,
+            )
         tensors = build_tensors_from_data(
             data=data,
             raw_data=data,
@@ -2276,44 +2327,91 @@ class BaseMatchModel(BaseModel):
             batch_size=batch_size,
             shuffle=False,
             collate_fn=collate_fn,
+            num_workers=num_workers,
         )
 
+    def build_feature_tensors(self, feature_source: dict, features: list) -> dict:
+        """Convert feature values to tensors on the model device."""
+        tensors = {}
+        for feature in features:
+            if feature.name not in feature_source:
+                raise KeyError(
+                    f"[BaseMatchModel-feature Error] Feature '{feature.name}' not found in input data."
+                )
+            feature_data = get_column_data(feature_source, feature.name)
+            tensors[feature.name] = to_tensor(
+                feature_data,
+                dtype=(
+                    torch.float32 if isinstance(feature, DenseFeature) else torch.long
+                ),
+                device=self.device,
+            )
+        return tensors
+
     def encode_user(
-        self, data: dict | pd.DataFrame | DataLoader, batch_size: int = 512
+        self,
+        data: (
+            dict
+            | pd.DataFrame
+            | DataLoader
+            | str
+            | os.PathLike
+            | list[str | os.PathLike]
+        ),
+        batch_size: int = 512,
+        num_workers: int = 0,
+        streaming_chunk_size: int = 10000,
     ) -> np.ndarray:
         self.eval()
-        all_user_features = (
-            self.user_dense_features
-            + self.user_sparse_features
-            + self.user_sequence_features
+        data_loader = self.prepare_feature_data(
+            data,
+            self.user_features_all,
+            batch_size,
+            num_workers=num_workers,
+            streaming_chunk_size=streaming_chunk_size,
         )
-        data_loader = self.prepare_feature_data(data, all_user_features, batch_size)
 
         embeddings_list = []
         with torch.no_grad():
             for batch_data in progress(data_loader, description="Encoding users"):
                 batch_dict = batch_to_dict(batch_data, include_ids=False)
-                user_input = self.get_user_features(batch_dict["features"])
+                user_input = self.build_feature_tensors(
+                    batch_dict["features"], self.user_features_all
+                )
                 user_emb = self.user_tower(user_input)
                 embeddings_list.append(user_emb.cpu().numpy())
         return np.concatenate(embeddings_list, axis=0)
 
     def encode_item(
-        self, data: dict | pd.DataFrame | DataLoader, batch_size: int = 512
+        self,
+        data: (
+            dict
+            | pd.DataFrame
+            | DataLoader
+            | str
+            | os.PathLike
+            | list[str | os.PathLike]
+        ),
+        batch_size: int = 512,
+        num_workers: int = 0,
+        streaming_chunk_size: int = 10000,
     ) -> np.ndarray:
         self.eval()
-        all_item_features = (
-            self.item_dense_features
-            + self.item_sparse_features
-            + self.item_sequence_features
+        data_loader = self.prepare_feature_data(
+            data,
+            self.item_features_all,
+            batch_size,
+            num_workers=num_workers,
+            streaming_chunk_size=streaming_chunk_size,
         )
-        data_loader = self.prepare_feature_data(data, all_item_features, batch_size)
 
         embeddings_list = []
         with torch.no_grad():
             for batch_data in progress(data_loader, description="Encoding items"):
                 batch_dict = batch_to_dict(batch_data, include_ids=False)
-                item_input = self.get_item_features(batch_dict["features"])
+                item_input = self.build_feature_tensors(
+                    batch_dict["features"], self.item_features_all
+                )
                 item_emb = self.item_tower(item_input)
                 embeddings_list.append(item_emb.cpu().numpy())
         return np.concatenate(embeddings_list, axis=0)

nextrec/cli.py

@@ -380,6 +380,7 @@ def train_model(train_config_path: str) -> None:
         optimizer_params=train_cfg.get("optimizer_params", {}),
         loss=train_cfg.get("loss", "focal"),
         loss_params=train_cfg.get("loss_params", {}),
+        loss_weights=train_cfg.get("loss_weights"),
     )
 
     model.fit(
@@ -416,7 +417,7 @@ def predict_model(predict_config_path: str) -> None:
         # Auto-infer session_id from checkpoint directory name
         session_cfg = cfg.get("session", {}) or {}
         session_id = session_cfg.get("id") or session_dir.name
-    
+
     setup_logger(session_id=session_id)
 
     log_cli_section("CLI")
@@ -436,7 +437,7 @@ def predict_model(predict_config_path: str) -> None:
         processor_path = session_dir / "processor" / "processor.pkl"
 
     predict_cfg = cfg.get("predict", {}) or {}
-    
+
     # Auto-find model_config in checkpoint directory if not specified
     if "model_config" in cfg:
         model_cfg_path = resolve_path(cfg["model_config"], config_dir)
@@ -563,7 +564,12 @@ def predict_model(predict_config_path: str) -> None:
     log_kv_lines(
         [
             ("Data path", data_path),
-            ("Format", predict_cfg.get("source_data_format", predict_cfg.get("data_format", "auto"))),
+            (
+                "Format",
+                predict_cfg.get(
+                    "source_data_format", predict_cfg.get("data_format", "auto")
+                ),
+            ),
             ("Batch size", batch_size),
             ("Chunk size", predict_cfg.get("chunk_size", 20000)),
             ("Streaming", predict_cfg.get("streaming", True)),
@@ -579,7 +585,9 @@ def predict_model(predict_config_path: str) -> None:
     )
 
     # Build output path: {checkpoint_path}/predictions/{name}.{save_data_format}
-    save_format = predict_cfg.get("save_data_format", predict_cfg.get("save_format", "csv"))
+    save_format = predict_cfg.get(
+        "save_data_format", predict_cfg.get("save_format", "csv")
+    )
     pred_name = predict_cfg.get("name", "pred")
     # Pass filename with extension to let model.predict handle path resolution
     save_path = f"{pred_name}.{save_format}"
@@ -597,7 +605,11 @@ def predict_model(predict_config_path: str) -> None:
     )
     duration = time.time() - start
     # When return_dataframe=False, result is the actual file path
-    output_path = result if isinstance(result, Path) else checkpoint_base / "predictions" / save_path
+    output_path = (
+        result
+        if isinstance(result, Path)
+        else checkpoint_base / "predictions" / save_path
+    )
     logger.info(f"Prediction completed, results saved to: {output_path}")
     logger.info(f"Total time: {duration:.2f} seconds")
 

nextrec/data/preprocessor.py

@@ -610,7 +610,7 @@ class DataProcessor(FeatureSet):
         save_format: Optional[Literal["csv", "parquet"]],
         output_path: Optional[str],
         warn_missing: bool = True,
-    ) -> Union[pd.DataFrame, Dict[str, np.ndarray]]:
+    ):
         logger = logging.getLogger()
         is_dataframe = isinstance(data, pd.DataFrame)
         data_dict = data if isinstance(data, dict) else None
@@ -705,7 +705,7 @@ class DataProcessor(FeatureSet):
         output_path: Optional[str],
         save_format: Optional[Literal["csv", "parquet"]],
         chunk_size: int = 200000,
-    ) -> list[str]:
+    ):
         """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.
@@ -852,7 +852,7 @@ class DataProcessor(FeatureSet):
         save_format: Optional[Literal["csv", "parquet"]] = None,
         output_path: Optional[str] = None,
         chunk_size: int = 200000,
-    ) -> Union[pd.DataFrame, Dict[str, np.ndarray], list[str]]:
+    ):
         if not self.is_fitted:
             raise ValueError(
                 "[Data Processor Error] DataProcessor must be fitted before transform"
@@ -880,7 +880,7 @@ class DataProcessor(FeatureSet):
         save_format: Optional[Literal["csv", "parquet"]] = None,
         output_path: Optional[str] = None,
         chunk_size: int = 200000,
-    ) -> Union[pd.DataFrame, Dict[str, np.ndarray], list[str]]:
+    ):
         self.fit(data, chunk_size=chunk_size)
         return self.transform(
             data,

nextrec/loss/__init__.py

@@ -5,6 +5,7 @@ from nextrec.loss.listwise import (
     ListNetLoss,
     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 (
@@ -30,6 +31,8 @@ __all__ = [
     "ListNetLoss",
     "ListMLELoss",
     "ApproxNDCGLoss",
+    # Multi-task weighting
+    "GradNormLossWeighting",
     # Utilities
     "get_loss_fn",
     "get_loss_kwargs",