nextrec 0.4.16__py3-none-any.whl → 0.4.18__py3-none-any.whl

nextrec/models/ranking/afm.py

@@ -1,6 +1,6 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 09/12/2025
+Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 [1] Xiao J, Ye H, He X, et al. Attentional factorization machines: Learning the weight of
@@ -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

nextrec/models/ranking/autoint.py

@@ -1,6 +1,6 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 09/12/2025
+Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 [1] Song W, Shi C, Xiao Z, et al. Autoint: Automatic feature interaction learning via
@@ -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(

nextrec/models/ranking/dcn.py

@@ -1,6 +1,6 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 09/12/2025
+Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 [1] Wang R, Fu B, Fu G, et al. Deep & cross network for ad click predictions[C]
@@ -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(

nextrec/models/ranking/dcn_v2.py

@@ -1,6 +1,6 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 09/12/2025
+Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 [1] R. Wang et al. DCN V2: Improved Deep & Cross Network and Practical Lessons for
@@ -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",

nextrec/models/ranking/deepfm.py

@@ -1,6 +1,6 @@
 """
 Date: create on 27/10/2025
-Checkpoint: edit on 24/11/2025
+Checkpoint: edit on 23/12/2025
 Author: Yang Zhou,zyaztec@gmail.com
 Reference:
 [1] Guo H, Tang R, Ye Y, et al. DeepFM: A factorization-machine based neural network
@@ -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(

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

@@ -1,6 +1,6 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 09/12/2025
+Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 [1] Zhao Z, Zhang H, Tang H, et al. EulerNet: Efficient and Effective Feature
@@ -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:

nextrec/models/ranking/ffm.py

@@ -1,6 +1,6 @@
 """
 Date: create on 19/12/2025
-Checkpoint: edit on 19/12/2025
+Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 [1] Juan Y, Zhuang Y, Chin W-S, et al. Field-aware Factorization Machines for CTR
@@ -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()

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

@@ -1,6 +1,6 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 09/12/2025
+Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 [1] Rendle S. Factorization machines[C]//ICDM. 2010: 995-1000.
@@ -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(

nextrec/models/ranking/lr.py

@@ -1,6 +1,6 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 09/12/2025
+Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 [1] Hosmer D W, Lemeshow S, Sturdivant R X. Applied Logistic Regression.
@@ -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"]

nextrec/models/ranking/masknet.py

@@ -1,6 +1,6 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 29/11/2025
+Checkpoint: edit on 23/12/2025
 Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 [1] Wang Z, She Q, Zhang J. MaskNet: Introducing Feature-Wise
@@ -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,14 +283,13 @@ 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(
                 embedding_attr="embedding",
                 include_modules=["mask_blocks", "output_layer"],
             )
-        # serial
         else:
             self.register_regularization_weights(
                 embedding_attr="embedding", include_modules=["mask_blocks", "final_mlp"]
@@ -314,7 +314,6 @@ class MaskNet(BaseModel):
                 block_outputs.append(h)
             concat_hidden = torch.cat(block_outputs, dim=-1)
             logit = self.final_mlp(concat_hidden)  # [B, 1]
-        # serial
         else:
             hidden = self.first_block(field_emb, v_emb_flat)
             hidden = self.block_dropout(hidden)

nextrec/models/ranking/pnn.py

@@ -1,7 +1,7 @@
 """
 Date: create on 09/11/2025
-Author:
-    Yang Zhou,zyaztec@gmail.com
+Checkpoint: edit on 23/12/2025
+Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 [1] Qu Y, Cai H, Ren K, et al. Product-based neural networks for user response
 prediction[C]//ICDM. 2016: 1149-1154. (https://arxiv.org/abs/1611.00144)
@@ -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:

nextrec/models/ranking/widedeep.py

@@ -1,12 +1,11 @@
 """
 Date: create on 09/11/2025
-Checkpoint: edit on 24/11/2025
-Author:
-    Yang Zhou,zyaztec@gmail.com
+Checkpoint: edit on 23/12/2025
+Author: Yang Zhou, zyaztec@gmail.com
 Reference:
-    [1] Cheng H T, Koc L, Harmsen J, et al. Wide & Deep learning for recommender systems[C]
-        //Proceedings of the 1st Workshop on Deep Learning for Recommender Systems. 2016: 7-10.
-        (https://arxiv.org/abs/1606.07792)
+[1] Cheng H T, Koc L, Harmsen J, et al. Wide & Deep learning for recommender systems[C]
+//Proceedings of the 1st Workshop on Deep Learning for Recommender Systems. 2016: 7-10.
+(https://arxiv.org/abs/1606.07792)
 
 Wide & Deep blends a linear wide component (memorization of cross features) with a
 deep neural network (generalization) sharing the same feature space. The wide part
@@ -42,7 +41,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 +114,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"]

nextrec/models/ranking/xdeepfm.py

@@ -1,7 +1,7 @@
 """
 Date: create on 09/11/2025
-Author:
-Yang Zhou,zyaztec@gmail.com
+Checkpoint: edit on 23/12/2025
+Author: Yang Zhou, zyaztec@gmail.com
 Reference:
 [1] Lian J, Zhou X, Zhang F, et al. xdeepfm: Combining explicit and implicit feature interactions
 for recommender systems[C]//Proceedings of the 24th ACM SIGKDD international conference on
@@ -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(

nextrec/utils/console.py

@@ -203,18 +203,32 @@ def progress(iterable, *, description=None, total=None, disable=False):
         console=console,
     )
 
-    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)
+    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(
     metrics: Mapping[str, Any] | None,
     target_names: list[str] | str | None,
     default_task_name: str = "overall",
-) -> tuple[list[str], dict[str, dict[str, float]]]:
+):
     if not metrics:
         return [], {}