nextrec 0.1.1__py3-none-any.whl

nextrec/models/match/youtube_dnn.py

@@ -0,0 +1,172 @@
+"""
+Date: create on 09/11/2025
+Author:
+    Yang Zhou,zyaztec@gmail.com
+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
+
+from nextrec.basic.model import BaseMatchModel
+from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
+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'):
+        
+        super(YoutubeDNN, self).__init__(
+            user_dense_features=user_dense_features,
+            user_sparse_features=user_sparse_features,
+            user_sequence_features=user_sequence_features,
+            item_dense_features=item_dense_features,
+            item_sparse_features=item_sparse_features,
+            item_sequence_features=item_sequence_features,
+            training_mode=training_mode,
+            num_negative_samples=num_negative_samples,
+            temperature=temperature,
+            similarity_metric=similarity_metric,
+            device=device,
+            embedding_l1_reg=embedding_l1_reg,
+            dense_l1_reg=dense_l1_reg,
+            embedding_l2_reg=embedding_l2_reg,
+            dense_l2_reg=dense_l2_reg,
+            early_stop_patience=early_stop_patience,
+            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:
+            user_features.extend(user_dense_features)
+        if user_sparse_features:
+            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
+            for feat in user_sparse_features or []:
+                user_input_dim += feat.embedding_dim
+            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
+            )
+        
+        # Item tower
+        item_features = []
+        if item_dense_features:
+            item_features.extend(item_dense_features)
+        if item_sparse_features:
+            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
+            for feat in item_sparse_features or []:
+                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
+            )
+        
+        self._register_regularization_weights(
+            embedding_attr='user_embedding',
+            include_modules=['user_dnn']
+        )
+        self._register_regularization_weights(
+            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
+        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
+        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

@@ -0,0 +1,129 @@
+"""
+Date: create on 09/11/2025
+Author:
+    Yang Zhou,zyaztec@gmail.com
+Reference:
+    [1] Ma X, Zhao L, Huang G, et al. Entire space multi-task model: An effective approach for estimating post-click conversion rate[C]//SIGIR. 2018: 1137-1140.
+"""
+
+import torch
+import torch.nn as nn
+
+from nextrec.basic.model import BaseModel
+from nextrec.basic.layers import EmbeddingLayer, MLP, PredictionLayer
+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"
+
+    @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):
+        
+        # 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)}")
+        
+        super(ESMM, self).__init__(
+            dense_features=dense_features,
+            sparse_features=sparse_features,
+            sequence_features=sequence_features,
+            target=target,
+            task=task,  # Both CTR and CTCVR are binary classification
+            device=device,
+            embedding_l1_reg=embedding_l1_reg,
+            dense_l1_reg=dense_l1_reg,
+            embedding_l2_reg=embedding_l2_reg,
+            dense_l2_reg=dense_l2_reg,
+            early_stop_patience=20,
+            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
+
+        # Shared embedding layer
+        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])
+        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]
+        )
+
+        # Register regularization weights
+        self._register_regularization_weights(
+            embedding_attr='embedding',
+            include_modules=['ctr_tower', 'cvr_tower']
+        )
+
+        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]
+        return y  # [B, 2], where y[:, 0] is CTR and y[:, 1] is CTCVR

nextrec/models/multi_task/mmoe.py

@@ -0,0 +1,161 @@
+"""
+Date: create on 09/11/2025
+Author:
+    Yang Zhou,zyaztec@gmail.com
+Reference:
+    [1] Ma J, Zhao Z, Yi X, et al. Modeling task relationships in multi-task learning with multi-gate mixture-of-experts[C]//KDD. 2018: 1930-1939.
+"""
+
+import torch
+import torch.nn as nn
+
+from nextrec.basic.model import BaseModel
+from nextrec.basic.layers import EmbeddingLayer, MLP, PredictionLayer
+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"
+
+    @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):
+        
+        super(MMOE, self).__init__(
+            dense_features=dense_features,
+            sparse_features=sparse_features,
+            sequence_features=sequence_features,
+            target=target,
+            task=task,
+            device=device,
+            embedding_l1_reg=embedding_l1_reg,
+            dense_l1_reg=dense_l1_reg,
+            embedding_l2_reg=embedding_l2_reg,
+            dense_l2_reg=dense_l2_reg,
+            early_stop_patience=20,
+            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})")
+            
+        # All features
+        self.all_features = dense_features + sparse_features + sequence_features
+
+        # Embedding layer
+        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])
+        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]
+        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)
+            )
+            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
+        )
+
+        # Register regularization weights
+        self._register_regularization_weights(
+            embedding_attr='embedding',
+            include_modules=['experts', 'gates', 'towers']
+        )
+
+        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]
+        
+        # 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]
+            
+            # 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)