nextrec 0.1.4__py3-none-any.whl → 0.1.8__py3-none-any.whl

nextrec/basic/model.py

@@ -23,64 +23,53 @@ from nextrec.basic.callback import EarlyStopper
 from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
 from nextrec.basic.metrics import configure_metrics, evaluate_metrics
 
-from nextrec.loss import get_loss_fn
 from nextrec.data import get_column_data
 from nextrec.basic.loggers import setup_logger, colorize
 from nextrec.utils import get_optimizer_fn, get_scheduler_fn
-
+from nextrec.loss import get_loss_fn
 
 
 class BaseModel(nn.Module):
     @property
     def model_name(self) -> str:
         raise NotImplementedError
-
+    
     @property
     def task_type(self) -> str:
         raise NotImplementedError
 
-    def __init__(
-        self,
-        dense_features: list[DenseFeature] | None = None,
-        sparse_features: list[SparseFeature] | None = None,
-        sequence_features: list[SequenceFeature] | None = None,
-        target: list[str] | str | None = None,
-        task: str | list[str] = "binary",
-        device: str = "cpu",
-        embedding_l1_reg: float = 0.0,
-        dense_l1_reg: float = 0.0,
-        embedding_l2_reg: float = 0.0,
-        dense_l2_reg: float = 0.0,
-        early_stop_patience: int = 20,
-        model_id: str = "baseline",
-    ):
-
+    def __init__(self, 
+                 dense_features: list[DenseFeature] | None = None, 
+                 sparse_features: list[SparseFeature] | None = None, 
+                 sequence_features: list[SequenceFeature] | None = None,
+                 target: list[str] | str | None = None,
+                 task: str|list[str] = 'binary',
+                 device: str = 'cpu',
+                 embedding_l1_reg: float = 0.0,
+                 dense_l1_reg: float = 0.0,
+                 embedding_l2_reg: float = 0.0, 
+                 dense_l2_reg: float = 0.0,
+                 early_stop_patience: int = 20, 
+                 model_id: str = 'baseline'): 
+        
         super(BaseModel, self).__init__()
 
         try:
             self.device = torch.device(device)
         except Exception as e:
-            logging.warning(
-                colorize("Invalid device , defaulting to CPU.", color="yellow")
-            )
-            self.device = torch.device("cpu")
+            logging.warning(colorize("Invalid device , defaulting to CPU.", color='yellow'))
+            self.device = torch.device('cpu')
 
         self.dense_features = list(dense_features) if dense_features is not None else []
-        self.sparse_features = (
-            list(sparse_features) if sparse_features is not None else []
-        )
-        self.sequence_features = (
-            list(sequence_features) if sequence_features is not None else []
-        )
-
+        self.sparse_features = list(sparse_features) if sparse_features is not None else []
+        self.sequence_features = list(sequence_features) if sequence_features is not None else []
+        
         if isinstance(target, str):
             self.target = [target]
         else:
             self.target = list(target) if target is not None else []
-
-        self.target_index = {
-            target_name: idx for idx, target_name in enumerate(self.target)
-        }
+        
+        self.target_index = {target_name: idx for idx, target_name in enumerate(self.target)}
 
         self.task = task
         self.nums_task = len(task) if isinstance(task, list) else 1
@@ -90,48 +79,33 @@ class BaseModel(nn.Module):
         self._embedding_l2_reg = embedding_l2_reg
         self._dense_l2_reg = dense_l2_reg
 
-        self._regularization_weights = (
-            []
-        )  # list of dense weights for regularization, used to compute reg loss
-        self._embedding_params = (
-            []
-        )  # list of embedding weights for regularization, used to compute reg loss
+        self._regularization_weights = [] # list of dense weights for regularization, used to compute reg loss
+        self._embedding_params = [] # list of embedding weights for regularization, used to compute reg loss
 
         self.early_stop_patience = early_stop_patience
-        self._max_gradient_norm = 1.0  # Maximum gradient norm for gradient clipping
+        self._max_gradient_norm = 1.0   # Maximum gradient norm for gradient clipping
 
         project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
         self.model_id = model_id
-
-        checkpoint_dir = os.path.abspath(
-            os.path.join(project_root, "..", "checkpoints")
-        )
+        
+        checkpoint_dir = os.path.abspath(os.path.join(project_root, "..", "checkpoints"))
         os.makedirs(checkpoint_dir, exist_ok=True)
-        self.checkpoint = os.path.join(
-            checkpoint_dir,
-            f"{self.model_name}_{self.model_id}_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.model",
-        )
-        self.best = os.path.join(
-            checkpoint_dir, f"{self.model_name}_{self.model_id}_best.model"
-        )
+        self.checkpoint = os.path.join(checkpoint_dir, f"{self.model_name}_{self.model_id}_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.model")
+        self.best = os.path.join(checkpoint_dir, f"{self.model_name}_{self.model_id}_best.model")
 
         self._logger_initialized = False
         self._verbose = 1
 
-    def _register_regularization_weights(
-        self,
-        embedding_attr: str = "embedding",
-        exclude_modules: (
-            list[str] | None
-        ) = [],  # modules wont add regularization, example: ['fm', 'lr'] / ['fm.fc'] / etc.
-        include_modules: list[str] | None = [],
-    ):
+    def _register_regularization_weights(self, 
+                                        embedding_attr: str = 'embedding',
+                                        exclude_modules: list[str] | None = [], # modules wont add regularization, example: ['fm', 'lr'] / ['fm.fc'] / etc.
+                                        include_modules: list[str] | None = []):
 
         exclude_modules = exclude_modules or []
-
+        
         if hasattr(self, embedding_attr):
             embedding_layer = getattr(self, embedding_attr)
-            if hasattr(embedding_layer, "embed_dict"):
+            if hasattr(embedding_layer, 'embed_dict'):
                 for embed in embedding_layer.embed_dict.values():
                     self._embedding_params.append(embed.weight)
 
@@ -139,49 +113,40 @@ class BaseModel(nn.Module):
             # Skip self module
             if module is self:
                 continue
-
+            
             # Skip embedding layers
             if embedding_attr in name:
                 continue
-
+            
             # Skip BatchNorm and Dropout by checking module type
-            if isinstance(
-                module,
-                (
-                    nn.BatchNorm1d,
-                    nn.BatchNorm2d,
-                    nn.BatchNorm3d,
-                    nn.Dropout,
-                    nn.Dropout2d,
-                    nn.Dropout3d,
-                ),
-            ):
+            if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, 
+                                   nn.Dropout, nn.Dropout2d, nn.Dropout3d)):
                 continue
-
+            
             # White-list: only include modules whose names contain specific keywords
             if include_modules is not None:
                 should_include = any(inc_name in name for inc_name in include_modules)
                 if not should_include:
                     continue
-
+            
             # Black-list: exclude modules whose names contain specific keywords
             if any(exc_name in name for exc_name in exclude_modules):
                 continue
-
+            
             # Only add regularization for Linear layers
             if isinstance(module, nn.Linear):
                 self._regularization_weights.append(module.weight)
 
     def add_reg_loss(self) -> torch.Tensor:
         reg_loss = torch.tensor(0.0, device=self.device)
-
+        
         if self._embedding_l1_reg > 0 and len(self._embedding_params) > 0:
             for param in self._embedding_params:
                 reg_loss += self._embedding_l1_reg * torch.sum(torch.abs(param))
 
         if self._embedding_l2_reg > 0 and len(self._embedding_params) > 0:
             for param in self._embedding_params:
-                reg_loss += self._embedding_l2_reg * torch.sum(param**2)
+                reg_loss += self._embedding_l2_reg * torch.sum(param ** 2)
 
         if self._dense_l1_reg > 0 and len(self._regularization_weights) > 0:
             for param in self._regularization_weights:
@@ -189,16 +154,11 @@ class BaseModel(nn.Module):
 
         if self._dense_l2_reg > 0 and len(self._regularization_weights) > 0:
             for param in self._regularization_weights:
-                reg_loss += self._dense_l2_reg * torch.sum(param**2)
-
+                reg_loss += self._dense_l2_reg * torch.sum(param ** 2)
+        
         return reg_loss
 
-    def _to_tensor(
-        self,
-        value,
-        dtype: torch.dtype | None = None,
-        device: str | torch.device | None = None,
-    ) -> torch.Tensor:
+    def _to_tensor(self, value, dtype: torch.dtype | None = None, device: str | torch.device | None = None) -> torch.Tensor:        
         if value is None:
             raise ValueError("Cannot convert None to tensor.")
         if isinstance(value, torch.Tensor):
@@ -210,13 +170,11 @@ class BaseModel(nn.Module):
         target_device = device if device is not None else self.device
         return tensor.to(target_device)
 
-    def get_input(self, input_data: dict | pd.DataFrame):
+    def get_input(self, input_data: dict|pd.DataFrame):
         X_input = {}
-
-        all_features = (
-            self.dense_features + self.sparse_features + self.sequence_features
-        )
-
+        
+        all_features = self.dense_features + self.sparse_features + self.sequence_features
+        
         for feature in all_features:
             if feature.name not in input_data:
                 continue
@@ -240,12 +198,10 @@ class BaseModel(nn.Module):
                 if target_data is None:
                     continue
                 target_tensor = self._to_tensor(target_data, dtype=torch.float32)
-
+                
                 if target_tensor.dim() > 1:
                     target_tensor = target_tensor.view(target_tensor.size(0), -1)
-                    target_tensors.extend(
-                        torch.chunk(target_tensor, chunks=target_tensor.shape[1], dim=1)
-                    )
+                    target_tensors.extend(torch.chunk(target_tensor, chunks=target_tensor.shape[1], dim=1))
                 else:
                     target_tensors.append(target_tensor.view(-1, 1))
 
@@ -260,14 +216,14 @@ class BaseModel(nn.Module):
 
     def _set_metrics(self, metrics: list[str] | dict[str, list[str]] | None = None):
         """Configure metrics for model evaluation using the metrics module."""
-        self.metrics, self.task_specific_metrics, self.best_metrics_mode = (
-            configure_metrics(task=self.task, metrics=metrics, target_names=self.target)
-        )  # ['auc', 'logloss'], {'target1': ['auc', 'logloss'], 'target2': ['mse']}, 'max'
-
-        if not hasattr(self, "early_stopper") or self.early_stopper is None:
-            self.early_stopper = EarlyStopper(
-                patience=self.early_stop_patience, mode=self.best_metrics_mode
-            )
+        self.metrics, self.task_specific_metrics, self.best_metrics_mode = configure_metrics(
+            task=self.task,
+            metrics=metrics,
+            target_names=self.target
+        ) # ['auc', 'logloss'], {'target1': ['auc', 'logloss'], 'target2': ['mse']}, 'max'
+        
+        if not hasattr(self, 'early_stopper') or self.early_stopper is None:
+            self.early_stopper = EarlyStopper(patience=self.early_stop_patience, mode=self.best_metrics_mode)
 
     def _validate_task_configuration(self):
         """Validate that task type, number of tasks, targets, and loss functions are consistent."""
@@ -276,35 +232,35 @@ class BaseModel(nn.Module):
             num_tasks_from_task = len(self.task)
         else:
             num_tasks_from_task = 1
-
+        
         num_targets = len(self.target)
-
+        
         if self.nums_task != num_tasks_from_task:
             raise ValueError(
                 f"Number of tasks mismatch: nums_task={self.nums_task}, "
                 f"but task list has {num_tasks_from_task} tasks."
             )
-
+        
         if self.nums_task != num_targets:
             raise ValueError(
                 f"Number of tasks ({self.nums_task}) does not match number of target columns ({num_targets}). "
                 f"Tasks: {self.task}, Targets: {self.target}"
             )
-
+        
         # Check loss function consistency
-        if hasattr(self, "loss_fn"):
+        if hasattr(self, 'loss_fn'):
             num_loss_fns = len(self.loss_fn)
             if num_loss_fns != self.nums_task:
                 raise ValueError(
                     f"Number of loss functions ({num_loss_fns}) does not match number of tasks ({self.nums_task})."
                 )
-
+        
         # Validate task types with metrics and loss functions
         from nextrec.loss import VALID_TASK_TYPES
         from nextrec.basic.metrics import CLASSIFICATION_METRICS, REGRESSION_METRICS
-
+        
         tasks_to_check = self.task if isinstance(self.task, list) else [self.task]
-
+        
         for i, task_type in enumerate(tasks_to_check):
             # Validate task type
             if task_type not in VALID_TASK_TYPES:
@@ -312,26 +268,26 @@ class BaseModel(nn.Module):
                     f"Invalid task type '{task_type}' for task {i}. "
                     f"Valid types: {VALID_TASK_TYPES}"
                 )
-
+            
             # Check metrics compatibility
-            if hasattr(self, "task_specific_metrics") and self.task_specific_metrics:
+            if hasattr(self, 'task_specific_metrics') and self.task_specific_metrics:
                 target_name = self.target[i] if i < len(self.target) else f"task_{i}"
                 task_metrics = self.task_specific_metrics.get(target_name, self.metrics)
-
+                
                 for metric in task_metrics:
                     metric_lower = metric.lower()
                     # Skip gauc as it's valid for both classification and regression in some contexts
-                    if metric_lower == "gauc":
+                    if metric_lower == 'gauc':
                         continue
-
-                    if task_type in ["binary", "multiclass"]:
+                    
+                    if task_type in ['binary', 'multiclass']:
                         # Classification task
                         if metric_lower in REGRESSION_METRICS:
                             raise ValueError(
                                 f"Metric '{metric}' is not compatible with classification task type '{task_type}' "
                                 f"for target '{target_name}'. Classification metrics: {CLASSIFICATION_METRICS}"
                             )
-                    elif task_type in ["regression", "multivariate_regression"]:
+                    elif task_type in ['regression', 'multivariate_regression']:
                         # Regression task
                         if metric_lower in CLASSIFICATION_METRICS:
                             raise ValueError(
@@ -339,72 +295,62 @@ class BaseModel(nn.Module):
                                 f"for target '{target_name}'. Regression metrics: {REGRESSION_METRICS}"
                             )
 
-    def _handle_validation_split(
-        self,
-        train_data: dict | pd.DataFrame | DataLoader,
-        validation_split: float,
-        batch_size: int,
-        shuffle: bool,
-    ) -> tuple[DataLoader, dict | pd.DataFrame]:
+    def _handle_validation_split(self, 
+                                 train_data: dict | pd.DataFrame | DataLoader,
+                                 validation_split: float,
+                                 batch_size: int,
+                                 shuffle: bool) -> tuple[DataLoader, dict | pd.DataFrame]:
         """Handle validation split logic for training data.
-
+        
         Args:
             train_data: Training data (dict, DataFrame, or DataLoader)
             validation_split: Fraction of data to use for validation (0 < validation_split < 1)
             batch_size: Batch size for DataLoader
             shuffle: Whether to shuffle training data
-
+            
         Returns:
             tuple: (train_loader, valid_data)
         """
         if not (0 < validation_split < 1):
-            raise ValueError(
-                f"validation_split must be between 0 and 1, got {validation_split}"
-            )
-
+            raise ValueError(f"validation_split must be between 0 and 1, got {validation_split}")
+        
         if isinstance(train_data, DataLoader):
             raise ValueError(
                 "validation_split cannot be used when train_data is a DataLoader. "
                 "Please provide dict or pd.DataFrame for train_data."
             )
-
+        
         if isinstance(train_data, pd.DataFrame):
             # Shuffle and split DataFrame
-            shuffled_df = train_data.sample(frac=1.0, random_state=42).reset_index(
-                drop=True
-            )
+            shuffled_df = train_data.sample(frac=1.0, random_state=42).reset_index(drop=True)
             split_idx = int(len(shuffled_df) * (1 - validation_split))
             train_split = shuffled_df.iloc[:split_idx]
             valid_split = shuffled_df.iloc[split_idx:]
-
-            train_loader = self._prepare_data_loader(
-                train_split, batch_size=batch_size, shuffle=shuffle
-            )
-
+            
+            train_loader = self._prepare_data_loader(train_split, batch_size=batch_size, shuffle=shuffle)
+            
             if self._verbose:
-                logging.info(
-                    colorize(
-                        f"Split data: {len(train_split)} training samples, {len(valid_split)} validation samples",
-                        color="cyan",
-                    )
-                )
-
+                logging.info(colorize(
+                    f"Split data: {len(train_split)} training samples, {len(valid_split)} validation samples",
+                    color="cyan"
+                ))
+            
             return train_loader, valid_split
-
+        
         elif isinstance(train_data, dict):
             # Get total length from any feature
             sample_key = list(train_data.keys())[0]
             total_length = len(train_data[sample_key])
-
+            
             # Create indices and shuffle
             indices = np.arange(total_length)
             np.random.seed(42)
             np.random.shuffle(indices)
-
+            
             split_idx = int(total_length * (1 - validation_split))
             train_indices = indices[:split_idx]
             valid_indices = indices[split_idx:]
-
+            
             # Split dict
             train_split = {}
             valid_split = {}
@@ -422,112 +368,82 @@ class BaseModel(nn.Module):
                 else:
                     train_split[key] = [value[i] for i in train_indices]
                     valid_split[key] = [value[i] for i in valid_indices]
-
-            train_loader = self._prepare_data_loader(
-                train_split, batch_size=batch_size, shuffle=shuffle
-            )
-
+            
+            train_loader = self._prepare_data_loader(train_split, batch_size=batch_size, shuffle=shuffle)
+            
             if self._verbose:
-                logging.info(
-                    colorize(
-                        f"Split data: {len(train_indices)} training samples, {len(valid_indices)} validation samples",
-                        color="cyan",
-                    )
-                )
-
+                logging.info(colorize(
+                    f"Split data: {len(train_indices)} training samples, {len(valid_indices)} validation samples",
+                    color="cyan"
+                ))
+            
             return train_loader, valid_split
-
+        
         else:
             raise TypeError(f"Unsupported train_data type: {type(train_data)}")
 
-    def compile(
-        self,
-        optimizer="adam",
-        optimizer_params: dict | None = None,
-        scheduler: (
-            str
-            | torch.optim.lr_scheduler._LRScheduler
-            | type[torch.optim.lr_scheduler._LRScheduler]
-            | None
-        ) = None,
-        scheduler_params: dict | None = None,
-        loss: str | nn.Module | list[str | nn.Module] | None = "bce",
-    ):
+
+    def compile(self, 
+                optimizer = "adam",
+                optimizer_params: dict | None = None,
+                scheduler: str | torch.optim.lr_scheduler._LRScheduler | type[torch.optim.lr_scheduler._LRScheduler] | None = None,
+                scheduler_params: dict | None = None,
+                loss: str | nn.Module | list[str | nn.Module] | None= "bce"):
         if optimizer_params is None:
             optimizer_params = {}
-
-        self._optimizer_name = (
-            optimizer if isinstance(optimizer, str) else optimizer.__class__.__name__
-        )
+        
+        self._optimizer_name = optimizer if isinstance(optimizer, str) else optimizer.__class__.__name__
         self._optimizer_params = optimizer_params
         if isinstance(scheduler, str):
             self._scheduler_name = scheduler
         elif scheduler is not None:
             # Try to get __name__ first (for class types), then __class__.__name__ (for instances)
-            self._scheduler_name = getattr(
-                scheduler,
-                "__name__",
-                getattr(scheduler.__class__, "__name__", str(scheduler)),
-            )
+            self._scheduler_name = getattr(scheduler, '__name__', getattr(scheduler.__class__, '__name__', str(scheduler)))
         else:
             self._scheduler_name = None
         self._scheduler_params = scheduler_params or {}
         self._loss_config = loss
-
+        
         # set optimizer
         self.optimizer_fn = get_optimizer_fn(
-            optimizer=optimizer, params=self.parameters(), **optimizer_params
+            optimizer=optimizer, 
+            params=self.parameters(), 
+            **optimizer_params
         )
-
+        
         # set loss functions
         if self.nums_task == 1:
             task_type = self.task if isinstance(self.task, str) else self.task[0]
             loss_value = loss[0] if isinstance(loss, list) else loss
             # For ranking and multitask, use pointwise training
-            training_mode = (
-                "pointwise" if self.task_type in ["ranking", "multitask"] else None
-            )
+            training_mode = 'pointwise' if self.task_type in ['ranking', 'multitask'] else None
             # Use task_type directly, not self.task_type for single task
-            self.loss_fn = [
-                get_loss_fn(
-                    task_type=task_type, training_mode=training_mode, loss=loss_value
-                )
-            ]
+            self.loss_fn = [get_loss_fn(task_type=task_type, training_mode=training_mode, loss=loss_value)]
         else:
             self.loss_fn = []
             for i in range(self.nums_task):
                 task_type = self.task[i] if isinstance(self.task, list) else self.task
-
+                
                 if isinstance(loss, list):
                     loss_value = loss[i] if i < len(loss) else None
                 else:
                     loss_value = loss
-
+                
                 # Multitask always uses pointwise training
-                training_mode = "pointwise"
-                self.loss_fn.append(
-                    get_loss_fn(
-                        task_type=task_type,
-                        training_mode=training_mode,
-                        loss=loss_value,
-                    )
-                )
-
+                training_mode = 'pointwise'
+                self.loss_fn.append(get_loss_fn(task_type=task_type, training_mode=training_mode, loss=loss_value))
+        
         # set scheduler
-        self.scheduler_fn = (
-            get_scheduler_fn(scheduler, self.optimizer_fn, **(scheduler_params or {}))
-            if scheduler
-            else None
-        )
+        self.scheduler_fn = get_scheduler_fn(scheduler, self.optimizer_fn, **(scheduler_params or {})) if scheduler else None
 
     def compute_loss(self, y_pred, y_true):
         if y_true is None:
             return torch.tensor(0.0, device=self.device)
-
+        
         if self.nums_task == 1:
             loss = self.loss_fn[0](y_pred, y_true)
             return loss
-
+        
         else:
             task_losses = []
             for i in range(self.nums_task):
@@ -535,69 +451,48 @@ class BaseModel(nn.Module):
                 task_losses.append(task_loss)
             return torch.stack(task_losses)
 
-    def _prepare_data_loader(
-        self,
-        data: dict | pd.DataFrame | DataLoader,
-        batch_size: int = 32,
-        shuffle: bool = True,
-    ):
+
+    def _prepare_data_loader(self, data: dict|pd.DataFrame|DataLoader, batch_size: int = 32, shuffle: bool = True):
         if isinstance(data, DataLoader):
             return data
         tensors = []
-        all_features = (
-            self.dense_features + self.sparse_features + self.sequence_features
-        )
-
+        all_features = self.dense_features + self.sparse_features + self.sequence_features
+        
         for feature in all_features:
             column = get_column_data(data, feature.name)
             if column is None:
                 raise KeyError(f"Feature {feature.name} not found in provided data.")
-
+            
             if isinstance(feature, SequenceFeature):
                 if isinstance(column, pd.Series):
                     column = column.values
                 if isinstance(column, np.ndarray) and column.dtype == object:
-                    column = np.array(
-                        [
-                            (
-                                np.array(seq, dtype=np.int64)
-                                if not isinstance(seq, np.ndarray)
-                                else seq
-                            )
-                            for seq in column
-                        ]
-                    )
-                if (
-                    isinstance(column, np.ndarray)
-                    and column.ndim == 1
-                    and column.dtype == object
-                ):
+                    column = np.array([np.array(seq, dtype=np.int64) if not isinstance(seq, np.ndarray) else seq for seq in column])
+                if isinstance(column, np.ndarray) and column.ndim == 1 and column.dtype == object:
                     column = np.vstack([c if isinstance(c, np.ndarray) else np.array(c) for c in column])  # type: ignore
-                tensor = torch.from_numpy(np.asarray(column, dtype=np.int64)).to("cpu")
+                tensor = torch.from_numpy(np.asarray(column, dtype=np.int64)).to('cpu')
             else:
-                dtype = (
-                    torch.float32 if isinstance(feature, DenseFeature) else torch.long
-                )
-                tensor = self._to_tensor(column, dtype=dtype, device="cpu")
-
+                dtype = torch.float32 if isinstance(feature, DenseFeature) else torch.long
+                tensor = self._to_tensor(column, dtype=dtype, device='cpu')
+            
             tensors.append(tensor)
-
+        
         label_tensors = []
         for target_name in self.target:
             column = get_column_data(data, target_name)
             if column is None:
                 continue
-            label_tensor = self._to_tensor(column, dtype=torch.float32, device="cpu")
-
+            label_tensor = self._to_tensor(column, dtype=torch.float32, device='cpu')
+            
             if label_tensor.dim() == 1:
                 # 1D tensor: (N,) -> (N, 1)
                 label_tensor = label_tensor.view(-1, 1)
             elif label_tensor.dim() == 2:
                 if label_tensor.shape[0] == 1 and label_tensor.shape[1] > 1:
                     label_tensor = label_tensor.t()
-
+            
             label_tensors.append(label_tensor)
-
+        
         if label_tensors:
             if len(label_tensors) == 1 and label_tensors[0].shape[1] > 1:
                 y_tensor = label_tensors[0]
@@ -607,23 +502,22 @@ class BaseModel(nn.Module):
             if y_tensor.shape[1] == 1:
                 y_tensor = y_tensor.squeeze(1)
             tensors.append(y_tensor)
-
+        
         dataset = TensorDataset(*tensors)
         return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
 
+
     def _batch_to_dict(self, batch_data: tuple) -> dict:
         result = {}
-        all_features = (
-            self.dense_features + self.sparse_features + self.sequence_features
-        )
-
+        all_features = self.dense_features + self.sparse_features + self.sequence_features
+        
         for i, feature in enumerate(all_features):
             if i < len(batch_data):
                 result[feature.name] = batch_data[i]
-
+        
         if len(batch_data) > len(all_features):
             labels = batch_data[-1]
-
+            
             if self.nums_task == 1:
                 result[self.target[0]] = labels
             else:
@@ -640,36 +534,28 @@ class BaseModel(nn.Module):
                     for i, target_name in enumerate(self.target):
                         if i < labels.shape[-1]:
                             result[target_name] = labels[..., i]
-
+        
         return result
 
-    def fit(
-        self,
-        train_data: dict | pd.DataFrame | DataLoader,
-        valid_data: dict | pd.DataFrame | DataLoader | None = None,
-        metrics: (
-            list[str] | dict[str, list[str]] | None
-        ) = None,  # ['auc', 'logloss'] or {'target1': ['auc', 'logloss'], 'target2': ['mse']}
-        epochs: int = 1,
-        verbose: int = 1,
-        shuffle: bool = True,
-        batch_size: int = 32,
-        user_id_column: str = "user_id",
-        validation_split: float | None = None,
-    ):
+
+    def fit(self, 
+            train_data: dict|pd.DataFrame|DataLoader, 
+            valid_data: dict|pd.DataFrame|DataLoader|None=None, 
+            metrics: list[str]|dict[str, list[str]]|None = None, # ['auc', 'logloss'] or {'target1': ['auc', 'logloss'], 'target2': ['mse']}
+            epochs:int=1, verbose:int=1, shuffle:bool=True, batch_size:int=32,
+            user_id_column: str = 'user_id',
+            validation_split: float | None = None):
 
         self.to(self.device)
         if not self._logger_initialized:
             setup_logger()
             self._logger_initialized = True
         self._verbose = verbose
-        self._set_metrics(
-            metrics
-        )  # add self.metrics, self.task_specific_metrics, self.best_metrics_mode, self.early_stopper
-
+        self._set_metrics(metrics) # add self.metrics, self.task_specific_metrics, self.best_metrics_mode, self.early_stopper
+        
         # Assert before training
         self._validate_task_configuration()
-
+        
         if self._verbose:
             self.summary()
 
@@ -680,85 +566,63 @@ class BaseModel(nn.Module):
                 train_data=train_data,
                 validation_split=validation_split,
                 batch_size=batch_size,
-                shuffle=shuffle,
+                shuffle=shuffle
             )
         else:
             if not isinstance(train_data, DataLoader):
-                train_loader = self._prepare_data_loader(
-                    train_data, batch_size=batch_size, shuffle=shuffle
-                )
+                train_loader = self._prepare_data_loader(train_data, batch_size=batch_size, shuffle=shuffle)
             else:
                 train_loader = train_data
+        
 
         valid_user_ids: np.ndarray | None = None
         needs_user_ids = self._needs_user_ids_for_metrics()
 
         if valid_loader is None:
             if valid_data is not None and not isinstance(valid_data, DataLoader):
-                valid_loader = self._prepare_data_loader(
-                    valid_data, batch_size=batch_size, shuffle=False
-                )
+                valid_loader = self._prepare_data_loader(valid_data, batch_size=batch_size, shuffle=False)
                 # Extract user_ids only if needed for GAUC
                 if needs_user_ids:
-                    if (
-                        isinstance(valid_data, pd.DataFrame)
-                        and user_id_column in valid_data.columns
-                    ):
+                    if isinstance(valid_data, pd.DataFrame) and user_id_column in valid_data.columns:
                         valid_user_ids = np.asarray(valid_data[user_id_column].values)
                     elif isinstance(valid_data, dict) and user_id_column in valid_data:
                         valid_user_ids = np.asarray(valid_data[user_id_column])
             elif valid_data is not None:
                 valid_loader = valid_data
-
+        
         try:
             self._steps_per_epoch = len(train_loader)
             is_streaming = False
         except TypeError:
             self._steps_per_epoch = None
             is_streaming = True
-
+        
         self._epoch_index = 0
         self._stop_training = False
-        self._best_metric = (
-            float("-inf") if self.best_metrics_mode == "max" else float("inf")
-        )
-
+        self._best_metric = float('-inf') if self.best_metrics_mode == 'max' else float('inf')
+        
         if self._verbose:
             logging.info("")
             logging.info(colorize("=" * 80, color="bright_green", bold=True))
             if is_streaming:
-                logging.info(
-                    colorize(
-                        f"Start training (Streaming Mode)",
-                        color="bright_green",
-                        bold=True,
-                    )
-                )
+                logging.info(colorize(f"Start training (Streaming Mode)", color="bright_green", bold=True))
             else:
-                logging.info(
-                    colorize(f"Start training", color="bright_green", bold=True)
-                )
+                logging.info(colorize(f"Start training", color="bright_green", bold=True))
             logging.info(colorize("=" * 80, color="bright_green", bold=True))
             logging.info("")
             logging.info(colorize(f"Model device: {self.device}", color="bright_green"))
-
+        
         for epoch in range(epochs):
             self._epoch_index = epoch
-
+            
             # In streaming mode, print epoch header before progress bar
             if self._verbose and is_streaming:
                 logging.info("")
-                logging.info(
-                    colorize(
-                        f"Epoch {epoch + 1}/{epochs}", color="bright_green", bold=True
-                    )
-                )
+                logging.info(colorize(f"Epoch {epoch + 1}/{epochs}", color="bright_green", bold=True))
 
             # Train with metrics computation
-            train_result = self.train_epoch(
-                train_loader, is_streaming=is_streaming, compute_metrics=True
-            )
-
+            train_result = self.train_epoch(train_loader, is_streaming=is_streaming, compute_metrics=True)
+            
             # Unpack results
             if isinstance(train_result, tuple):
                 train_loss, train_metrics = train_result
@@ -768,13 +632,9 @@ class BaseModel(nn.Module):
 
             if self._verbose:
                 if self.nums_task == 1:
-                    log_str = (
-                        f"Epoch {epoch + 1}/{epochs} - Train: loss={train_loss:.4f}"
-                    )
+                    log_str = f"Epoch {epoch + 1}/{epochs} - Train: loss={train_loss:.4f}"
                     if train_metrics:
-                        metrics_str = ", ".join(
-                            [f"{k}={v:.4f}" for k, v in train_metrics.items()]
-                        )
+                        metrics_str = ", ".join([f"{k}={v:.4f}" for k, v in train_metrics.items()])
                         log_str += f", {metrics_str}"
                     logging.info(colorize(log_str, color="white"))
                 else:
@@ -784,12 +644,10 @@ class BaseModel(nn.Module):
                             task_labels.append(self.target[i])
                         else:
                             task_labels.append(f"task_{i}")
-
+                    
                     total_loss_val = np.sum(train_loss) if isinstance(train_loss, np.ndarray) else train_loss  # type: ignore
-                    log_str = (
-                        f"Epoch {epoch + 1}/{epochs} - Train: loss={total_loss_val:.4f}"
-                    )
-
+                    log_str = f"Epoch {epoch + 1}/{epochs} - Train: loss={total_loss_val:.4f}"
+                    
                     if train_metrics:
                         # Group metrics by task
                         task_metrics = {}
@@ -798,50 +656,28 @@ class BaseModel(nn.Module):
                                 if metric_key.endswith(f"_{target_name}"):
                                     if target_name not in task_metrics:
                                         task_metrics[target_name] = {}
-                                    metric_name = metric_key.rsplit(
-                                        f"_{target_name}", 1
-                                    )[0]
-                                    task_metrics[target_name][
-                                        metric_name
-                                    ] = metric_value
+                                    metric_name = metric_key.rsplit(f"_{target_name}", 1)[0]
+                                    task_metrics[target_name][metric_name] = metric_value
                                     break
-
+                        
                         if task_metrics:
                             task_metric_strs = []
                             for target_name in self.target:
                                 if target_name in task_metrics:
-                                    metrics_str = ", ".join(
-                                        [
-                                            f"{k}={v:.4f}"
-                                            for k, v in task_metrics[
-                                                target_name
-                                            ].items()
-                                        ]
-                                    )
-                                    task_metric_strs.append(
-                                        f"{target_name}[{metrics_str}]"
-                                    )
+                                    metrics_str = ", ".join([f"{k}={v:.4f}" for k, v in task_metrics[target_name].items()])
+                                    task_metric_strs.append(f"{target_name}[{metrics_str}]")
                             log_str += ", " + ", ".join(task_metric_strs)
-
+                    
                     logging.info(colorize(log_str, color="white"))
-
+            
             if valid_loader is not None:
                 # Pass user_ids only if needed for GAUC metric
-                val_metrics = self.evaluate(
-                    valid_loader, user_ids=valid_user_ids if needs_user_ids else None
-                )  # {'auc': 0.75, 'logloss': 0.45} or {'auc_target1': 0.75, 'logloss_target1': 0.45, 'mse_target2': 3.2}
-
+                val_metrics = self.evaluate(valid_loader, user_ids=valid_user_ids if needs_user_ids else None) # {'auc': 0.75, 'logloss': 0.45} or {'auc_target1': 0.75, 'logloss_target1': 0.45, 'mse_target2': 3.2}
+            
                 if self._verbose:
                     if self.nums_task == 1:
-                        metrics_str = ", ".join(
-                            [f"{k}={v:.4f}" for k, v in val_metrics.items()]
-                        )
-                        logging.info(
-                            colorize(
-                                f"Epoch {epoch + 1}/{epochs} - Valid: {metrics_str}",
-                                color="cyan",
-                            )
-                        )
+                        metrics_str = ", ".join([f"{k}={v:.4f}" for k, v in val_metrics.items()])
+                        logging.info(colorize(f"Epoch {epoch + 1}/{epochs} - Valid: {metrics_str}", color="cyan"))
                     else:
                         # multi task metrics
                         task_metrics = {}
@@ -850,55 +686,33 @@ class BaseModel(nn.Module):
                                 if metric_key.endswith(f"_{target_name}"):
                                     if target_name not in task_metrics:
                                         task_metrics[target_name] = {}
-                                    metric_name = metric_key.rsplit(
-                                        f"_{target_name}", 1
-                                    )[0]
-                                    task_metrics[target_name][
-                                        metric_name
-                                    ] = metric_value
+                                    metric_name = metric_key.rsplit(f"_{target_name}", 1)[0]
+                                    task_metrics[target_name][metric_name] = metric_value
                                     break
 
                         task_metric_strs = []
                         for target_name in self.target:
                             if target_name in task_metrics:
-                                metrics_str = ", ".join(
-                                    [
-                                        f"{k}={v:.4f}"
-                                        for k, v in task_metrics[target_name].items()
-                                    ]
-                                )
+                                metrics_str = ", ".join([f"{k}={v:.4f}" for k, v in task_metrics[target_name].items()])
                                 task_metric_strs.append(f"{target_name}[{metrics_str}]")
-
-                        logging.info(
-                            colorize(
-                                f"Epoch {epoch + 1}/{epochs} - Valid: "
-                                + ", ".join(task_metric_strs),
-                                color="cyan",
-                            )
-                        )
-
+                        
+                        logging.info(colorize(f"Epoch {epoch + 1}/{epochs} - Valid: " + ", ".join(task_metric_strs), color="cyan"))
+                
                 # Handle empty validation metrics
                 if not val_metrics:
                     if self._verbose:
-                        logging.info(
-                            colorize(
-                                f"Warning: No validation metrics computed. Skipping validation for this epoch.",
-                                color="yellow",
-                            )
-                        )
+                        logging.info(colorize(f"Warning: No validation metrics computed. Skipping validation for this epoch.", color="yellow"))
                     continue
-
+                
                 if self.nums_task == 1:
                     primary_metric_key = self.metrics[0]
                 else:
                     primary_metric_key = f"{self.metrics[0]}_{self.target[0]}"
-
-                primary_metric = val_metrics.get(
-                    primary_metric_key, val_metrics[list(val_metrics.keys())[0]]
-                )
+                
+                primary_metric = val_metrics.get(primary_metric_key, val_metrics[list(val_metrics.keys())[0]])
                 improved = False
-
-                if self.best_metrics_mode == "max":
+                
+                if self.best_metrics_mode == 'max':
                     if primary_metric > self._best_metric:
                         self._best_metric = primary_metric
                         self.save_weights(self.best)
@@ -907,85 +721,56 @@ class BaseModel(nn.Module):
                     if primary_metric < self._best_metric:
                         self._best_metric = primary_metric
                         improved = True
-
+                
                 if improved:
                     if self._verbose:
-                        logging.info(
-                            colorize(
-                                f"Validation {primary_metric_key} improved to {self._best_metric:.4f}",
-                                color="yellow",
-                            )
-                        )
+                        logging.info(colorize(f"Validation {primary_metric_key} improved to {self._best_metric:.4f}", color="yellow"))
                     self.save_weights(self.checkpoint)
                     self.early_stopper.trial_counter = 0
                 else:
                     self.early_stopper.trial_counter += 1
                     if self._verbose:
-                        logging.info(
-                            colorize(
-                                f"No improvement for {self.early_stopper.trial_counter} epoch(s)",
-                                color="yellow",
-                            )
-                        )
-
+                        logging.info(colorize(f"No improvement for {self.early_stopper.trial_counter} epoch(s)", color="yellow"))
+                
                 if self.early_stopper.trial_counter >= self.early_stopper.patience:
                     self._stop_training = True
                     if self._verbose:
-                        logging.info(
-                            colorize(
-                                f"Early stopping triggered after {epoch + 1} epochs",
-                                color="bright_red",
-                                bold=True,
-                            )
-                        )
+                        logging.info(colorize(f"Early stopping triggered after {epoch + 1} epochs", color="bright_red", bold=True))
                     break
             else:
                 self.save_weights(self.checkpoint)
-
+            
             if self._stop_training:
                 break
-
+            
             if self.scheduler_fn is not None:
-                if isinstance(
-                    self.scheduler_fn, torch.optim.lr_scheduler.ReduceLROnPlateau
-                ):
+                if isinstance(self.scheduler_fn, torch.optim.lr_scheduler.ReduceLROnPlateau):
                     if valid_loader is not None:
                         self.scheduler_fn.step(primary_metric)
                 else:
                     self.scheduler_fn.step()
-
+                    
         if self._verbose:
             logging.info("\n")
-            logging.info(
-                colorize("Training finished.", color="bright_green", bold=True)
-            )
+            logging.info(colorize("Training finished.", color="bright_green", bold=True))
             logging.info("\n")
-
+        
         if valid_loader is not None:
             if self._verbose:
-                logging.info(
-                    colorize(
-                        f"Load best model from: {self.checkpoint}", color="bright_blue"
-                    )
-                )
+                logging.info(colorize(f"Load best model from: {self.checkpoint}", color="bright_blue"))
             self.load_weights(self.checkpoint)
-
+        
         return self
 
-    def train_epoch(
-        self,
-        train_loader: DataLoader,
-        is_streaming: bool = False,
-        compute_metrics: bool = False,
-    ) -> Union[float, np.ndarray, tuple[Union[float, np.ndarray], dict]]:
+    def train_epoch(self, train_loader: DataLoader, is_streaming: bool = False, compute_metrics: bool = False) -> Union[float, np.ndarray, tuple[Union[float, np.ndarray], dict]]:
         if self.nums_task == 1:
             accumulated_loss = 0.0
         else:
             accumulated_loss = np.zeros(self.nums_task, dtype=np.float64)
-
+        
         self.train()
         num_batches = 0
-
+        
         # Lists to store predictions and labels for metric computation
         y_true_list = []
         y_pred_list = []
@@ -993,29 +778,19 @@ class BaseModel(nn.Module):
         if self._verbose:
             # For streaming datasets without known length, set total=None to show progress without percentage
             if self._steps_per_epoch is not None:
-                batch_iter = enumerate(
-                    tqdm.tqdm(
-                        train_loader,
-                        desc=f"Epoch {self._epoch_index + 1}",
-                        total=self._steps_per_epoch,
-                    )
-                )
+                batch_iter = enumerate(tqdm.tqdm(train_loader, desc=f"Epoch {self._epoch_index + 1}", total=self._steps_per_epoch))
             else:
                 # Streaming mode: show batch/file progress without epoch in desc
                 if is_streaming:
-                    batch_iter = enumerate(
-                        tqdm.tqdm(
-                            train_loader,
-                            desc="Batches",
-                            # position=1,
-                            # leave=False,
-                            # unit="batch"
-                        )
-                    )
+                    batch_iter = enumerate(tqdm.tqdm(
+                        train_loader, 
+                        desc="Batches", 
+                        # position=1,
+                        # leave=False,
+                        # unit="batch"
+                    ))
                 else:
-                    batch_iter = enumerate(
-                        tqdm.tqdm(train_loader, desc=f"Epoch {self._epoch_index + 1}")
-                    )
+                    batch_iter = enumerate(tqdm.tqdm(train_loader, desc=f"Epoch {self._epoch_index + 1}"))
         else:
             batch_iter = enumerate(train_loader)
 
@@ -1031,17 +806,17 @@ class BaseModel(nn.Module):
                 total_loss = loss + reg_loss
             else:
                 total_loss = loss.sum() + reg_loss
-
+            
             self.optimizer_fn.zero_grad()
             total_loss.backward()
             nn.utils.clip_grad_norm_(self.parameters(), self._max_gradient_norm)
             self.optimizer_fn.step()
-
+            
             if self.nums_task == 1:
                 accumulated_loss += loss.item()
             else:
                 accumulated_loss += loss.detach().cpu().numpy()
-
+            
             # Collect predictions and labels for metrics if requested
             if compute_metrics:
                 if y_true is not None:
@@ -1049,52 +824,49 @@ class BaseModel(nn.Module):
                 # For pairwise/listwise mode, y_pred is a tuple of embeddings, skip metric collection during training
                 if y_pred is not None and isinstance(y_pred, torch.Tensor):
                     y_pred_list.append(y_pred.detach().cpu().numpy())
-
+            
             num_batches += 1
 
         if self.nums_task == 1:
             avg_loss = accumulated_loss / num_batches
         else:
             avg_loss = accumulated_loss / num_batches
-
+        
         # Compute metrics if requested
         if compute_metrics and len(y_true_list) > 0 and len(y_pred_list) > 0:
             y_true_all = np.concatenate(y_true_list, axis=0)
             y_pred_all = np.concatenate(y_pred_list, axis=0)
-            metrics_dict = self.evaluate_metrics(
-                y_true_all, y_pred_all, self.metrics, user_ids=None
-            )
+            metrics_dict = self.evaluate_metrics(y_true_all, y_pred_all, self.metrics, user_ids=None)
             return avg_loss, metrics_dict
-
+        
         return avg_loss
 
+
     def _needs_user_ids_for_metrics(self) -> bool:
         """Check if any configured metric requires user_ids (e.g., gauc)."""
         all_metrics = set()
-
+        
         # Collect all metrics from different sources
-        if hasattr(self, "metrics") and self.metrics:
+        if hasattr(self, 'metrics') and self.metrics:
             all_metrics.update(m.lower() for m in self.metrics)
-
-        if hasattr(self, "task_specific_metrics") and self.task_specific_metrics:
+        
+        if hasattr(self, 'task_specific_metrics') and self.task_specific_metrics:
             for task_metrics in self.task_specific_metrics.values():
                 if isinstance(task_metrics, list):
                     all_metrics.update(m.lower() for m in task_metrics)
-
+        
         # Check if gauc is in any of the metrics
-        return "gauc" in all_metrics
-
-    def evaluate(
-        self,
-        data: dict | pd.DataFrame | DataLoader,
-        metrics: list[str] | dict[str, list[str]] | None = None,
-        batch_size: int = 32,
-        user_ids: np.ndarray | None = None,
-        user_id_column: str = "user_id",
-    ) -> dict:
+        return 'gauc' in all_metrics
+
+    def evaluate(self, 
+                 data: dict | pd.DataFrame | DataLoader, 
+                 metrics: list[str] | dict[str, list[str]] | None = None,
+                 batch_size: int = 32,
+                 user_ids: np.ndarray | None = None,
+                 user_id_column: str = 'user_id') -> dict:
         """
         Evaluate the model on validation data.
-
+        
         Args:
             data: Evaluation data (dict, DataFrame, or DataLoader)
             metrics: Optional metrics to use for evaluation. If None, uses metrics from fit()
@@ -1102,19 +874,17 @@ class BaseModel(nn.Module):
             user_ids: Optional user IDs for computing GAUC metric. If None and gauc is needed,
                      will try to extract from data using user_id_column
             user_id_column: Column name for user IDs (default: 'user_id')
-
+            
         Returns:
             Dictionary of metric values
         """
         self.eval()
-
+        
         # Use provided metrics or fall back to configured metrics
         eval_metrics = metrics if metrics is not None else self.metrics
         if eval_metrics is None:
-            raise ValueError(
-                "No metrics specified for evaluation. Please provide metrics parameter or call fit() first."
-            )
-
+            raise ValueError("No metrics specified for evaluation. Please provide metrics parameter or call fit() first.")
+        
         # Prepare DataLoader if needed
         if isinstance(data, DataLoader):
             data_loader = data
@@ -1122,13 +892,11 @@ class BaseModel(nn.Module):
             if user_ids is None and self._needs_user_ids_for_metrics():
                 # Cannot extract user_ids from DataLoader, user must provide them
                 if self._verbose:
-                    logging.warning(
-                        colorize(
-                            "GAUC metric requires user_ids, but data is a DataLoader. "
-                            "Please provide user_ids parameter or use dict/DataFrame format.",
-                            color="yellow",
-                        )
-                    )
+                    logging.warning(colorize(
+                        "GAUC metric requires user_ids, but data is a DataLoader. "
+                        "Please provide user_ids parameter or use dict/DataFrame format.",
+                        color="yellow"
+                    ))
         else:
             # Extract user_ids if needed and not provided
             if user_ids is None and self._needs_user_ids_for_metrics():
@@ -1136,14 +904,12 @@ class BaseModel(nn.Module):
                     user_ids = np.asarray(data[user_id_column].values)
                 elif isinstance(data, dict) and user_id_column in data:
                     user_ids = np.asarray(data[user_id_column])
-
-            data_loader = self._prepare_data_loader(
-                data, batch_size=batch_size, shuffle=False
-            )
-
+            
+            data_loader = self._prepare_data_loader(data, batch_size=batch_size, shuffle=False)
+        
         y_true_list = []
         y_pred_list = []
-
+        
         batch_count = 0
         with torch.no_grad():
             for batch_data in data_loader:
@@ -1151,7 +917,7 @@ class BaseModel(nn.Module):
                 batch_dict = self._batch_to_dict(batch_data)
                 X_input, y_true = self.get_input(batch_dict)
                 y_pred = self.forward(X_input)
-
+                
                 if y_true is not None:
                     y_true_list.append(y_true.cpu().numpy())
                 # Skip if y_pred is not a tensor (e.g., tuple in pairwise mode, though this shouldn't happen in eval mode)
@@ -1159,40 +925,24 @@ class BaseModel(nn.Module):
                     y_pred_list.append(y_pred.cpu().numpy())
 
         if self._verbose:
-            logging.info(
-                colorize(f"  Evaluation batches processed: {batch_count}", color="cyan")
-            )
-
+            logging.info(colorize(f"  Evaluation batches processed: {batch_count}", color="cyan"))
+        
         if len(y_true_list) > 0:
             y_true_all = np.concatenate(y_true_list, axis=0)
             if self._verbose:
-                logging.info(
-                    colorize(
-                        f"  Evaluation samples: {y_true_all.shape[0]}", color="cyan"
-                    )
-                )
+                logging.info(colorize(f"  Evaluation samples: {y_true_all.shape[0]}", color="cyan"))
         else:
             y_true_all = None
             if self._verbose:
-                logging.info(
-                    colorize(
-                        f"  Warning: No y_true collected from evaluation data",
-                        color="yellow",
-                    )
-                )
-
+                logging.info(colorize(f"  Warning: No y_true collected from evaluation data", color="yellow"))
+            
         if len(y_pred_list) > 0:
             y_pred_all = np.concatenate(y_pred_list, axis=0)
         else:
             y_pred_all = None
             if self._verbose:
-                logging.info(
-                    colorize(
-                        f"  Warning: No y_pred collected from evaluation data",
-                        color="yellow",
-                    )
-                )
-
+                logging.info(colorize(f"  Warning: No y_pred collected from evaluation data", color="yellow"))
+        
         # Convert metrics to list if it's a dict
         if isinstance(eval_metrics, dict):
             # For dict metrics, we need to collect all unique metric names
@@ -1204,23 +954,16 @@ class BaseModel(nn.Module):
             metrics_to_use = unique_metrics
         else:
             metrics_to_use = eval_metrics
-
-        metrics_dict = self.evaluate_metrics(
-            y_true_all, y_pred_all, metrics_to_use, user_ids
-        )
-
+        
+        metrics_dict = self.evaluate_metrics(y_true_all, y_pred_all, metrics_to_use, user_ids)
+        
         return metrics_dict
 
-    def evaluate_metrics(
-        self,
-        y_true: np.ndarray | None,
-        y_pred: np.ndarray | None,
-        metrics: list[str],
-        user_ids: np.ndarray | None = None,
-    ) -> dict:
-        """Evaluate metrics using the metrics module."""
-        task_specific_metrics = getattr(self, "task_specific_metrics", None)
 
+    def evaluate_metrics(self, y_true: np.ndarray|None, y_pred: np.ndarray|None, metrics: list[str], user_ids: np.ndarray|None = None) -> dict:
+        """Evaluate metrics using the metrics module."""
+        task_specific_metrics = getattr(self, 'task_specific_metrics', None)
+        
         return evaluate_metrics(
             y_true=y_true,
             y_pred=y_pred,
@@ -1228,29 +971,24 @@ class BaseModel(nn.Module):
             task=self.task,
             target_names=self.target,
             task_specific_metrics=task_specific_metrics,
-            user_ids=user_ids,
+            user_ids=user_ids
         )
 
-    def predict(
-        self, data: str | dict | pd.DataFrame | DataLoader, batch_size: int = 32
-    ) -> np.ndarray:
+
+    def predict(self, data: str|dict|pd.DataFrame|DataLoader, batch_size: int = 32) -> np.ndarray:
         self.eval()
         # todo: handle file path input later
         if isinstance(data, (str, os.PathLike)):
             pass
         if not isinstance(data, DataLoader):
-            data_loader = self._prepare_data_loader(
-                data, batch_size=batch_size, shuffle=False
-            )
+            data_loader = self._prepare_data_loader(data, batch_size=batch_size, shuffle=False)
         else:
             data_loader = data
-
+        
         y_pred_list = []
-
+        
         with torch.no_grad():
-            for batch_data in tqdm.tqdm(
-                data_loader, desc="Predicting", disable=self._verbose == 0
-            ):
+            for batch_data in tqdm.tqdm(data_loader, desc="Predicting", disable=self._verbose == 0):
                 batch_dict = self._batch_to_dict(batch_data)
                 X_input, _ = self.get_input(batch_dict)
                 y_pred = self.forward(X_input)
@@ -1263,10 +1001,10 @@ class BaseModel(nn.Module):
             return y_pred_all
         else:
             return np.array([])
-
+    
     def save_weights(self, model_path: str):
         torch.save(self.state_dict(), model_path)
-
+    
     def load_weights(self, checkpoint):
         self.to(self.device)
         state_dict = torch.load(checkpoint, map_location="cpu")
@@ -1274,136 +1012,112 @@ class BaseModel(nn.Module):
 
     def summary(self):
         logger = logging.getLogger()
-
+        
         logger.info(colorize("=" * 80, color="bright_blue", bold=True))
-        logger.info(
-            colorize(
-                f"Model Summary: {self.model_name}", color="bright_blue", bold=True
-            )
-        )
+        logger.info(colorize(f"Model Summary: {self.model_name}", color="bright_blue", bold=True))
         logger.info(colorize("=" * 80, color="bright_blue", bold=True))
-
+        
         logger.info("")
         logger.info(colorize("[1] Feature Configuration", color="cyan", bold=True))
         logger.info(colorize("-" * 80, color="cyan"))
-
+        
         if self.dense_features:
             logger.info(f"Dense Features ({len(self.dense_features)}):")
             for i, feat in enumerate(self.dense_features, 1):
-                embed_dim = feat.embedding_dim if hasattr(feat, "embedding_dim") else 1
+                embed_dim = feat.embedding_dim if hasattr(feat, 'embedding_dim') else 1
                 logger.info(f"  {i}. {feat.name:20s}")
-
+        
         if self.sparse_features:
             logger.info(f"Sparse Features ({len(self.sparse_features)}):")
 
             max_name_len = max(len(feat.name) for feat in self.sparse_features)
-            max_embed_name_len = max(
-                len(feat.embedding_name) for feat in self.sparse_features
-            )
+            max_embed_name_len = max(len(feat.embedding_name) for feat in self.sparse_features)
             name_width = max(max_name_len, 10) + 2
             embed_name_width = max(max_embed_name_len, 15) + 2
-
-            logger.info(
-                f"  {'#':<4} {'Name':<{name_width}} {'Vocab Size':>12} {'Embed Name':>{embed_name_width}} {'Embed Dim':>10}"
-            )
-            logger.info(
-                f"  {'-'*4} {'-'*name_width} {'-'*12} {'-'*embed_name_width} {'-'*10}"
-            )
+            
+            logger.info(f"  {'#':<4} {'Name':<{name_width}} {'Vocab Size':>12} {'Embed Name':>{embed_name_width}} {'Embed Dim':>10}")
+            logger.info(f"  {'-'*4} {'-'*name_width} {'-'*12} {'-'*embed_name_width} {'-'*10}")
             for i, feat in enumerate(self.sparse_features, 1):
-                vocab_size = feat.vocab_size if hasattr(feat, "vocab_size") else "N/A"
-                embed_dim = (
-                    feat.embedding_dim if hasattr(feat, "embedding_dim") else "N/A"
-                )
-                logger.info(
-                    f"  {i:<4} {feat.name:<{name_width}} {str(vocab_size):>12} {feat.embedding_name:>{embed_name_width}} {str(embed_dim):>10}"
-                )
-
+                vocab_size = feat.vocab_size if hasattr(feat, 'vocab_size') else 'N/A'
+                embed_dim = feat.embedding_dim if hasattr(feat, 'embedding_dim') else 'N/A'
+                logger.info(f"  {i:<4} {feat.name:<{name_width}} {str(vocab_size):>12} {feat.embedding_name:>{embed_name_width}} {str(embed_dim):>10}")
+        
         if self.sequence_features:
             logger.info(f"Sequence Features ({len(self.sequence_features)}):")
 
             max_name_len = max(len(feat.name) for feat in self.sequence_features)
-            max_embed_name_len = max(
-                len(feat.embedding_name) for feat in self.sequence_features
-            )
+            max_embed_name_len = max(len(feat.embedding_name) for feat in self.sequence_features)
             name_width = max(max_name_len, 10) + 2
             embed_name_width = max(max_embed_name_len, 15) + 2
-
-            logger.info(
-                f"  {'#':<4} {'Name':<{name_width}} {'Vocab Size':>12} {'Embed Name':>{embed_name_width}} {'Embed Dim':>10} {'Max Len':>10}"
-            )
-            logger.info(
-                f"  {'-'*4} {'-'*name_width} {'-'*12} {'-'*embed_name_width} {'-'*10} {'-'*10}"
-            )
+            
+            logger.info(f"  {'#':<4} {'Name':<{name_width}} {'Vocab Size':>12} {'Embed Name':>{embed_name_width}} {'Embed Dim':>10} {'Max Len':>10}")
+            logger.info(f"  {'-'*4} {'-'*name_width} {'-'*12} {'-'*embed_name_width} {'-'*10} {'-'*10}")
             for i, feat in enumerate(self.sequence_features, 1):
-                vocab_size = feat.vocab_size if hasattr(feat, "vocab_size") else "N/A"
-                embed_dim = (
-                    feat.embedding_dim if hasattr(feat, "embedding_dim") else "N/A"
-                )
-                max_len = feat.max_len if hasattr(feat, "max_len") else "N/A"
-                logger.info(
-                    f"  {i:<4} {feat.name:<{name_width}} {str(vocab_size):>12} {feat.embedding_name:>{embed_name_width}} {str(embed_dim):>10} {str(max_len):>10}"
-                )
-
+                vocab_size = feat.vocab_size if hasattr(feat, 'vocab_size') else 'N/A'
+                embed_dim = feat.embedding_dim if hasattr(feat, 'embedding_dim') else 'N/A'
+                max_len = feat.max_len if hasattr(feat, 'max_len') else 'N/A'
+                logger.info(f"  {i:<4} {feat.name:<{name_width}} {str(vocab_size):>12} {feat.embedding_name:>{embed_name_width}} {str(embed_dim):>10} {str(max_len):>10}")
+        
         logger.info("")
         logger.info(colorize("[2] Model Parameters", color="cyan", bold=True))
         logger.info(colorize("-" * 80, color="cyan"))
-
+        
         # Model Architecture
         logger.info("Model Architecture:")
         logger.info(str(self))
         logger.info("")
-
+        
         total_params = sum(p.numel() for p in self.parameters())
         trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
         non_trainable_params = total_params - trainable_params
-
+        
         logger.info(f"Total Parameters:        {total_params:,}")
         logger.info(f"Trainable Parameters:    {trainable_params:,}")
         logger.info(f"Non-trainable Parameters: {non_trainable_params:,}")
-
+        
         logger.info("Layer-wise Parameters:")
         for name, module in self.named_children():
             layer_params = sum(p.numel() for p in module.parameters())
             if layer_params > 0:
                 logger.info(f"  {name:30s}: {layer_params:,}")
-
+        
         logger.info("")
         logger.info(colorize("[3] Training Configuration", color="cyan", bold=True))
         logger.info(colorize("-" * 80, color="cyan"))
-
+        
         logger.info(f"Task Type:               {self.task}")
         logger.info(f"Number of Tasks:         {self.nums_task}")
         logger.info(f"Metrics:                 {self.metrics}")
         logger.info(f"Target Columns:          {self.target}")
         logger.info(f"Device:                  {self.device}")
-
-        if hasattr(self, "_optimizer_name"):
+        
+        if hasattr(self, '_optimizer_name'):
             logger.info(f"Optimizer:               {self._optimizer_name}")
             if self._optimizer_params:
                 for key, value in self._optimizer_params.items():
                     logger.info(f"  {key:25s}: {value}")
-
-        if hasattr(self, "_scheduler_name") and self._scheduler_name:
+        
+        if hasattr(self, '_scheduler_name') and self._scheduler_name:
             logger.info(f"Scheduler:               {self._scheduler_name}")
             if self._scheduler_params:
                 for key, value in self._scheduler_params.items():
                     logger.info(f"  {key:25s}: {value}")
-
-        if hasattr(self, "_loss_config"):
+        
+        if hasattr(self, '_loss_config'):
             logger.info(f"Loss Function:           {self._loss_config}")
-
+        
         logger.info("Regularization:")
         logger.info(f"  Embedding L1:          {self._embedding_l1_reg}")
         logger.info(f"  Embedding L2:          {self._embedding_l2_reg}")
         logger.info(f"  Dense L1:              {self._dense_l1_reg}")
         logger.info(f"  Dense L2:              {self._dense_l2_reg}")
-
+        
         logger.info("Other Settings:")
         logger.info(f"  Early Stop Patience:   {self.early_stop_patience}")
         logger.info(f"  Max Gradient Norm:     {self._max_gradient_norm}")
         logger.info(f"  Model ID:              {self.model_id}")
         logger.info(f"  Checkpoint Path:       {self.checkpoint}")
-
+        
         logger.info("")
         logger.info("")
 
@@ -1413,47 +1127,45 @@ class BaseMatchModel(BaseModel):
     Base class for match (retrieval/recall) models
     Supports pointwise, pairwise, and listwise training modes
     """
-
+    
     @property
     def task_type(self) -> str:
-        return "match"
-
+        return 'match'
+    
     @property
     def support_training_modes(self) -> list[str]:
         """
         Returns list of supported training modes for this model.
         Override in subclasses to restrict training modes.
-
+        
         Returns:
             List of supported modes: ['pointwise', 'pairwise', 'listwise']
         """
-        return ["pointwise", "pairwise", "listwise"]
-
-    def __init__(
-        self,
-        user_dense_features: list[DenseFeature] | None = None,
-        user_sparse_features: list[SparseFeature] | None = None,
-        user_sequence_features: list[SequenceFeature] | None = None,
-        item_dense_features: list[DenseFeature] | None = None,
-        item_sparse_features: list[SparseFeature] | None = None,
-        item_sequence_features: list[SequenceFeature] | None = None,
-        training_mode: Literal["pointwise", "pairwise", "listwise"] = "pointwise",
-        num_negative_samples: int = 4,
-        temperature: float = 1.0,
-        similarity_metric: Literal["dot", "cosine", "euclidean"] = "dot",
-        device: str = "cpu",
-        embedding_l1_reg: float = 0.0,
-        dense_l1_reg: float = 0.0,
-        embedding_l2_reg: float = 0.0,
-        dense_l2_reg: float = 0.0,
-        early_stop_patience: int = 20,
-        model_id: str = "baseline",
-    ):
-
+        return ['pointwise', 'pairwise', 'listwise']
+    
+    def __init__(self,
+                 user_dense_features: list[DenseFeature] | None = None,
+                 user_sparse_features: list[SparseFeature] | None = None,
+                 user_sequence_features: list[SequenceFeature] | None = None,
+                 item_dense_features: list[DenseFeature] | None = None,
+                 item_sparse_features: list[SparseFeature] | None = None,
+                 item_sequence_features: list[SequenceFeature] | None = None,
+                 training_mode: Literal['pointwise', 'pairwise', 'listwise'] = 'pointwise',
+                 num_negative_samples: int = 4,
+                 temperature: float = 1.0,
+                 similarity_metric: Literal['dot', 'cosine', 'euclidean'] = 'dot',
+                 device: str = 'cpu',
+                 embedding_l1_reg: float = 0.0,
+                 dense_l1_reg: float = 0.0,
+                 embedding_l2_reg: float = 0.0,
+                 dense_l2_reg: float = 0.0,
+                 early_stop_patience: int = 20,
+                 model_id: str = 'baseline'):
+        
         all_dense_features = []
         all_sparse_features = []
         all_sequence_features = []
-
+        
         if user_dense_features:
             all_dense_features.extend(user_dense_features)
         if item_dense_features:
@@ -1466,158 +1178,117 @@ class BaseMatchModel(BaseModel):
             all_sequence_features.extend(user_sequence_features)
         if item_sequence_features:
             all_sequence_features.extend(item_sequence_features)
-
+        
         super(BaseMatchModel, self).__init__(
             dense_features=all_dense_features,
             sparse_features=all_sparse_features,
             sequence_features=all_sequence_features,
-            target=["label"],
-            task="binary",
+            target=['label'],  
+            task='binary',  
             device=device,
             embedding_l1_reg=embedding_l1_reg,
             dense_l1_reg=dense_l1_reg,
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=early_stop_patience,
-            model_id=model_id,
-        )
-
-        self.user_dense_features = (
-            list(user_dense_features) if user_dense_features else []
-        )
-        self.user_sparse_features = (
-            list(user_sparse_features) if user_sparse_features else []
-        )
-        self.user_sequence_features = (
-            list(user_sequence_features) if user_sequence_features else []
-        )
-
-        self.item_dense_features = (
-            list(item_dense_features) if item_dense_features else []
+            model_id=model_id
         )
-        self.item_sparse_features = (
-            list(item_sparse_features) if item_sparse_features else []
-        )
-        self.item_sequence_features = (
-            list(item_sequence_features) if item_sequence_features else []
-        )
-
+        
+        self.user_dense_features = list(user_dense_features) if user_dense_features else []
+        self.user_sparse_features = list(user_sparse_features) if user_sparse_features else []
+        self.user_sequence_features = list(user_sequence_features) if user_sequence_features else []
+        
+        self.item_dense_features = list(item_dense_features) if item_dense_features else []
+        self.item_sparse_features = list(item_sparse_features) if item_sparse_features else []
+        self.item_sequence_features = list(item_sequence_features) if item_sequence_features else []
+        
         self.training_mode = training_mode
         self.num_negative_samples = num_negative_samples
         self.temperature = temperature
         self.similarity_metric = similarity_metric
-
+    
     def get_user_features(self, X_input: dict) -> dict:
         user_input = {}
-        all_user_features = (
-            self.user_dense_features
-            + self.user_sparse_features
-            + self.user_sequence_features
-        )
+        all_user_features = self.user_dense_features + self.user_sparse_features + self.user_sequence_features
         for feature in all_user_features:
             if feature.name in X_input:
                 user_input[feature.name] = X_input[feature.name]
         return user_input
-
+    
     def get_item_features(self, X_input: dict) -> dict:
         item_input = {}
-        all_item_features = (
-            self.item_dense_features
-            + self.item_sparse_features
-            + self.item_sequence_features
-        )
+        all_item_features = self.item_dense_features + self.item_sparse_features + self.item_sequence_features
         for feature in all_item_features:
             if feature.name in X_input:
                 item_input[feature.name] = X_input[feature.name]
         return item_input
-
-    def compile(
-        self,
-        optimizer="adam",
-        optimizer_params: dict | None = None,
-        scheduler: (
-            str
-            | torch.optim.lr_scheduler._LRScheduler
-            | type[torch.optim.lr_scheduler._LRScheduler]
-            | None
-        ) = None,
-        scheduler_params: dict | None = None,
-        loss: str | nn.Module | list[str | nn.Module] | None = None,
-    ):
+    
+    def compile(self, 
+                optimizer = "adam",
+                optimizer_params: dict | None = None,
+                scheduler: str | torch.optim.lr_scheduler._LRScheduler | type[torch.optim.lr_scheduler._LRScheduler] | None = None,
+                scheduler_params: dict | None = None,
+                loss: str | nn.Module | list[str | nn.Module] | None= None):
         """
         Compile match model with optimizer, scheduler, and loss function.
         Validates that training_mode is supported by the model.
         """
         from nextrec.loss import validate_training_mode
-
+        
         # Validate training mode is supported
         validate_training_mode(
             training_mode=self.training_mode,
             support_training_modes=self.support_training_modes,
-            model_name=self.model_name,
+            model_name=self.model_name
         )
-
+        
         # Call parent compile with match-specific logic
         if optimizer_params is None:
             optimizer_params = {}
-
-        self._optimizer_name = (
-            optimizer if isinstance(optimizer, str) else optimizer.__class__.__name__
-        )
+        
+        self._optimizer_name = optimizer if isinstance(optimizer, str) else optimizer.__class__.__name__
         self._optimizer_params = optimizer_params
         if isinstance(scheduler, str):
             self._scheduler_name = scheduler
         elif scheduler is not None:
             # Try to get __name__ first (for class types), then __class__.__name__ (for instances)
-            self._scheduler_name = getattr(
-                scheduler,
-                "__name__",
-                getattr(scheduler.__class__, "__name__", str(scheduler)),
-            )
+            self._scheduler_name = getattr(scheduler, '__name__', getattr(scheduler.__class__, '__name__', str(scheduler)))
         else:
             self._scheduler_name = None
         self._scheduler_params = scheduler_params or {}
         self._loss_config = loss
-
+        
         # set optimizer
         self.optimizer_fn = get_optimizer_fn(
-            optimizer=optimizer, params=self.parameters(), **optimizer_params
+            optimizer=optimizer, 
+            params=self.parameters(), 
+            **optimizer_params
         )
-
+        
         # Set loss function based on training mode
         loss_value = loss[0] if isinstance(loss, list) else loss
-        self.loss_fn = [
-            get_loss_fn(
-                task_type="match", training_mode=self.training_mode, loss=loss_value
-            )
-        ]
-
+        self.loss_fn = [get_loss_fn(
+            task_type='match',
+            training_mode=self.training_mode,
+            loss=loss_value
+        )]
+        
         # set scheduler
-        self.scheduler_fn = (
-            get_scheduler_fn(scheduler, self.optimizer_fn, **(scheduler_params or {}))
-            if scheduler
-            else None
-        )
+        self.scheduler_fn = get_scheduler_fn(scheduler, self.optimizer_fn, **(scheduler_params or {})) if scheduler else None
 
-    def compute_similarity(
-        self, user_emb: torch.Tensor, item_emb: torch.Tensor
-    ) -> torch.Tensor:
-        if self.similarity_metric == "dot":
+    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]
+                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]
+                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]
-
-        elif self.similarity_metric == "cosine":
+        
+        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:
@@ -1625,8 +1296,8 @@ class BaseMatchModel(BaseModel):
                 similarity = F.cosine_similarity(user_emb_expanded, item_emb, dim=-1)
             else:
                 similarity = F.cosine_similarity(user_emb, item_emb, dim=-1)
-
-        elif self.similarity_metric == "euclidean":
+        
+        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:
@@ -1634,96 +1305,84 @@ class BaseMatchModel(BaseModel):
                 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 = -distance 
+        
         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:
         raise NotImplementedError
-
+    
     def item_tower(self, item_input: dict) -> torch.Tensor:
         raise NotImplementedError
-
-    def forward(
-        self, X_input: dict
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    
+    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)
-
-        user_emb = self.user_tower(user_input)  # [B, D]
-        item_emb = self.item_tower(item_input)  # [B, D]
-
-        if self.training and self.training_mode in ["pairwise", "listwise"]:
+        
+        user_emb = self.user_tower(user_input)   # [B, D]
+        item_emb = self.item_tower(item_input)   # [B, D]
+        
+        if self.training and self.training_mode in ['pairwise', 'listwise']:
             return user_emb, item_emb
 
         similarity = self.compute_similarity(user_emb, item_emb)  # [B]
-
-        if self.training_mode == "pointwise":
+        
+        if self.training_mode == 'pointwise':
             return torch.sigmoid(similarity)
         else:
             return similarity
-
+    
     def compute_loss(self, y_pred, y_true):
-        if self.training_mode == "pointwise":
+        if self.training_mode == 'pointwise':
             if y_true is None:
                 return torch.tensor(0.0, device=self.device)
             return self.loss_fn[0](y_pred, y_true)
-
+        
         # pairwise / listwise using inbatch neg
-        elif self.training_mode in ["pairwise", "listwise"]:
+        elif self.training_mode in ['pairwise', 'listwise']:
             if not isinstance(y_pred, (tuple, list)) or len(y_pred) != 2:
                 raise ValueError(
                     "For pairwise/listwise training, forward should return (user_emb, item_emb). "
                     "Please check BaseMatchModel.forward implementation."
                 )
-
+            
             user_emb, item_emb = y_pred  # [B, D], [B, D]
-
+            
             logits = torch.matmul(user_emb, item_emb.t())  # [B, B]
-            logits = logits / self.temperature
-
+            logits = logits / self.temperature            
+            
             batch_size = logits.size(0)
-            targets = torch.arange(
-                batch_size, device=logits.device
-            )  # [0, 1, 2, ..., B-1]
-
+            targets = torch.arange(batch_size, device=logits.device)  # [0, 1, 2, ..., B-1]
+            
             # Cross-Entropy = InfoNCE
             loss = F.cross_entropy(logits, targets)
             return loss
-
+        
         else:
             raise ValueError(f"Unknown training mode: {self.training_mode}")
-
+    
     def _set_metrics(self, metrics: list[str] | None = None):
         if metrics is not None and len(metrics) > 0:
             self.metrics = [m.lower() for m in metrics]
         else:
-            self.metrics = ["auc", "logloss"]
-
-        self.best_metrics_mode = "max"
-
-        if not hasattr(self, "early_stopper") or self.early_stopper is None:
-            self.early_stopper = EarlyStopper(
-                patience=self.early_stop_patience, mode=self.best_metrics_mode
-            )
-
-    def encode_user(
-        self, data: dict | pd.DataFrame | DataLoader, batch_size: int = 512
-    ) -> np.ndarray:
+            self.metrics = ['auc', 'logloss']
+        
+        self.best_metrics_mode = 'max'  
+        
+        if not hasattr(self, 'early_stopper') or self.early_stopper is None:
+            self.early_stopper = EarlyStopper(patience=self.early_stop_patience, mode=self.best_metrics_mode)
+    
+    def encode_user(self, data: dict | pd.DataFrame | DataLoader, batch_size: int = 512) -> np.ndarray:
         self.eval()
-
+        
         if not isinstance(data, DataLoader):
             user_data = {}
-            all_user_features = (
-                self.user_dense_features
-                + self.user_sparse_features
-                + self.user_sequence_features
-            )
+            all_user_features = self.user_dense_features + self.user_sparse_features + self.user_sequence_features
             for feature in all_user_features:
                 if isinstance(data, dict):
                     if feature.name in data:
@@ -1731,39 +1390,29 @@ class BaseMatchModel(BaseModel):
                 elif isinstance(data, pd.DataFrame):
                     if feature.name in data.columns:
                         user_data[feature.name] = data[feature.name].values
-
-            data_loader = self._prepare_data_loader(
-                user_data, batch_size=batch_size, shuffle=False
-            )
+            
+            data_loader = self._prepare_data_loader(user_data, batch_size=batch_size, shuffle=False)
         else:
             data_loader = data
-
+        
         embeddings_list = []
-
+        
         with torch.no_grad():
-            for batch_data in tqdm.tqdm(
-                data_loader, desc="Encoding users", disable=self._verbose == 0
-            ):
+            for batch_data in tqdm.tqdm(data_loader, desc="Encoding users", disable=self._verbose == 0):
                 batch_dict = self._batch_to_dict(batch_data)
                 user_input = self.get_user_features(batch_dict)
                 user_emb = self.user_tower(user_input)
                 embeddings_list.append(user_emb.cpu().numpy())
-
+        
         embeddings = np.concatenate(embeddings_list, axis=0)
         return embeddings
-
-    def encode_item(
-        self, data: dict | pd.DataFrame | DataLoader, batch_size: int = 512
-    ) -> np.ndarray:
+    
+    def encode_item(self, data: dict | pd.DataFrame | DataLoader, batch_size: int = 512) -> np.ndarray:
         self.eval()
-
+        
         if not isinstance(data, DataLoader):
             item_data = {}
-            all_item_features = (
-                self.item_dense_features
-                + self.item_sparse_features
-                + self.item_sequence_features
-            )
+            all_item_features = self.item_dense_features + self.item_sparse_features + self.item_sequence_features
             for feature in all_item_features:
                 if isinstance(data, dict):
                     if feature.name in data:
@@ -1772,22 +1421,18 @@ class BaseMatchModel(BaseModel):
                     if feature.name in data.columns:
                         item_data[feature.name] = data[feature.name].values
 
-            data_loader = self._prepare_data_loader(
-                item_data, batch_size=batch_size, shuffle=False
-            )
+            data_loader = self._prepare_data_loader(item_data, batch_size=batch_size, shuffle=False)
         else:
             data_loader = data
-
+        
         embeddings_list = []
-
+        
         with torch.no_grad():
-            for batch_data in tqdm.tqdm(
-                data_loader, desc="Encoding items", disable=self._verbose == 0
-            ):
+            for batch_data in tqdm.tqdm(data_loader, desc="Encoding items", disable=self._verbose == 0):
                 batch_dict = self._batch_to_dict(batch_data)
                 item_input = self.get_item_features(batch_dict)
                 item_emb = self.item_tower(item_input)
                 embeddings_list.append(item_emb.cpu().numpy())
-
+        
         embeddings = np.concatenate(embeddings_list, axis=0)
         return embeddings