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

nextrec/__version__.py

@@ -1 +1 @@
-__version__ = "0.4.22"
+__version__ = "0.4.23"

nextrec/basic/metrics.py

@@ -2,7 +2,7 @@
 Metrics computation and configuration for model evaluation.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 20/12/2025
+Checkpoint: edit on 29/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 """
 
@@ -39,7 +39,6 @@ REGRESSION_METRICS = {"mse", "mae", "rmse", "r2", "mape", "msle"}
 TASK_DEFAULT_METRICS = {
     "binary": ["auc", "gauc", "ks", "logloss", "accuracy", "precision", "recall", "f1"],
     "regression": ["mse", "mae", "rmse", "r2", "mape"],
-    "multilabel": ["auc", "hamming_loss", "subset_accuracy", "micro_f1", "macro_f1"],
     "matching": ["auc", "gauc", "precision@10", "hitrate@10", "map@10", "cosine"]
     + [f"recall@{k}" for k in (5, 10, 20)]
     + [f"ndcg@{k}" for k in (5, 10, 20)]

nextrec/basic/model.py

@@ -2,13 +2,14 @@
 Base Model & Base Match Model Class
 
 Date: create on 27/10/2025
-Checkpoint: edit on 28/12/2025
+Checkpoint: edit on 29/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 """
 
 import getpass
 import logging
 import os
+import sys
 import pickle
 import socket
 from pathlib import Path
@@ -16,6 +17,16 @@ from typing import Any, Literal
 
 import numpy as np
 import pandas as pd
+
+try:
+    import swanlab  # type: ignore
+except ModuleNotFoundError:
+    swanlab = None
+try:
+    import wandb  # type: ignore
+except ModuleNotFoundError:
+    wandb = None
+
 import torch
 import torch.distributed as dist
 import torch.nn as nn
@@ -74,13 +85,19 @@ from nextrec.utils.torch_utils import (
     to_tensor,
 )
 from nextrec.utils.config import safe_value
-from nextrec.utils.model import compute_ranking_loss
+from nextrec.utils.model import (
+    compute_ranking_loss,
+    get_loss_list,
+    resolve_loss_weights,
+    get_training_modes,
+)
 from nextrec.utils.types import (
     LossName,
     OptimizerName,
     SchedulerName,
     TrainingModeName,
     TaskTypeName,
+    MetricsName,
 )
 
 
@@ -90,7 +107,7 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
         raise NotImplementedError
 
     @property
-    def default_task(self) -> str | list[str]:
+    def default_task(self) -> TaskTypeName | list[TaskTypeName]:
         raise NotImplementedError
 
     def __init__(
@@ -139,6 +156,9 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             world_size: Number of processes (defaults to env WORLD_SIZE).
             local_rank: Local rank for selecting CUDA device (defaults to env LOCAL_RANK).
             ddp_find_unused_parameters: Default False, set it True only when exist unused parameters in ddp model, in most cases should be False.
+
+        Note:
+            Optimizer, scheduler, and loss are configured via compile().
         """
         super(BaseModel, self).__init__()
 
@@ -171,24 +191,12 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             dense_features, sparse_features, sequence_features, target, id_columns
         )
 
-        self.task = self.default_task if task is None else task
+        self.task = task or self.default_task
         self.nums_task = len(self.task) if isinstance(self.task, list) else 1
-        if isinstance(training_mode, list):
-            training_modes = list(training_mode)
-            if len(training_modes) != self.nums_task:
-                raise ValueError(
-                    "[BaseModel-init Error] training_mode list length must match number of tasks."
-                )
-        else:
-            training_modes = [training_mode] * self.nums_task
-        if any(
-            mode not in {"pointwise", "pairwise", "listwise"} for mode in training_modes
-        ):
-            raise ValueError(
-                "[BaseModel-init Error] training_mode must be one of {'pointwise', 'pairwise', 'listwise'}."
-            )
-        self.training_modes = training_modes
-        self.training_mode = training_modes if self.nums_task > 1 else training_modes[0]
+        self.training_modes = get_training_modes(training_mode, self.nums_task)
+        self.training_mode = (
+            self.training_modes if self.nums_task > 1 else self.training_modes[0]
+        )
 
         self.embedding_l1_reg = embedding_l1_reg
         self.dense_l1_reg = dense_l1_reg
@@ -196,8 +204,9 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
         self.dense_l2_reg = dense_l2_reg
         self.regularization_weights = []
         self.embedding_params = []
-        self.loss_weight = None
+
         self.ignore_label = None
+        self.compiled = False
 
         self.max_gradient_norm = 1.0
         self.logger_initialized = False
@@ -431,28 +440,9 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             "pairwise": "bpr",
             "listwise": "listnet",
         }
-        effective_loss = loss
-        if effective_loss is None:
-            loss_list = [default_losses[mode] for mode in self.training_modes]
-        elif isinstance(effective_loss, list):
-            if not effective_loss:
-                loss_list = [default_losses[mode] for mode in self.training_modes]
-            else:
-                if len(effective_loss) != self.nums_task:
-                    raise ValueError(
-                        f"[BaseModel-compile Error] Number of loss functions ({len(effective_loss)}) must match number of tasks ({self.nums_task})."
-                    )
-                loss_list = list(effective_loss)
-        else:
-            loss_list = [effective_loss] * self.nums_task
-
-        for idx, mode in enumerate(self.training_modes):
-            if isinstance(loss_list[idx], str) and loss_list[idx] in {
-                "bce",
-                "binary_crossentropy",
-            }:
-                if mode in {"pairwise", "listwise"}:
-                    loss_list[idx] = default_losses[mode]
+        loss_list = get_loss_list(
+            loss, self.training_modes, self.nums_task, default_losses
+        )
         self.loss_params = loss_params or {}
         optimizer_params = optimizer_params or {}
         self.optimizer_name = (
@@ -516,30 +506,9 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
                 nums_task=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)):
-                if len(loss_weights) != 1:
-                    raise ValueError(
-                        "[BaseModel-compile Error] loss_weights list must have exactly one element for single-task setup."
-                    )
-                loss_weights = loss_weights[0]
-            self.loss_weights = [float(loss_weights)]  # type: ignore
         else:
-            if isinstance(loss_weights, (int, float)):
-                weights = [float(loss_weights)] * self.nums_task
-            elif isinstance(loss_weights, (list, tuple)):
-                weights = [float(w) for w in loss_weights]
-                if len(weights) != self.nums_task:
-                    raise ValueError(
-                        f"[BaseModel-compile Error] Number of loss_weights ({len(weights)}) must match number of tasks ({self.nums_task})."
-                    )
-            else:
-                raise TypeError(
-                    f"[BaseModel-compile Error] loss_weights must be int, float, list or tuple, got {type(loss_weights)}"
-                )
-            self.loss_weights = weights
+            self.loss_weights = resolve_loss_weights(loss_weights, self.nums_task)
+        self.compiled = True
 
     def compute_loss(self, y_pred, y_true):
         if y_true is None:
@@ -602,9 +571,6 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
         for i, (start, end) in enumerate(slices):  # type: ignore
             y_pred_i = y_pred[:, start:end]
             y_true_i = y_true[:, start:end]
-            total_count = y_true_i.shape[0]
-            # valid_count = None
-
             # mask ignored labels
             if self.ignore_label is not None:
                 valid_mask = y_true_i != self.ignore_label
@@ -613,11 +579,8 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
                 if not torch.any(valid_mask):
                     task_losses.append(y_pred_i.sum() * 0.0)
                     continue
-                # valid_count = valid_mask.sum().to(dtype=y_true_i.dtype)
                 y_pred_i = y_pred_i[valid_mask]
                 y_true_i = y_true_i[valid_mask]
-            # else:
-            # valid_count = y_true_i.new_tensor(float(total_count))
 
             mode = self.training_modes[i]
 
@@ -691,7 +654,7 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
         train_data=None,
         valid_data=None,
         metrics: (
-            list[str] | dict[str, list[str]] | None
+            list[MetricsName] | dict[str, list[MetricsName]] | None
         ) = None,  # ['auc', 'logloss'] or {'target1': ['auc', 'logloss'], 'target2': ['mse']}
         epochs: int = 1,
         shuffle: bool = True,
@@ -705,6 +668,8 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
         use_tensorboard: bool = True,
         use_wandb: bool = False,
         use_swanlab: bool = False,
+        wandb_api: str | None = None,
+        swanlab_api: str | None = None,
         wandb_kwargs: dict | None = None,
         swanlab_kwargs: dict | None = None,
         auto_ddp_sampler: bool = True,
@@ -734,6 +699,8 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             use_tensorboard: Enable tensorboard logging.
             use_wandb: Enable Weights & Biases logging.
             use_swanlab: Enable SwanLab logging.
+            wandb_api: W&B API key for non-tty login.
+            swanlab_api: SwanLab API key for non-tty login.
             wandb_kwargs: Optional kwargs for wandb.init(...).
             swanlab_kwargs: Optional kwargs for swanlab.init(...).
             auto_ddp_sampler: Attach DistributedSampler automatically when distributed, set False to when data is already sharded per rank.
@@ -751,6 +718,16 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
         )
         self.to(self.device)
 
+        if not self.compiled:
+            self.compile(
+                optimizer="adam",
+                optimizer_params={},
+                scheduler=None,
+                scheduler_params={},
+                loss=None,
+                loss_params={},
+            )
+
         if (
             self.distributed
             and dist.is_available()
@@ -825,6 +802,24 @@ class BaseModel(SummarySet, FeatureSet, nn.Module):
             }
             training_config: dict = safe_value(training_config)  # type: ignore
 
+        if self.is_main_process:
+            is_tty = sys.stdin.isatty() and sys.stdout.isatty()
+            if not is_tty:
+                if use_wandb and wandb_api:
+                    if wandb is None:
+                        logging.warning(
+                            "[BaseModel-fit] wandb not installed, skip wandb login."
+                        )
+                    else:
+                        wandb.login(key=wandb_api)
+                if use_swanlab and swanlab_api:
+                    if swanlab is None:
+                        logging.warning(
+                            "[BaseModel-fit] swanlab not installed, skip swanlab login."
+                        )
+                    else:
+                        swanlab.login(api_key=swanlab_api)
+
         self.training_logger = (
             TrainingLogger(
                 session=self.session,

nextrec/basic/summary.py

@@ -1,5 +1,9 @@
 """
 Summary utilities for BaseModel.
+
+Date: create on 03/12/2025
+Checkpoint: edit on 29/12/2025
+Author: Yang Zhou,zyaztec@gmail.com
 """
 
 from __future__ import annotations
@@ -12,9 +16,39 @@ from torch.utils.data import DataLoader
 
 from nextrec.basic.loggers import colorize, format_kv
 from nextrec.data.data_processing import extract_label_arrays, get_data_length
+from nextrec.utils.types import TaskTypeName
 
 
 class SummarySet:
+    model_name: str
+    dense_features: list[Any]
+    sparse_features: list[Any]
+    sequence_features: list[Any]
+    task: TaskTypeName | list[TaskTypeName]
+    target_columns: list[str]
+    nums_task: int
+    metrics: Any
+    device: Any
+    optimizer_name: str
+    optimizer_params: dict[str, Any]
+    scheduler_name: str | None
+    scheduler_params: dict[str, Any]
+    loss_config: Any
+    loss_weights: Any
+    grad_norm: Any
+    embedding_l1_reg: float
+    embedding_l2_reg: float
+    dense_l1_reg: float
+    dense_l2_reg: float
+    early_stop_patience: int
+    max_gradient_norm: float | None
+    metrics_sample_limit: int | None
+    session_id: str | None
+    features_config_path: str
+    checkpoint_path: str
+    train_data_summary: dict[str, Any] | None
+    valid_data_summary: dict[str, Any] | None
+
     def build_data_summary(
         self, data: Any, data_loader: DataLoader | None, sample_key: str
     ):
@@ -305,7 +339,7 @@ class SummarySet:
                         lines = details.get("lines", [])
                         logger.info(f"{target_name}:")
                         for label, value in lines:
-                            logger.info(format_kv(label, value))
+                            logger.info(f"  {format_kv(label, value)}")
 
             if self.valid_data_summary:
                 if self.train_data_summary:
@@ -320,4 +354,4 @@ class SummarySet:
                         lines = details.get("lines", [])
                         logger.info(f"{target_name}:")
                         for label, value in lines:
-                            logger.info(format_kv(label, value))
+                            logger.info(f"  {format_kv(label, value)}")

nextrec/data/preprocessor.py

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

nextrec/loss/listwise.py

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