nextrec 0.1.4__py3-none-any.whl → 0.1.7__py3-none-any.whl

nextrec/models/match/youtube_dnn.py

@@ -6,7 +6,6 @@ Reference:
     [1] Covington P, Adams J, Sargin E. Deep neural networks for youtube recommendations[C]
         //Proceedings of the 10th ACM conference on recommender systems. 2016: 191-198.
 """
-
 import torch
 import torch.nn as nn
 from typing import Literal
@@ -19,42 +18,40 @@ from nextrec.basic.layers import MLP, EmbeddingLayer, AveragePooling
 class YoutubeDNN(BaseMatchModel):
     """
     YouTube Deep Neural Network for Recommendations
-
+    
     用户塔：历史行为序列 + 用户特征 -> 用户embedding
     物品塔：物品特征 -> 物品embedding
     训练：sampled softmax loss (listwise)
     """
-
+    
     @property
     def model_name(self) -> str:
         return "YouTubeDNN"
-
-    def __init__(
-        self,
-        user_dense_features: list[DenseFeature] | None = None,
-        user_sparse_features: list[SparseFeature] | None = None,
-        user_sequence_features: list[SequenceFeature] | None = None,
-        item_dense_features: list[DenseFeature] | None = None,
-        item_sparse_features: list[SparseFeature] | None = None,
-        item_sequence_features: list[SequenceFeature] | None = None,
-        user_dnn_hidden_units: list[int] = [256, 128, 64],
-        item_dnn_hidden_units: list[int] = [256, 128, 64],
-        embedding_dim: int = 64,
-        dnn_activation: str = "relu",
-        dnn_dropout: float = 0.0,
-        training_mode: Literal["pointwise", "pairwise", "listwise"] = "listwise",
-        num_negative_samples: int = 100,
-        temperature: float = 1.0,
-        similarity_metric: Literal["dot", "cosine", "euclidean"] = "dot",
-        device: str = "cpu",
-        embedding_l1_reg: float = 0.0,
-        dense_l1_reg: float = 0.0,
-        embedding_l2_reg: float = 0.0,
-        dense_l2_reg: float = 0.0,
-        early_stop_patience: int = 20,
-        model_id: str = "youtube_dnn",
-    ):
-
+    
+    def __init__(self,
+                 user_dense_features: list[DenseFeature] | None = None,
+                 user_sparse_features: list[SparseFeature] | None = None,
+                 user_sequence_features: list[SequenceFeature] | None = None,
+                 item_dense_features: list[DenseFeature] | None = None,
+                 item_sparse_features: list[SparseFeature] | None = None,
+                 item_sequence_features: list[SequenceFeature] | None = None,
+                 user_dnn_hidden_units: list[int] = [256, 128, 64],
+                 item_dnn_hidden_units: list[int] = [256, 128, 64],
+                 embedding_dim: int = 64,
+                 dnn_activation: str = 'relu',
+                 dnn_dropout: float = 0.0,
+                 training_mode: Literal['pointwise', 'pairwise', 'listwise'] = 'listwise',
+                 num_negative_samples: int = 100,
+                 temperature: float = 1.0,
+                 similarity_metric: Literal['dot', 'cosine', 'euclidean'] = 'dot',
+                 device: str = 'cpu',
+                 embedding_l1_reg: float = 0.0,
+                 dense_l1_reg: float = 0.0,
+                 embedding_l2_reg: float = 0.0,
+                 dense_l2_reg: float = 0.0,
+                 early_stop_patience: int = 20,
+                 model_id: str = 'youtube_dnn'):
+        
         super(YoutubeDNN, self).__init__(
             user_dense_features=user_dense_features,
             user_sparse_features=user_sparse_features,
@@ -72,13 +69,13 @@ class YoutubeDNN(BaseMatchModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=early_stop_patience,
-            model_id=model_id,
+            model_id=model_id
         )
-
+        
         self.embedding_dim = embedding_dim
         self.user_dnn_hidden_units = user_dnn_hidden_units
         self.item_dnn_hidden_units = item_dnn_hidden_units
-
+        
         # User tower
         user_features = []
         if user_dense_features:
@@ -87,10 +84,10 @@ class YoutubeDNN(BaseMatchModel):
             user_features.extend(user_sparse_features)
         if user_sequence_features:
             user_features.extend(user_sequence_features)
-
+        
         if len(user_features) > 0:
             self.user_embedding = EmbeddingLayer(user_features)
-
+            
             user_input_dim = 0
             for feat in user_dense_features or []:
                 user_input_dim += 1
@@ -99,16 +96,16 @@ class YoutubeDNN(BaseMatchModel):
             for feat in user_sequence_features or []:
                 # 序列特征通过平均池化聚合
                 user_input_dim += feat.embedding_dim
-
+            
             user_dnn_units = user_dnn_hidden_units + [embedding_dim]
             self.user_dnn = MLP(
                 input_dim=user_input_dim,
                 dims=user_dnn_units,
                 output_layer=False,
                 dropout=dnn_dropout,
-                activation=dnn_activation,
+                activation=dnn_activation
             )
-
+        
         # Item tower
         item_features = []
         if item_dense_features:
@@ -117,10 +114,10 @@ class YoutubeDNN(BaseMatchModel):
             item_features.extend(item_sparse_features)
         if item_sequence_features:
             item_features.extend(item_sequence_features)
-
+        
         if len(item_features) > 0:
             self.item_embedding = EmbeddingLayer(item_features)
-
+            
             item_input_dim = 0
             for feat in item_dense_features or []:
                 item_input_dim += 1
@@ -128,54 +125,48 @@ class YoutubeDNN(BaseMatchModel):
                 item_input_dim += feat.embedding_dim
             for feat in item_sequence_features or []:
                 item_input_dim += feat.embedding_dim
-
+            
             item_dnn_units = item_dnn_hidden_units + [embedding_dim]
             self.item_dnn = MLP(
                 input_dim=item_input_dim,
                 dims=item_dnn_units,
                 output_layer=False,
                 dropout=dnn_dropout,
-                activation=dnn_activation,
+                activation=dnn_activation
             )
-
+        
         self._register_regularization_weights(
-            embedding_attr="user_embedding", include_modules=["user_dnn"]
+            embedding_attr='user_embedding',
+            include_modules=['user_dnn']
         )
         self._register_regularization_weights(
-            embedding_attr="item_embedding", include_modules=["item_dnn"]
+            embedding_attr='item_embedding',
+            include_modules=['item_dnn']
         )
-
+        
         self.to(device)
-
+    
     def user_tower(self, user_input: dict) -> torch.Tensor:
         """
         User tower
         处理用户历史行为序列和其他用户特征
         """
-        all_user_features = (
-            self.user_dense_features
-            + self.user_sparse_features
-            + self.user_sequence_features
-        )
+        all_user_features = self.user_dense_features + self.user_sparse_features + self.user_sequence_features
         user_emb = self.user_embedding(user_input, all_user_features, squeeze_dim=True)
         user_emb = self.user_dnn(user_emb)
-
+        
         # L2 normalization
         user_emb = torch.nn.functional.normalize(user_emb, p=2, dim=1)
-
+        
         return user_emb
-
+    
     def item_tower(self, item_input: dict) -> torch.Tensor:
         """Item tower"""
-        all_item_features = (
-            self.item_dense_features
-            + self.item_sparse_features
-            + self.item_sequence_features
-        )
+        all_item_features = self.item_dense_features + self.item_sparse_features + self.item_sequence_features
         item_emb = self.item_embedding(item_input, all_item_features, squeeze_dim=True)
         item_emb = self.item_dnn(item_emb)
-
+        
         # L2 normalization
         item_emb = torch.nn.functional.normalize(item_emb, p=2, dim=1)
-
+        
         return item_emb

nextrec/models/multi_task/esmm.py

@@ -17,15 +17,15 @@ from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
 class ESMM(BaseModel):
     """
     Entire Space Multi-Task Model
-
+    
     ESMM is designed for CVR (Conversion Rate) prediction. It models two related tasks:
     - CTR task: P(click | impression)
     - CVR task: P(conversion | click)
     - CTCVR task (auxiliary): P(click & conversion | impression) = P(click) * P(conversion | click)
-
+    
     This design addresses the sample selection bias and data sparsity issues in CVR modeling.
     """
-
+    
     @property
     def model_name(self):
         return "ESMM"
@@ -33,34 +33,30 @@ class ESMM(BaseModel):
     @property
     def task_type(self):
         # ESMM has fixed task types: CTR (binary) and CVR (binary)
-        return ["binary", "binary"]
-
-    def __init__(
-        self,
-        dense_features: list[DenseFeature],
-        sparse_features: list[SparseFeature],
-        sequence_features: list[SequenceFeature],
-        ctr_params: dict,
-        cvr_params: dict,
-        target: list[str] = ["ctr", "ctcvr"],  # Note: ctcvr = ctr * cvr
-        task: str | list[str] = "binary",
-        optimizer: str = "adam",
-        optimizer_params: dict = {},
-        loss: str | nn.Module | list[str | nn.Module] | None = "bce",
-        device: str = "cpu",
-        model_id: str = "baseline",
-        embedding_l1_reg=1e-6,
-        dense_l1_reg=1e-5,
-        embedding_l2_reg=1e-5,
-        dense_l2_reg=1e-4,
-    ):
-
+        return ['binary', 'binary']
+    
+    def __init__(self,
+                 dense_features: list[DenseFeature],
+                 sparse_features: list[SparseFeature],
+                 sequence_features: list[SequenceFeature],
+                 ctr_params: dict,
+                 cvr_params: dict,
+                 target: list[str] = ['ctr', 'ctcvr'],  # Note: ctcvr = ctr * cvr
+                 task: str | list[str] = 'binary',
+                 optimizer: str = "adam",
+                 optimizer_params: dict = {},
+                 loss: str | nn.Module | list[str | nn.Module] | None = "bce",
+                 device: str = 'cpu',
+                 model_id: str = "baseline",
+                 embedding_l1_reg=1e-6,
+                 dense_l1_reg=1e-5,
+                 embedding_l2_reg=1e-5,
+                 dense_l2_reg=1e-4):
+        
         # ESMM requires exactly 2 targets: ctr and ctcvr
         if len(target) != 2:
-            raise ValueError(
-                f"ESMM requires exactly 2 targets (ctr and ctcvr), got {len(target)}"
-            )
-
+            raise ValueError(f"ESMM requires exactly 2 targets (ctr and ctcvr), got {len(target)}")
+        
         super(ESMM, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -73,13 +69,13 @@ class ESMM(BaseModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=20,
-            model_id=model_id,
+            model_id=model_id
         )
 
         self.loss = loss
         if self.loss is None:
             self.loss = "bce"
-
+            
         # All features
         self.all_features = dense_features + sparse_features + sequence_features
 
@@ -87,48 +83,46 @@ class ESMM(BaseModel):
         self.embedding = EmbeddingLayer(features=self.all_features)
 
         # Calculate input dimension
-        emb_dim_total = sum(
-            [
-                f.embedding_dim
-                for f in self.all_features
-                if not isinstance(f, DenseFeature)
-            ]
-        )
-        dense_input_dim = sum(
-            [getattr(f, "embedding_dim", 1) or 1 for f in dense_features]
-        )
+        emb_dim_total = sum([f.embedding_dim for f in self.all_features if not isinstance(f, DenseFeature)])
+        dense_input_dim = sum([getattr(f, "embedding_dim", 1) or 1 for f in dense_features])
         input_dim = emb_dim_total + dense_input_dim
-
+        
         # CTR tower
         self.ctr_tower = MLP(input_dim=input_dim, output_layer=True, **ctr_params)
-
+        
         # CVR tower
         self.cvr_tower = MLP(input_dim=input_dim, output_layer=True, **cvr_params)
         self.prediction_layer = PredictionLayer(
-            task_type=self.task_type, task_dims=[1, 1]
+            task_type=self.task_type,
+            task_dims=[1, 1]
         )
 
         # Register regularization weights
         self._register_regularization_weights(
-            embedding_attr="embedding", include_modules=["ctr_tower", "cvr_tower"]
+            embedding_attr='embedding',
+            include_modules=['ctr_tower', 'cvr_tower']
         )
 
-        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss)
+        self.compile(
+            optimizer=optimizer,
+            optimizer_params=optimizer_params,
+            loss=loss
+        )
 
     def forward(self, x):
         # Get all embeddings and flatten
         input_flat = self.embedding(x=x, features=self.all_features, squeeze_dim=True)
-
+        
         # CTR prediction: P(click | impression)
         ctr_logit = self.ctr_tower(input_flat)  # [B, 1]
         cvr_logit = self.cvr_tower(input_flat)  # [B, 1]
         logits = torch.cat([ctr_logit, cvr_logit], dim=1)
         preds = self.prediction_layer(logits)
         ctr, cvr = preds.chunk(2, dim=1)
-
+        
         # CTCVR prediction: P(click & conversion | impression) = P(click) * P(conversion | click)
         ctcvr = ctr * cvr  # [B, 1]
-
+        
         # Output: [CTR, CTCVR]
         # Note: We supervise CTR with click labels and CTCVR with conversion labels
         y = torch.cat([ctr, ctcvr], dim=1)  # [B, 2]

nextrec/models/multi_task/mmoe.py

@@ -17,13 +17,13 @@ from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
 class MMOE(BaseModel):
     """
     Multi-gate Mixture-of-Experts
-
+    
     MMOE improves upon shared-bottom architecture by using multiple expert networks
     and task-specific gating networks. Each task has its own gate that learns to
     weight the contributions of different experts, allowing for both task-specific
     and shared representations.
     """
-
+    
     @property
     def model_name(self):
         return "MMOE"
@@ -31,28 +31,26 @@ class MMOE(BaseModel):
     @property
     def task_type(self):
         return self.task if isinstance(self.task, list) else [self.task]
-
-    def __init__(
-        self,
-        dense_features: list[DenseFeature] = [],
-        sparse_features: list[SparseFeature] = [],
-        sequence_features: list[SequenceFeature] = [],
-        expert_params: dict = {},
-        num_experts: int = 3,
-        tower_params_list: list[dict] = [],
-        target: list[str] = [],
-        task: str | list[str] = "binary",
-        optimizer: str = "adam",
-        optimizer_params: dict = {},
-        loss: str | nn.Module | list[str | nn.Module] | None = "bce",
-        device: str = "cpu",
-        model_id: str = "baseline",
-        embedding_l1_reg=1e-6,
-        dense_l1_reg=1e-5,
-        embedding_l2_reg=1e-5,
-        dense_l2_reg=1e-4,
-    ):
-
+    
+    def __init__(self,
+                 dense_features: list[DenseFeature]=[],
+                 sparse_features: list[SparseFeature]=[],
+                 sequence_features: list[SequenceFeature]=[],
+                 expert_params: dict={},
+                 num_experts: int=3,
+                 tower_params_list: list[dict]=[],
+                 target: list[str]=[],
+                 task: str | list[str] = 'binary',
+                 optimizer: str = "adam",
+                 optimizer_params: dict = {},
+                 loss: str | nn.Module | list[str | nn.Module] | None = "bce",
+                 device: str = 'cpu',
+                 model_id: str = "baseline",
+                 embedding_l1_reg=1e-6,
+                 dense_l1_reg=1e-5,
+                 embedding_l2_reg=1e-5,
+                 dense_l2_reg=1e-4):
+        
         super(MMOE, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -65,22 +63,20 @@ class MMOE(BaseModel):
             embedding_l2_reg=embedding_l2_reg,
             dense_l2_reg=dense_l2_reg,
             early_stop_patience=20,
-            model_id=model_id,
+            model_id=model_id
         )
 
         self.loss = loss
         if self.loss is None:
             self.loss = "bce"
-
+        
         # Number of tasks and experts
         self.num_tasks = len(target)
         self.num_experts = num_experts
-
+        
         if len(tower_params_list) != self.num_tasks:
-            raise ValueError(
-                f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.num_tasks})"
-            )
-
+            raise ValueError(f"Number of tower params ({len(tower_params_list)}) must match number of tasks ({self.num_tasks})")
+            
         # All features
         self.all_features = dense_features + sparse_features + sequence_features
 
@@ -88,83 +84,78 @@ class MMOE(BaseModel):
         self.embedding = EmbeddingLayer(features=self.all_features)
 
         # Calculate input dimension
-        emb_dim_total = sum(
-            [
-                f.embedding_dim
-                for f in self.all_features
-                if not isinstance(f, DenseFeature)
-            ]
-        )
-        dense_input_dim = sum(
-            [getattr(f, "embedding_dim", 1) or 1 for f in dense_features]
-        )
+        emb_dim_total = sum([f.embedding_dim for f in self.all_features if not isinstance(f, DenseFeature)])
+        dense_input_dim = sum([getattr(f, "embedding_dim", 1) or 1 for f in dense_features])
         input_dim = emb_dim_total + dense_input_dim
-
+        
         # Expert networks (shared by all tasks)
         self.experts = nn.ModuleList()
         for _ in range(num_experts):
             expert = MLP(input_dim=input_dim, output_layer=False, **expert_params)
             self.experts.append(expert)
-
+        
         # Get expert output dimension
-        if "dims" in expert_params and len(expert_params["dims"]) > 0:
-            expert_output_dim = expert_params["dims"][-1]
+        if 'dims' in expert_params and len(expert_params['dims']) > 0:
+            expert_output_dim = expert_params['dims'][-1]
         else:
             expert_output_dim = input_dim
-
+        
         # Task-specific gates
         self.gates = nn.ModuleList()
         for _ in range(self.num_tasks):
-            gate = nn.Sequential(nn.Linear(input_dim, num_experts), nn.Softmax(dim=1))
+            gate = nn.Sequential(
+                nn.Linear(input_dim, num_experts),
+                nn.Softmax(dim=1)
+            )
             self.gates.append(gate)
-
+        
         # Task-specific towers
         self.towers = nn.ModuleList()
         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(
-            task_type=self.task_type, task_dims=[1] * self.num_tasks
+            task_type=self.task_type,
+            task_dims=[1] * self.num_tasks
         )
 
         # Register regularization weights
         self._register_regularization_weights(
-            embedding_attr="embedding", include_modules=["experts", "gates", "towers"]
+            embedding_attr='embedding',
+            include_modules=['experts', 'gates', 'towers']
         )
 
-        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, loss=loss)
+        self.compile(
+            optimizer=optimizer,
+            optimizer_params=optimizer_params,
+            loss=loss
+        )
 
     def forward(self, x):
         # Get all embeddings and flatten
         input_flat = self.embedding(x=x, features=self.all_features, squeeze_dim=True)
-
+        
         # Expert outputs: [num_experts, B, expert_dim]
         expert_outputs = [expert(input_flat) for expert in self.experts]
-        expert_outputs = torch.stack(
-            expert_outputs, dim=0
-        )  # [num_experts, B, expert_dim]
-
+        expert_outputs = torch.stack(expert_outputs, dim=0)  # [num_experts, B, expert_dim]
+        
         # Task-specific processing
         task_outputs = []
         for task_idx in range(self.num_tasks):
             # Gate weights for this task: [B, num_experts]
             gate_weights = self.gates[task_idx](input_flat)  # [B, num_experts]
-
+            
             # Weighted sum of expert outputs
             # gate_weights: [B, num_experts, 1]
             # expert_outputs: [num_experts, B, expert_dim]
             gate_weights = gate_weights.unsqueeze(2)  # [B, num_experts, 1]
-            expert_outputs_t = expert_outputs.permute(
-                1, 0, 2
-            )  # [B, num_experts, expert_dim]
-            gated_output = torch.sum(
-                gate_weights * expert_outputs_t, dim=1
-            )  # [B, expert_dim]
-
+            expert_outputs_t = expert_outputs.permute(1, 0, 2)  # [B, num_experts, expert_dim]
+            gated_output = torch.sum(gate_weights * expert_outputs_t, dim=1)  # [B, expert_dim]
+            
             # Tower output
             tower_output = self.towers[task_idx](gated_output)  # [B, 1]
             task_outputs.append(tower_output)
-
+        
         # Stack outputs: [B, num_tasks]
         y = torch.cat(task_outputs, dim=1)
         return self.prediction_layer(y)