nextrec 0.4.15__py3-none-any.whl → 0.4.17__py3-none-any.whl

nextrec/__version__.py

@@ -1 +1 @@
-__version__ = "0.4.15"
+__version__ = "0.4.17"

nextrec/basic/heads.py

@@ -0,0 +1,101 @@
+"""
+Task head implementations for NextRec models.
+
+Date: create on 23/12/2025
+Author: Yang Zhou, zyaztec@gmail.com
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from nextrec.basic.layers import PredictionLayer
+
+
+class TaskHead(nn.Module):
+    """
+    Unified task head for ranking/regression/multi-task outputs.
+
+    This wraps PredictionLayer so models can depend on a "Head" abstraction
+    without changing their existing forward signatures.
+    """
+
+    def __init__(
+        self,
+        task_type: str | list[str] = "binary",
+        task_dims: int | list[int] | None = None,
+        use_bias: bool = True,
+        return_logits: bool = False,
+    ) -> None:
+        super().__init__()
+        self.prediction = PredictionLayer(
+            task_type=task_type,
+            task_dims=task_dims,
+            use_bias=use_bias,
+            return_logits=return_logits,
+        )
+        # Expose commonly used attributes for compatibility with PredictionLayer.
+        self.task_types = self.prediction.task_types
+        self.task_dims = self.prediction.task_dims
+        self.task_slices = self.prediction.task_slices
+        self.total_dim = self.prediction.total_dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.prediction(x)
+
+
+class RetrievalHead(nn.Module):
+    """
+    Retrieval head for two-tower models.
+
+    It computes similarity for pointwise training/inference, and returns
+    raw embeddings for in-batch negative sampling in pairwise/listwise modes.
+    """
+
+    def __init__(
+        self,
+        similarity_metric: Literal["dot", "cosine", "euclidean"] = "dot",
+        temperature: float = 1.0,
+        training_mode: Literal["pointwise", "pairwise", "listwise"] = "pointwise",
+        apply_sigmoid: bool = True,
+    ) -> None:
+        super().__init__()
+        self.similarity_metric = similarity_metric
+        self.temperature = temperature
+        self.training_mode = training_mode
+        self.apply_sigmoid = apply_sigmoid
+
+    def forward(
+        self,
+        user_emb: torch.Tensor,
+        item_emb: torch.Tensor,
+        similarity_fn=None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if self.training and self.training_mode in {"pairwise", "listwise"}:
+            return user_emb, item_emb
+
+        if similarity_fn is not None:
+            similarity = similarity_fn(user_emb, item_emb)
+        else:
+            if user_emb.dim() == 2 and item_emb.dim() == 3:
+                user_emb = user_emb.unsqueeze(1)
+
+            if self.similarity_metric == "dot":
+                similarity = torch.sum(user_emb * item_emb, dim=-1)
+            elif self.similarity_metric == "cosine":
+                similarity = F.cosine_similarity(user_emb, item_emb, dim=-1)
+            elif self.similarity_metric == "euclidean":
+                similarity = -torch.sum((user_emb - item_emb) ** 2, dim=-1)
+            else:
+                raise ValueError(
+                    f"Unknown similarity metric: {self.similarity_metric}"
+                )
+
+            similarity = similarity / self.temperature
+        if self.training_mode == "pointwise" and self.apply_sigmoid:
+            return torch.sigmoid(similarity)
+        return similarity

nextrec/basic/metrics.py

@@ -77,6 +77,8 @@ def check_user_id(*metric_sources: Any) -> bool:
 
 def compute_ks(y_true: np.ndarray, y_pred: np.ndarray) -> float:
     """Compute Kolmogorov-Smirnov statistic."""
+    y_true = np.asarray(y_true).reshape(-1)
+    y_pred = np.asarray(y_pred).reshape(-1)
     sorted_indices = np.argsort(y_pred)[::-1]
     y_true_sorted = y_true[sorted_indices]
 

nextrec/basic/model.py

@@ -38,6 +38,7 @@ from nextrec.basic.features import (
     SequenceFeature,
     SparseFeature,
 )
+from nextrec.basic.heads import RetrievalHead
 from nextrec.basic.loggers import TrainingLogger, colorize, format_kv, setup_logger
 from nextrec.basic.metrics import check_user_id, configure_metrics, evaluate_metrics
 from nextrec.basic.session import create_session, resolve_save_path
@@ -481,7 +482,7 @@ class BaseModel(FeatureSet, nn.Module):
                         "[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
+            self.loss_weights = [float(loss_weights)]  # type: ignore
         else:
             if isinstance(loss_weights, (int, float)):
                 weights = [float(loss_weights)] * self.nums_task
@@ -591,8 +592,8 @@ class BaseModel(FeatureSet, nn.Module):
 
     def fit(
         self,
-        train_data = None,
-        valid_data = None,
+        train_data=None,
+        valid_data=None,
         metrics: (
             list[str] | dict[str, list[str]] | None
         ) = None,  # ['auc', 'logloss'] or {'target1': ['auc', 'logloss'], 'target2': ['mse']}
@@ -1583,8 +1584,11 @@ class BaseModel(FeatureSet, nn.Module):
         else:
             data_loader = data
 
-            if hasattr(data_loader, 'num_workers') and data_loader.num_workers > 0:
-                if hasattr(data_loader.dataset, '__class__') and 'Streaming' in data_loader.dataset.__class__.__name__:
+            if hasattr(data_loader, "num_workers") and data_loader.num_workers > 0:
+                if (
+                    hasattr(data_loader.dataset, "__class__")
+                    and "Streaming" in data_loader.dataset.__class__.__name__
+                ):
                     logging.warning(
                         f"[Predict Streaming Warning] Detected DataLoader with num_workers={data_loader.num_workers} "
                         "and streaming dataset. This may cause data duplication! "
@@ -2112,6 +2116,12 @@ class BaseMatchModel(BaseModel):
         )
         self.user_feature_names = {feature.name for feature in self.user_features_all}
         self.item_feature_names = {feature.name for feature in self.item_features_all}
+        self.head = RetrievalHead(
+            similarity_metric=self.similarity_metric,
+            temperature=self.temperature,
+            training_mode=self.training_mode,
+            apply_sigmoid=True,
+        )
 
     def compile(
         self,
@@ -2241,15 +2251,9 @@ class BaseMatchModel(BaseModel):
         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":
-            return torch.sigmoid(similarity)
-        else:
-            return similarity
+        return self.head(
+            user_emb, item_emb, similarity_fn=self.compute_similarity
+        )
 
     def compute_loss(self, y_pred, y_true):
         if self.training_mode == "pointwise":

nextrec/models/multi_task/esmm.py

@@ -45,7 +45,8 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -139,7 +140,7 @@ class ESMM(BaseModel):
         # CVR tower
         self.cvr_tower = MLP(input_dim=input_dim, output_layer=True, **cvr_params)
         self.grad_norm_shared_modules = ["embedding"]
-        self.prediction_layer = PredictionLayer(
+        self.prediction_layer = TaskHead(
             task_type=self.default_task, task_dims=[1, 1]
         )
         # Register regularization weights
@@ -167,4 +168,4 @@ class ESMM(BaseModel):
 
         # Output: [CTR, CTCVR], We supervise CTR with click labels and CTCVR with conversion labels
         y = torch.cat([ctr, ctcvr], dim=1)  # [B, 2]
-        return y  # [B, 2], where y[:, 0] is CTR and y[:, 1] is CTCVR
+        return y  # [B, 2], where y[:, 0] is CTR and y[:, 1] is CTCVR

nextrec/models/multi_task/mmoe.py

@@ -46,7 +46,8 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -172,7 +173,7 @@ class MMOE(BaseModel):
         for tower_params in tower_params_list:
             tower = MLP(input_dim=expert_output_dim, output_layer=True, **tower_params)
             self.towers.append(tower)
-        self.prediction_layer = PredictionLayer(
+        self.prediction_layer = TaskHead(
             task_type=self.default_task, task_dims=[1] * self.num_tasks
         )
         # Register regularization weights
@@ -219,4 +220,4 @@ class MMOE(BaseModel):
 
         # Stack outputs: [B, num_tasks]
         y = torch.cat(task_outputs, dim=1)
-        return self.prediction_layer(y)
+        return self.prediction_layer(y)

nextrec/models/multi_task/ple.py

@@ -49,7 +49,8 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 from nextrec.utils.model import get_mlp_output_dim
 
@@ -302,7 +303,7 @@ class PLE(BaseModel):
         for tower_params in tower_params_list:
             tower = MLP(input_dim=expert_output_dim, output_layer=True, **tower_params)
             self.towers.append(tower)
-        self.prediction_layer = PredictionLayer(
+        self.prediction_layer = TaskHead(
             task_type=self.default_task, task_dims=[1] * self.num_tasks
         )
         # Register regularization weights
@@ -336,4 +337,4 @@ class PLE(BaseModel):
 
         # [B, num_tasks]
         y = torch.cat(task_outputs, dim=1)
-        return self.prediction_layer(y)
+        return self.prediction_layer(y)

nextrec/models/multi_task/poso.py

@@ -44,7 +44,8 @@ import torch.nn.functional as F
 
 from nextrec.basic.activation import activation_layer
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 from nextrec.utils.model import select_features
 
@@ -487,7 +488,7 @@ class POSO(BaseModel):
             self.grad_norm_shared_modules = ["embedding"]
         else:
             self.grad_norm_shared_modules = ["embedding", "mmoe"]
-        self.prediction_layer = PredictionLayer(
+        self.prediction_layer = TaskHead(
             task_type=self.default_task,
             task_dims=[1] * self.num_tasks,
         )
@@ -524,4 +525,4 @@ class POSO(BaseModel):
                 task_outputs.append(logit)
 
         y = torch.cat(task_outputs, dim=1)
-        return self.prediction_layer(y)
+        return self.prediction_layer(y)

nextrec/models/multi_task/share_bottom.py

@@ -43,7 +43,8 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -142,7 +143,7 @@ class ShareBottom(BaseModel):
         for tower_params in tower_params_list:
             tower = MLP(input_dim=bottom_output_dim, output_layer=True, **tower_params)
             self.towers.append(tower)
-        self.prediction_layer = PredictionLayer(
+        self.prediction_layer = TaskHead(
             task_type=self.default_task, task_dims=[1] * self.num_tasks
         )
         # Register regularization weights
@@ -171,4 +172,4 @@ class ShareBottom(BaseModel):
 
         # Stack outputs: [B, num_tasks]
         y = torch.cat(task_outputs, dim=1)
-        return self.prediction_layer(y)
+        return self.prediction_layer(y)

nextrec/models/ranking/afm.py

@@ -40,7 +40,8 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import EmbeddingLayer, InputMask, PredictionLayer
+from nextrec.basic.layers import EmbeddingLayer, InputMask
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -141,7 +142,7 @@ class AFM(BaseModel):
         self.attention_p = nn.Linear(attention_dim, 1, bias=False)
         self.attention_dropout = nn.Dropout(attention_dropout)
         self.output_projection = nn.Linear(self.embedding_dim, 1, bias=False)
-        self.prediction_layer = PredictionLayer(task_type=self.default_task)
+        self.prediction_layer = TaskHead(task_type=self.default_task)
         self.input_mask = InputMask()
 
         # Register regularization weights
@@ -243,4 +244,4 @@ class AFM(BaseModel):
         y_afm = self.output_projection(weighted_sum)
 
         y = y_linear + y_afm
-        return self.prediction_layer(y)
+        return self.prediction_layer(y)

nextrec/models/ranking/autoint.py

@@ -58,7 +58,8 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import EmbeddingLayer, MultiHeadSelfAttention, PredictionLayer
+from nextrec.basic.layers import EmbeddingLayer, MultiHeadSelfAttention
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -162,7 +163,7 @@ class AutoInt(BaseModel):
 
         # Final prediction layer
         self.fc = nn.Linear(num_fields * att_embedding_dim, 1)
-        self.prediction_layer = PredictionLayer(task_type=self.default_task)
+        self.prediction_layer = TaskHead(task_type=self.default_task)
 
         # Register regularization weights
         self.register_regularization_weights(
@@ -206,4 +207,4 @@ class AutoInt(BaseModel):
             start_dim=1
         )  # [B, num_fields * att_embedding_dim]
         y = self.fc(attention_output_flat)  # [B, 1]
-        return self.prediction_layer(y)
+        return self.prediction_layer(y)

nextrec/models/ranking/dcn.py

@@ -54,7 +54,8 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -163,7 +164,7 @@ class DCN(BaseModel):
             # Final layer only uses cross network output
             self.final_layer = nn.Linear(input_dim, 1)
 
-        self.prediction_layer = PredictionLayer(task_type=self.task)
+        self.prediction_layer = TaskHead(task_type=self.task)
 
         # Register regularization weights
         self.register_regularization_weights(
@@ -197,4 +198,4 @@ class DCN(BaseModel):
 
         # Final prediction
         y = self.final_layer(combined)
-        return self.prediction_layer(y)
+        return self.prediction_layer(y)

nextrec/models/ranking/dcn_v2.py

@@ -47,7 +47,8 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -272,7 +273,7 @@ class DCNv2(BaseModel):
             final_input_dim = input_dim
 
         self.final_layer = nn.Linear(final_input_dim, 1)
-        self.prediction_layer = PredictionLayer(task_type=self.default_task)
+        self.prediction_layer = TaskHead(task_type=self.default_task)
 
         self.register_regularization_weights(
             embedding_attr="embedding",
@@ -301,4 +302,4 @@ class DCNv2(BaseModel):
             combined = cross_out
 
         logit = self.final_layer(combined)
-        return self.prediction_layer(logit)
+        return self.prediction_layer(logit)

nextrec/models/ranking/deepfm.py

@@ -45,7 +45,8 @@ embedding，无需手工构造交叉特征即可端到端训练，常用于 CTR/
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import FM, LR, MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import FM, LR, MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -111,7 +112,7 @@ class DeepFM(BaseModel):
         self.linear = LR(fm_emb_dim_total)
         self.fm = FM(reduce_sum=True)
         self.mlp = MLP(input_dim=mlp_input_dim, **mlp_params)
-        self.prediction_layer = PredictionLayer(task_type=self.default_task)
+        self.prediction_layer = TaskHead(task_type=self.default_task)
 
         # Register regularization weights
         self.register_regularization_weights(
@@ -133,4 +134,4 @@ class DeepFM(BaseModel):
         y_deep = self.mlp(input_deep)  # [B, 1]
 
         y = y_linear + y_fm + y_deep
-        return self.prediction_layer(y)
+        return self.prediction_layer(y)

nextrec/models/ranking/dien.py

@@ -55,8 +55,8 @@ from nextrec.basic.layers import (
     MLP,
     AttentionPoolingLayer,
     EmbeddingLayer,
-    PredictionLayer,
 )
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -346,7 +346,7 @@ class DIEN(BaseModel):
         )
 
         self.mlp = MLP(input_dim=mlp_input_dim, **mlp_params)
-        self.prediction_layer = PredictionLayer(task_type=self.task)
+        self.prediction_layer = TaskHead(task_type=self.task)
 
         self.register_regularization_weights(
             embedding_attr="embedding",

nextrec/models/ranking/din.py

@@ -55,8 +55,8 @@ from nextrec.basic.layers import (
     MLP,
     AttentionPoolingLayer,
     EmbeddingLayer,
-    PredictionLayer,
 )
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -173,7 +173,7 @@ class DIN(BaseModel):
 
         # MLP for final prediction
         self.mlp = MLP(input_dim=mlp_input_dim, **mlp_params)
-        self.prediction_layer = PredictionLayer(task_type=self.task)
+        self.prediction_layer = TaskHead(task_type=self.task)
 
         # Register regularization weights
         self.register_regularization_weights(

nextrec/models/ranking/eulernet.py

@@ -38,7 +38,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import LR, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import LR, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -295,7 +296,7 @@ class EulerNet(BaseModel):
         else:
             self.linear = None
 
-        self.prediction_layer = PredictionLayer(task_type=self.task)
+        self.prediction_layer = TaskHead(task_type=self.task)
 
         modules = ["mapping", "layers", "w", "w_im"]
         if self.use_linear:
@@ -331,4 +332,4 @@ class EulerNet(BaseModel):
             r, p = layer(r, p)
         r_flat = r.reshape(r.size(0), self.num_orders * self.embedding_dim)
         p_flat = p.reshape(p.size(0), self.num_orders * self.embedding_dim)
-        return self.w(r_flat) + self.w_im(p_flat)
+        return self.w(r_flat) + self.w_im(p_flat)

nextrec/models/ranking/ffm.py

@@ -43,7 +43,8 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import AveragePooling, InputMask, PredictionLayer, SumPooling
+from nextrec.basic.layers import AveragePooling, InputMask, SumPooling
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 from nextrec.utils.torch_utils import get_initializer
 
@@ -140,7 +141,7 @@ class FFM(BaseModel):
             nn.Linear(dense_input_dim, 1, bias=True) if dense_input_dim > 0 else None
         )
 
-        self.prediction_layer = PredictionLayer(task_type=self.task)
+        self.prediction_layer = TaskHead(task_type=self.task)
         self.input_mask = InputMask()
         self.mean_pool = AveragePooling()
         self.sum_pool = SumPooling()
@@ -272,4 +273,4 @@ class FFM(BaseModel):
                 )
 
         y = y_linear + y_interaction
-        return self.prediction_layer(y)
+        return self.prediction_layer(y)

nextrec/models/ranking/fibinet.py

@@ -50,9 +50,9 @@ from nextrec.basic.layers import (
     BiLinearInteractionLayer,
     EmbeddingLayer,
     HadamardInteractionLayer,
-    PredictionLayer,
     SENETLayer,
 )
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -168,7 +168,7 @@ class FiBiNET(BaseModel):
         num_pairs = self.num_fields * (self.num_fields - 1) // 2
         interaction_dim = num_pairs * self.embedding_dim * 2
         self.mlp = MLP(input_dim=interaction_dim, **mlp_params)
-        self.prediction_layer = PredictionLayer(task_type=self.default_task)
+        self.prediction_layer = TaskHead(task_type=self.default_task)
 
         # Register regularization weights
         self.register_regularization_weights(

nextrec/models/ranking/fm.py

@@ -42,7 +42,8 @@ import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
 from nextrec.basic.layers import FM as FMInteraction
-from nextrec.basic.layers import LR, EmbeddingLayer, PredictionLayer
+from nextrec.basic.heads import TaskHead
+from nextrec.basic.layers import LR, EmbeddingLayer
 from nextrec.basic.model import BaseModel
 
 
@@ -105,7 +106,7 @@ class FM(BaseModel):
         fm_input_dim = sum([f.embedding_dim for f in self.fm_features])
         self.linear = LR(fm_input_dim)
         self.fm = FMInteraction(reduce_sum=True)
-        self.prediction_layer = PredictionLayer(task_type=self.task)
+        self.prediction_layer = TaskHead(task_type=self.task)
 
         # Register regularization weights
         self.register_regularization_weights(
@@ -124,4 +125,4 @@ class FM(BaseModel):
         y_linear = self.linear(input_fm.flatten(start_dim=1))
         y_fm = self.fm(input_fm)
         y = y_linear + y_fm
-        return self.prediction_layer(y)
+        return self.prediction_layer(y)

nextrec/models/ranking/lr.py

@@ -41,7 +41,8 @@ LR 是 CTR/排序任务中最经典的线性基线模型。它将稠密、稀疏
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import EmbeddingLayer, LR as LinearLayer, PredictionLayer
+from nextrec.basic.layers import EmbeddingLayer, LR as LinearLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -99,7 +100,7 @@ class LR(BaseModel):
         self.embedding = EmbeddingLayer(features=self.all_features)
         linear_input_dim = self.embedding.input_dim
         self.linear = LinearLayer(linear_input_dim)
-        self.prediction_layer = PredictionLayer(task_type=self.task)
+        self.prediction_layer = TaskHead(task_type=self.task)
 
         self.register_regularization_weights(
             embedding_attr="embedding", include_modules=["linear"]
@@ -115,4 +116,4 @@ class LR(BaseModel):
     def forward(self, x):
         input_linear = self.embedding(x=x, features=self.all_features, squeeze_dim=True)
         y = self.linear(input_linear)
-        return self.prediction_layer(y)
+        return self.prediction_layer(y)

nextrec/models/ranking/masknet.py

@@ -58,7 +58,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -282,7 +283,7 @@ class MaskNet(BaseModel):
                 input_dim=self.num_blocks * block_hidden_dim, **mlp_params
             )
             self.output_layer = None
-        self.prediction_layer = PredictionLayer(task_type=self.task)
+        self.prediction_layer = TaskHead(task_type=self.task)
 
         if self.architecture == "serial":
             self.register_regularization_weights(
@@ -323,4 +324,4 @@ class MaskNet(BaseModel):
                 hidden = self.block_dropout(hidden)
             logit = self.output_layer(hidden)  # [B, 1]
         y = self.prediction_layer(logit)
-        return y
+        return y

nextrec/models/ranking/pnn.py

@@ -38,7 +38,8 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -136,7 +137,7 @@ class PNN(BaseModel):
             product_dim = 2 * self.num_pairs
 
         self.mlp = MLP(input_dim=linear_dim + product_dim, **mlp_params)
-        self.prediction_layer = PredictionLayer(task_type=self.task)
+        self.prediction_layer = TaskHead(task_type=self.task)
 
         modules = ["mlp"]
         if self.kernel is not None:
@@ -197,4 +198,4 @@ class PNN(BaseModel):
 
         deep_input = torch.cat([linear_signal, product_signal], dim=1)
         y = self.mlp(deep_input)
-        return self.prediction_layer(y)
+        return self.prediction_layer(y)

nextrec/models/ranking/widedeep.py

@@ -42,7 +42,8 @@ Wide & Deep 同时使用宽线性部分（记忆共现/手工交叉）与深网
 import torch.nn as nn
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import LR, MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import LR, MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -114,7 +115,7 @@ class WideDeep(BaseModel):
         # deep_emb_dim_total = sum([f.embedding_dim for f in self.deep_features if not isinstance(f, DenseFeature)])
         # dense_input_dim = sum([getattr(f, "embedding_dim", 1) or 1 for f in dense_features])
         self.mlp = MLP(input_dim=input_dim, **mlp_params)
-        self.prediction_layer = PredictionLayer(task_type=self.task)
+        self.prediction_layer = TaskHead(task_type=self.task)
         # Register regularization weights
         self.register_regularization_weights(
             embedding_attr="embedding", include_modules=["linear", "mlp"]
@@ -137,4 +138,4 @@ class WideDeep(BaseModel):
 
         # Combine wide and deep
         y = y_wide + y_deep
-        return self.prediction_layer(y)
+        return self.prediction_layer(y)

nextrec/models/ranking/xdeepfm.py

@@ -56,7 +56,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
-from nextrec.basic.layers import LR, MLP, EmbeddingLayer, PredictionLayer
+from nextrec.basic.layers import LR, MLP, EmbeddingLayer
+from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 
 
@@ -186,7 +187,7 @@ class xDeepFM(BaseModel):
             [getattr(f, "embedding_dim", 1) or 1 for f in dense_features]
         )
         self.mlp = MLP(input_dim=deep_emb_dim_total + dense_input_dim, **mlp_params)
-        self.prediction_layer = PredictionLayer(task_type=self.task)
+        self.prediction_layer = TaskHead(task_type=self.task)
 
         # Register regularization weights
         self.register_regularization_weights(
@@ -218,4 +219,4 @@ class xDeepFM(BaseModel):
 
         # Combine all parts
         y = y_linear + y_cin + y_deep
-        return self.prediction_layer(y)
+        return self.prediction_layer(y)

nextrec/utils/config.py

@@ -28,7 +28,9 @@ if TYPE_CHECKING:
     from nextrec.data.preprocessor import DataProcessor
 
 
-def resolve_path(path_str: str | Path | None = None, base_dir: Path | None = None) -> Path:
+def resolve_path(
+    path_str: str | Path | None = None, base_dir: Path | None = None
+) -> Path:
     if path_str is None:
         return Path.cwd()
     path = Path(path_str).expanduser()
@@ -36,9 +38,16 @@ def resolve_path(path_str: str | Path | None = None, base_dir: Path | None = Non
         return path
     # Prefer resolving relative to current working directory when the path (or its parent)
     # already exists there; otherwise fall back to the config file's directory.
-    candidates = ((Path.cwd() / path).resolve(), ((base_dir or Path.cwd()) / path).resolve())
+    candidates = (
+        (Path.cwd() / path).resolve(),
+        ((base_dir or Path.cwd()) / path).resolve(),
+    )
     return next(
-        (candidate for candidate in candidates if candidate.exists() or candidate.parent.exists()),
+        (
+            candidate
+            for candidate in candidates
+            if candidate.exists() or candidate.parent.exists()
+        ),
         candidates[0],
     )
 

nextrec/utils/console.py

@@ -16,8 +16,9 @@ import numbers
 import os
 import platform
 import sys
+import time
 from datetime import datetime, timedelta
-from typing import Any, Callable, Iterable, Mapping, TypeVar
+from typing import Any, Callable, Mapping, TypeVar
 
 import numpy as np
 from rich import box
@@ -128,45 +129,99 @@ class BlackMofNCompleteColumn(MofNCompleteColumn):
         )
 
 
-def progress(
-    iterable: Iterable[T],
-    *,
-    description: str | None = None,
-    total: int | None = None,
-    disable: bool = False,
-) -> Iterable[T]:
+def progress(iterable, *, description=None, total=None, disable=False):
     if disable:
-        for item in iterable:
-            yield item
+        yield from iterable
         return
+
     resolved_total = total
     if resolved_total is None:
         try:
-            resolved_total = len(iterable)  # type: ignore[arg-type]
+            resolved_total = len(iterable)
         except TypeError:
             resolved_total = None
 
+    stream = sys.stderr
+
+    if not stream.isatty():
+        start_time = time.monotonic()
+        last_tick = start_time
+        min_interval_seconds = 10.0
+        max_interval_seconds = 300.0
+        target_steps = (
+            max(1, resolved_total // 20) if resolved_total is not None else 500
+        )
+        interval_seconds = min_interval_seconds
+        completed = 0
+
+        def emit(now: float):
+            elapsed = max(0.0, now - start_time)
+            speed = completed / elapsed if elapsed > 0 else 0.0
+            if resolved_total is not None and speed > 0:
+                remaining = max(0.0, resolved_total - completed)
+                eta_seconds = remaining / speed
+                eta_text = str(timedelta(seconds=int(eta_seconds)))
+            else:
+                eta_text = "--:--:--"
+            total_text = str(resolved_total) if resolved_total is not None else "?"
+            stream.write(
+                f"{description or 'Working'}: {completed}/{total_text} "
+                f"elapsed={timedelta(seconds=int(elapsed))} "
+                f"speed={speed:.2f}/s ETA={eta_text}\n"
+            )
+            stream.flush()
+            return speed
+
+        for item in iterable:
+            yield item
+            completed += 1
+            now = time.monotonic()
+            if now - last_tick >= interval_seconds:
+                speed = emit(now)
+                last_tick = now
+                if speed > 0:
+                    interval_seconds = min(
+                        max_interval_seconds,
+                        max(min_interval_seconds, target_steps / speed),
+                    )
+        end_now = time.monotonic()
+        if end_now - last_tick >= 1e-6:
+            emit(end_now)
+        return
+
+    # TTY: rich
+    console = Console(file=stream, force_terminal=True)
     progress_bar = Progress(
-        SpinnerColumn(style="black"),
-        TextColumn("{task.description}", style="black"),
-        BarColumn(
-            bar_width=36, style="black", complete_style="black", finished_style="black"
-        ),
-        TaskProgressColumn(style="black"),
-        BlackMofNCompleteColumn(),
-        BlackTimeElapsedColumn(),
-        BlackTimeRemainingColumn(),
+        SpinnerColumn(),
+        TextColumn("{task.description}"),
+        BarColumn(bar_width=36),
+        TaskProgressColumn(),
+        MofNCompleteColumn(),
+        TimeElapsedColumn(),
+        TimeRemainingColumn(),
         refresh_per_second=12,
+        console=console,
     )
 
-    task_id = progress_bar.add_task(description or "Working", total=resolved_total)
-    progress_bar.start()
-    try:
-        for item in iterable:
-            yield item
-            progress_bar.advance(task_id, 1)
-    finally:
-        progress_bar.stop()
+    if hasattr(progress_bar, "__enter__"):
+        with progress_bar:
+            task_id = progress_bar.add_task(
+                description or "Working", total=resolved_total
+            )
+            for item in iterable:
+                yield item
+                progress_bar.advance(task_id, 1)
+    else:
+        progress_bar.start()
+        try:
+            task_id = progress_bar.add_task(
+                description or "Working", total=resolved_total
+            )
+            for item in iterable:
+                yield item
+                progress_bar.advance(task_id, 1)
+        finally:
+            progress_bar.stop()
 
 
 def group_metrics_by_task(

{nextrec-0.4.15.dist-info → nextrec-0.4.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nextrec
-Version: 0.4.15
+Version: 0.4.17
 Summary: A comprehensive recommendation library with match, ranking, and multi-task learning models
 Project-URL: Homepage, https://github.com/zerolovesea/NextRec
 Project-URL: Repository, https://github.com/zerolovesea/NextRec
@@ -68,7 +68,7 @@ Description-Content-Type: text/markdown
 ![PyTorch](https://img.shields.io/badge/PyTorch-1.10+-ee4c2c.svg)
 
 ![License](https://img.shields.io/badge/License-Apache%202.0-green.svg)
-![Version](https://img.shields.io/badge/Version-0.4.15-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.4.17-orange.svg)
 
 
 中文文档 | [English Version](README_en.md)
@@ -102,7 +102,7 @@ NextRec是一个基于PyTorch的现代推荐系统框架，旨在为研究工程
 
 ## NextRec近期进展
 
-- **21/12/2025** 在v0.4.15中加入了对[GradNorm](/nextrec/loss/grad_norm.py)的支持，通过compile的`loss_weight='grad_norm'`进行配置
+- **21/12/2025** 在v0.4.16中加入了对[GradNorm](/nextrec/loss/grad_norm.py)的支持，通过compile的`loss_weight='grad_norm'`进行配置
 - **12/12/2025** 在v0.4.9中加入了[RQ-VAE](/nextrec/models/representation/rqvae.py)模块。配套的[数据集](/dataset/ecommerce_task.csv)和[代码](tutorials/notebooks/zh/使用RQ-VAE构建语义ID.ipynb)已经同步在仓库中
 - **07/12/2025** 发布了NextRec CLI命令行工具，它允许用户根据配置文件进行一键训练和推理，我们提供了相关的[教程](/nextrec_cli_preset/NextRec-CLI_zh.md)和[教学代码](/nextrec_cli_preset)
 - **03/12/2025** NextRec获得了100颗🌟！感谢大家的支持
@@ -244,11 +244,11 @@ nextrec --mode=train --train_config=path/to/train_config.yaml
 nextrec --mode=predict --predict_config=path/to/predict_config.yaml
 ```
 
-> 截止当前版本0.4.15，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
+> 截止当前版本0.4.17，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
 
 ## 兼容平台
 
-当前最新版本为0.4.15，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
+当前最新版本为0.4.17，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
 
 | 平台 | 配置 | 
 |------|------|

{nextrec-0.4.15.dist-info → nextrec-0.4.17.dist-info}/RECORD

@@ -1,14 +1,15 @@
 nextrec/__init__.py,sha256=_M3oUqyuvQ5k8Th_3wId6hQ_caclh7M5ad51XN09m98,235
-nextrec/__version__.py,sha256=yHOqz5A9VYGVIjRAkE5ZWR9IpLeDo8sygF-I11UMLv0,23
+nextrec/__version__.py,sha256=rJbmNVFVceaHhRo-pizixUQW0k5sqoDVdU1TfXfG8CA,23
 nextrec/cli.py,sha256=JUprwpoVbT4tXsGgMpj9Y_5yYByQXYMliMdWd38ReKo,24441
 nextrec/basic/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nextrec/basic/activation.py,sha256=uzTWfCOtBSkbu_Gk9XBNTj8__s241CaYLJk6l8nGX9I,2885
 nextrec/basic/callback.py,sha256=nn1f8FG9c52vJ-gvwteqPbk3-1QuNS1vmhBlkENdb0I,14636
 nextrec/basic/features.py,sha256=GyCUzGPuizUofrZSSOdqHK84YhnX4MGTdu7Cx2OGhUA,4654
+nextrec/basic/heads.py,sha256=0TqttvqkaUFkGg2bkMbFhOK63AuHs8sZVzvtbgpo1CE,3318
 nextrec/basic/layers.py,sha256=ZM3Nka3e2cit3e3peL0ukJCMgKZK1ovNFfAWvVOwlos,28556
 nextrec/basic/loggers.py,sha256=Zh1A5DVAFqlGglyaQ4_IMgvFbWAcXX5H3aHbCWA82nE,6524
-nextrec/basic/metrics.py,sha256=saNgM7kuHk9xqDxZF6x33irTaxeXCU-hxYTUQauuGgg,23074
-nextrec/basic/model.py,sha256=WGdGBK0oI8amRIor48xqIMij9UwrSkOMs0LhFczalzY,103710
+nextrec/basic/metrics.py,sha256=1r6efTc9TpARNBt5X9ISoppTZflej6EdFkjPYHV-YZI,23162
+nextrec/basic/model.py,sha256=MFJr6_FAT8-niEwZrqmo8oWzkjuspdjXRsizt1R056Q,103814
 nextrec/basic/session.py,sha256=UOG_-EgCOxvqZwCkiEd8sgNV2G1sm_HbzKYVQw8yYDI,4483
 nextrec/data/__init__.py,sha256=YZQjpty1pDCM7q_YNmiA2sa5kbujUw26ObLHWjMPjKY,1194
 nextrec/data/batch_utils.py,sha256=0bYGVX7RlhnHv_ZBaUngjDIpBNw-igCk98DgOsF7T6o,2879
@@ -25,28 +26,28 @@ nextrec/loss/pointwise.py,sha256=o9J3OznY0hlbDsUXqn3k-BBzYiuUH5dopz8QBFqS_kQ,734
 nextrec/models/generative/__init__.py,sha256=0MV3P-_ainPaTxmRBGWKUVCEt14KJvuvEHmRB3OQ1Fs,176
 nextrec/models/generative/tiger.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nextrec/models/multi_task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-nextrec/models/multi_task/esmm.py,sha256=AqesBZ4tOFNm7POCrHZ90h1zWWSViZAYfydUVOh2dEU,6545
-nextrec/models/multi_task/mmoe.py,sha256=aaQKcx4PL_mAanW3tkjAR886KmMCHTdBuu4p9EIKQJo,8657
-nextrec/models/multi_task/ple.py,sha256=fqkujPFGxxQOO_6nBZEz_UcxLEUoX_vCJsk0YOpxTg4,13084
-nextrec/models/multi_task/poso.py,sha256=J_Btxhm9JpFJMdQQHNNf9mMRHOgO7j1ts6VN5o4qJnk,19193
-nextrec/models/multi_task/share_bottom.py,sha256=DTWm6fpLCLiXimD-qk_0YIKT_9THMFDrnx4GDViXc_g,6583
+nextrec/models/multi_task/esmm.py,sha256=zVHuDwVP0HNx1tgjJOc9iz_uFXfUwoy5pVPn-XX0zBs,6561
+nextrec/models/multi_task/mmoe.py,sha256=lXK_kSw9TFTY6437yDA90_jerZi02lKdVM8zTCH3QMU,8673
+nextrec/models/multi_task/ple.py,sha256=9X0eSMIFGj4Dj9yy2tpKEI7vp9TL5XvgKehMLm4wpJs,13100
+nextrec/models/multi_task/poso.py,sha256=BomMkRx-tNtFiZ0mFBrRYms6PnffMvRlPkO9c5QZaiQ,19209
+nextrec/models/multi_task/share_bottom.py,sha256=i1vt3z6Ycj4SjucHrzmJeuIfo2_VDwSF9ubJeNzKy6c,6599
 nextrec/models/ranking/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-nextrec/models/ranking/afm.py,sha256=96jGUPL4yTWobMIVBjHpOxl9AtAzCAGR8yw7Sy2JmdQ,10125
-nextrec/models/ranking/autoint.py,sha256=S6Cxnp1q2OErSYqmIix5P-b4qLWR-0dY6TMStuU6WLg,8109
-nextrec/models/ranking/dcn.py,sha256=whkjiKEuadl6oSP-NJdSOCOqvWZGX4EsId9oqlfVpa8,7299
-nextrec/models/ranking/dcn_v2.py,sha256=QnqQbJsrtQp4mtvnBXFUVefKyr4dw-gHNWrCbO26oHw,11163
-nextrec/models/ranking/deepfm.py,sha256=aXoK59e2KaaPe5vfyFW4YiHbX4E2iG3gxFCxmWo8RHk,5200
-nextrec/models/ranking/dien.py,sha256=c7Zs85vxhOgKHg5s0QcSLCn1xXCCSD177TMERgM_v8g,18958
-nextrec/models/ranking/din.py,sha256=gdUhuKiKXBNOALbK8fGhlbSeuDT8agcEdNSrC_wveHc,9422
-nextrec/models/ranking/eulernet.py,sha256=SQr7M_6GI1u09jpxzRasQLFAPLXcmqff69waER6fiD8,12201
-nextrec/models/ranking/ffm.py,sha256=9t685SViSU1J0ESz-lrYSXhf4d4BWLNYZXReeVEk3e8,11262
-nextrec/models/ranking/fibinet.py,sha256=_eroddVHooJcaGT8MqS4mUrtv5j4pnTmfI3FoAKOZhs,7919
-nextrec/models/ranking/fm.py,sha256=SsrSKK3y4xg5Lv-t3JLnZan55Hzze2AxAiVPuscy0bk,4536
-nextrec/models/ranking/lr.py,sha256=MUonlKyA77_bfshTupFjOhY5tiuSJxApFM-_yOk4Nwk,4008
-nextrec/models/ranking/masknet.py,sha256=tY1y2lO0iq82oylPN0SBnL5Bikc8weinFXpURyVT1hE,12373
-nextrec/models/ranking/pnn.py,sha256=FcNIFAw5J0ORGSR6L8ZK7NeXlJPpojwe_SpsxMQqCFw,8174
-nextrec/models/ranking/widedeep.py,sha256=-ghKfe_0puvlI9fBQr8lK3gXkfVvslGwP40AJTGqc7w,5077
-nextrec/models/ranking/xdeepfm.py,sha256=FMtl_zYO1Ty_2d9VWRsz6Jo-Xjw8vikpIQPZCDVavVY,8156
+nextrec/models/ranking/afm.py,sha256=p3rNc957TaEc2tkhpbQyTup8a1oH34UJ2YKLqCuUD-4,10141
+nextrec/models/ranking/autoint.py,sha256=OgNh-NpC9ragHbIJfMCnQJRS6mh0wINkw961rGp32_Q,8125
+nextrec/models/ranking/dcn.py,sha256=c22ZBbBw3aEAISgtnsKZqhCvcYhM68AmEPZFS6MYO44,7315
+nextrec/models/ranking/dcn_v2.py,sha256=PqhkSPhGUx52hZMd2JPBqYQr1o5twCg6M6hCaaTl3aY,11179
+nextrec/models/ranking/deepfm.py,sha256=9K_TEptWwIr_SNC4yqN8PA2SzUKafZL94PVT2qH5Bag,5216
+nextrec/models/ranking/dien.py,sha256=URKWeOPeRD5OIWNsAxgVvbetOSrBHoq2eO5rR5UJ0jU,18971
+nextrec/models/ranking/din.py,sha256=Y8v0gONRt1OZORmn0hqMuzMfkvX0Nz1gByJ94jo3MUw,9435
+nextrec/models/ranking/eulernet.py,sha256=Mz2TObgYmYR7jIzwkafaySIxSn7P0ZVVFJTHZ0QEPFs,12217
+nextrec/models/ranking/ffm.py,sha256=0NB-ygR486lAJgarhOsdbaE8E0rbc1o8h5d8spWs5io,11278
+nextrec/models/ranking/fibinet.py,sha256=VP0gNoQwoLKxniv2HmHzxlnR3YlrnQJt6--CwmAgsW4,7932
+nextrec/models/ranking/fm.py,sha256=wX0VBvme41mFxSaT4Y1w4JfROL8EYT5KEe-hwRNr3zc,4552
+nextrec/models/ranking/lr.py,sha256=hc7XKa1GJhRa2CJY8-sYR0AB_fvmVLYCFoDRNCg04rU,4024
+nextrec/models/ranking/masknet.py,sha256=NNtIyAd-UHT78TvLcGbMvzR_aJRPzG3HBri1adNBrc4,12389
+nextrec/models/ranking/pnn.py,sha256=6a78M6n6B9aMBSvm0qICyBKjccoZ9TWVkDo9enc5d08,8190
+nextrec/models/ranking/widedeep.py,sha256=LNhLBmFP7M_z6ldxp99uMMK9s07EkP_9pjCAl6Hqp_Q,5093
+nextrec/models/ranking/xdeepfm.py,sha256=q4ITrgFfLe7NABxJQuchs8yqW_-jKwgXIwmpB2Lvoio,8172
 nextrec/models/representation/__init__.py,sha256=O3QHMMXBszwM-mTl7bA3wawNZvDGet-QIv6Ys5GHGJ8,190
 nextrec/models/representation/autorec.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nextrec/models/representation/bpr.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -64,15 +65,15 @@ nextrec/models/retrieval/youtube_dnn.py,sha256=xtGPV6_5LeSZBKkrTaU1CmtxlhgYLvZmj
 nextrec/models/sequential/hstu.py,sha256=P2Kl7HEL3afwiCApGKQ6UbUNO9eNXXrB10H7iiF8cI0,19735
 nextrec/models/sequential/sasrec.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nextrec/utils/__init__.py,sha256=C-1l-suSsN_MlPlj_5LApyCRQLOao5l7bO0SccwKHw4,2598
-nextrec/utils/config.py,sha256=2zcjK4TeN8ow-JSXbWpqyh9C1vFeKnEsYHPg1x564KU,19969
-nextrec/utils/console.py,sha256=e94SiwA0gKn2pfpP94mY_jl-kFok3TCjxo298KdFuP4,11696
+nextrec/utils/config.py,sha256=VgCh5fto8HGodwXPJacenqjxre3Aw6tw-mntW9n3OYA,20044
+nextrec/utils/console.py,sha256=wHAGQT0HmPumWlMlqQuBVXPVGCWubgYez-vuzkOMiMQ,13646
 nextrec/utils/data.py,sha256=alruiWZFbmwy3kO12q42VXmtHmXFFjVULpHa43fx_mI,21098
 nextrec/utils/embedding.py,sha256=akAEc062MG2cD7VIOllHaqtwzAirQR2gq5iW7oKpGAU,1449
 nextrec/utils/feature.py,sha256=rsUAv3ELyDpehVw8nPEEsLCCIjuKGTJJZuFaWB_wrPk,633
 nextrec/utils/model.py,sha256=3B85a0IJCggI26dxv25IX8R_5yQPo7wXI0JIAns6bkQ,1727
 nextrec/utils/torch_utils.py,sha256=AKfYbSOJjEw874xsDB5IO3Ote4X7vnqzt_E0jJny0o8,13468
-nextrec-0.4.15.dist-info/METADATA,sha256=fNo7--haI7VdwCBY-FXqtrDXTK6pVrkMde-opJJsLjg,21298
-nextrec-0.4.15.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-nextrec-0.4.15.dist-info/entry_points.txt,sha256=NN-dNSdfMRTv86bNXM7d3ZEPW2BQC6bRi7QP7i9cIps,45
-nextrec-0.4.15.dist-info/licenses/LICENSE,sha256=2fQfVKeafywkni7MYHyClC6RGGC3laLTXCNBx-ubtp0,1064
-nextrec-0.4.15.dist-info/RECORD,,
+nextrec-0.4.17.dist-info/METADATA,sha256=X-ezN2q05f0SKvwE80Y8ThecPGQI1D7RJCl6FEDXyjA,21298
+nextrec-0.4.17.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+nextrec-0.4.17.dist-info/entry_points.txt,sha256=NN-dNSdfMRTv86bNXM7d3ZEPW2BQC6bRi7QP7i9cIps,45
+nextrec-0.4.17.dist-info/licenses/LICENSE,sha256=2fQfVKeafywkni7MYHyClC6RGGC3laLTXCNBx-ubtp0,1064
+nextrec-0.4.17.dist-info/RECORD,,