nextrec 0.4.15__tar.gz → 0.4.17__tar.gz

{nextrec-0.4.15 → nextrec-0.4.17}/PKG-INFO

@@ -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 → nextrec-0.4.17}/README.md

@@ -9,7 +9,7 @@
 ![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)
@@ -43,7 +43,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颗🌟！感谢大家的支持
@@ -185,11 +185,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 → nextrec-0.4.17}/README_en.md

@@ -9,7 +9,7 @@
 ![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 | [中文文档](README.md)
 
@@ -44,7 +44,7 @@ NextRec is a modern recommendation framework built on PyTorch, delivering a unif
 
 ## NextRec Progress
 
-- **21/12/2025** Added support for [GradNorm](/nextrec/loss/grad_norm.py) in v0.4.15, configurable via `loss_weight='grad_norm'` in the compile method
+- **21/12/2025** Added support for [GradNorm](/nextrec/loss/grad_norm.py) in v0.4.16, configurable via `loss_weight='grad_norm'` in the compile method
 - **12/12/2025** Added [RQ-VAE](/nextrec/models/representation/rqvae.py), a common module for generative retrieval in v0.4.9. Paired [dataset](/dataset/ecommerce_task.csv) and [notebook code](tutorials/notebooks/en/Build%20semantic%20ID%20with%20RQ-VAE.ipynb) are available.
 - **07/12/2025** Released the NextRec CLI tool to run training/inference from configs. See the [guide](/nextrec_cli_preset/NextRec-CLI.md) and [reference code](/nextrec_cli_preset).
 - **03/12/2025** NextRec reached 100 ⭐—thanks for the support!
@@ -188,11 +188,11 @@ nextrec --mode=train --train_config=path/to/train_config.yaml
 nextrec --mode=predict --predict_config=path/to/predict_config.yaml
 ```
 
-> As of version 0.4.15, NextRec CLI supports single-machine training; distributed training features are currently under development.
+> As of version 0.4.17, NextRec CLI supports single-machine training; distributed training features are currently under development.
 
 ## Platform Compatibility
 
-The current version is 0.4.15. All models and test code have been validated on the following platforms. If you encounter compatibility issues, please report them in the issue tracker with your system version:
+The current version is 0.4.17. All models and test code have been validated on the following platforms. If you encounter compatibility issues, please report them in the issue tracker with your system version:
 
 | Platform | Configuration | 
 |----------|---------------|

{nextrec-0.4.15 → nextrec-0.4.17}/docs/rtd/conf.py

@@ -11,7 +11,7 @@ sys.path.insert(0, str(PROJECT_ROOT / "nextrec"))
 project = "NextRec"
 copyright = "2025, Yang Zhou"
 author = "Yang Zhou"
-release = "0.4.15"
+release = "0.4.17"
 
 extensions = [
     "myst_parser",

nextrec-0.4.17/nextrec/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.4.17"

nextrec-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-0.4.15 → nextrec-0.4.17}/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-0.4.15 → nextrec-0.4.17}/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-0.4.15 → nextrec-0.4.17}/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-0.4.15 → nextrec-0.4.17}/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-0.4.15 → nextrec-0.4.17}/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-0.4.15 → nextrec-0.4.17}/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-0.4.15 → nextrec-0.4.17}/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-0.4.15 → nextrec-0.4.17}/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-0.4.15 → nextrec-0.4.17}/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-0.4.15 → nextrec-0.4.17}/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-0.4.15 → nextrec-0.4.17}/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-0.4.15 → nextrec-0.4.17}/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-0.4.15 → nextrec-0.4.17}/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-0.4.15 → nextrec-0.4.17}/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-0.4.15 → nextrec-0.4.17}/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-0.4.15 → nextrec-0.4.17}/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)