nextrec 0.4.16__tar.gz → 0.4.17__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nextrec
-Version: 0.4.16
+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.16-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.4.17-orange.svg)
 
 
 中文文档 | [English Version](README_en.md)
@@ -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.16，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
+> 截止当前版本0.4.17，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
 
 ## 兼容平台
 
-当前最新版本为0.4.16，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
+当前最新版本为0.4.17，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
 
 | 平台 | 配置 | 
 |------|------|

{nextrec-0.4.16 → 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.16-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.4.17-orange.svg)
 
 
 中文文档 | [English Version](README_en.md)
@@ -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.16，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
+> 截止当前版本0.4.17，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
 
 ## 兼容平台
 
-当前最新版本为0.4.16，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
+当前最新版本为0.4.17，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
 
 | 平台 | 配置 | 
 |------|------|

{nextrec-0.4.16 → 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.16-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.4.17-orange.svg)
 
 English | [中文文档](README.md)
 
@@ -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.16, 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.16. 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.16 → 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.16"
+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.16 → 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
@@ -2115,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,
@@ -2244,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.16 → 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.16 → 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.16 → 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.16 → 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.16 → 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.16 → 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.16 → 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.16 → 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.16 → 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.16 → 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.16 → 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.16 → 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.16 → 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.16 → 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)

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