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

nextrec/models/ranking/dien.py

@@ -12,14 +12,7 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.model import BaseModel
-from nextrec.basic.layers import (
-    EmbeddingLayer,
-    MLP,
-    AttentionPoolingLayer,
-    DynamicGRU,
-    AUGRU,
-    PredictionLayer,
-)
+from nextrec.basic.layers import EmbeddingLayer, MLP, AttentionPoolingLayer, DynamicGRU, AUGRU, PredictionLayer
 from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
 
 
@@ -31,29 +24,27 @@ class DIEN(BaseModel):
     @property
     def task_type(self):
         return "binary"
-
-    def __init__(
-        self,
-        dense_features: list[DenseFeature],
-        sparse_features: list[SparseFeature],
-        sequence_features: list[SequenceFeature],
-        mlp_params: dict,
-        gru_hidden_size: int = 64,
-        attention_hidden_units: list[int] = [80, 40],
-        attention_activation: str = "sigmoid",
-        use_negsampling: bool = False,
-        target: list[str] = [],
-        optimizer: str = "adam",
-        optimizer_params: dict = {},
-        loss: 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],
+                 mlp_params: dict,
+                 gru_hidden_size: int = 64,
+                 attention_hidden_units: list[int] = [80, 40],
+                 attention_activation: str = 'sigmoid',
+                 use_negsampling: bool = False,
+                 target: list[str] = [],
+                 optimizer: str = "adam",
+                 optimizer_params: dict = {},
+                 loss: 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(DIEN, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -66,185 +57,158 @@ class DIEN(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"
-
+        
         self.use_negsampling = use_negsampling
-
+        
         # Features classification
         if len(sequence_features) == 0:
-            raise ValueError(
-                "DIEN requires at least one sequence feature for user behavior history"
-            )
-
+            raise ValueError("DIEN requires at least one sequence feature for user behavior history")
+        
         self.behavior_feature = sequence_features[0]  # User behavior sequence
-        self.candidate_feature = (
-            sparse_features[-1] if sparse_features else None
-        )  # Candidate item
-
-        self.other_sparse_features = (
-            sparse_features[:-1] if self.candidate_feature else sparse_features
-        )
+        self.candidate_feature = sparse_features[-1] if sparse_features else None  # Candidate item
+        
+        self.other_sparse_features = sparse_features[:-1] if self.candidate_feature else sparse_features
         self.dense_features_list = dense_features
-
+        
         # All features for embedding
         self.all_features = dense_features + sparse_features + sequence_features
 
         # Embedding layer
         self.embedding = EmbeddingLayer(features=self.all_features)
-
+        
         behavior_emb_dim = self.behavior_feature.embedding_dim
         self.candidate_proj = None
-        if (
-            self.candidate_feature is not None
-            and self.candidate_feature.embedding_dim != gru_hidden_size
-        ):
-            self.candidate_proj = nn.Linear(
-                self.candidate_feature.embedding_dim, gru_hidden_size
-            )
-
+        if self.candidate_feature is not None and self.candidate_feature.embedding_dim != gru_hidden_size:
+            self.candidate_proj = nn.Linear(self.candidate_feature.embedding_dim, gru_hidden_size)
+        
         # Interest Extractor Layer (GRU)
         self.interest_extractor = DynamicGRU(
-            input_size=behavior_emb_dim, hidden_size=gru_hidden_size
+            input_size=behavior_emb_dim,
+            hidden_size=gru_hidden_size
         )
-
+        
         # Attention layer for computing attention scores
         self.attention_layer = AttentionPoolingLayer(
             embedding_dim=gru_hidden_size,
             hidden_units=attention_hidden_units,
             activation=attention_activation,
-            use_softmax=False,  # We'll use scores directly for AUGRU
+            use_softmax=False  # We'll use scores directly for AUGRU
         )
-
+        
         # Interest Evolution Layer (AUGRU)
         self.interest_evolution = AUGRU(
-            input_size=gru_hidden_size, hidden_size=gru_hidden_size
+            input_size=gru_hidden_size,
+            hidden_size=gru_hidden_size
         )
-
+        
         # Calculate MLP input dimension
         mlp_input_dim = 0
         if self.candidate_feature:
             mlp_input_dim += self.candidate_feature.embedding_dim
         mlp_input_dim += gru_hidden_size  # final interest state
         mlp_input_dim += sum([f.embedding_dim for f in self.other_sparse_features])
-        mlp_input_dim += sum(
-            [getattr(f, "embedding_dim", 1) or 1 for f in dense_features]
-        )
-
+        mlp_input_dim += sum([getattr(f, "embedding_dim", 1) or 1 for f in dense_features])
+        
         # MLP for final prediction
         self.mlp = MLP(input_dim=mlp_input_dim, **mlp_params)
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
         # Register regularization weights
         self._register_regularization_weights(
-            embedding_attr="embedding",
-            include_modules=[
-                "interest_extractor",
-                "interest_evolution",
-                "attention_layer",
-                "mlp",
-                "candidate_proj",
-            ],
+            embedding_attr='embedding',
+            include_modules=['interest_extractor', 'interest_evolution', 'attention_layer', 'mlp', 'candidate_proj']
         )
 
-        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 candidate item embedding
         if self.candidate_feature:
-            candidate_emb = self.embedding.embed_dict[
-                self.candidate_feature.embedding_name
-            ](
+            candidate_emb = self.embedding.embed_dict[self.candidate_feature.embedding_name](
                 x[self.candidate_feature.name].long()
             )  # [B, emb_dim]
         else:
             raise ValueError("DIEN requires a candidate item feature")
-
+        
         # Get behavior sequence embedding
         behavior_seq = x[self.behavior_feature.name].long()  # [B, seq_len]
         behavior_emb = self.embedding.embed_dict[self.behavior_feature.embedding_name](
             behavior_seq
         )  # [B, seq_len, emb_dim]
-
+        
         # Create mask for padding
         if self.behavior_feature.padding_idx is not None:
-            mask = (
-                (behavior_seq != self.behavior_feature.padding_idx)
-                .unsqueeze(-1)
-                .float()
-            )
+            mask = (behavior_seq != self.behavior_feature.padding_idx).unsqueeze(-1).float()
         else:
             mask = (behavior_seq != 0).unsqueeze(-1).float()
-
+        
         # Step 1: Interest Extractor (GRU)
-        interest_states, _ = self.interest_extractor(
-            behavior_emb
-        )  # [B, seq_len, hidden_size]
-
+        interest_states, _ = self.interest_extractor(behavior_emb)  # [B, seq_len, hidden_size]
+        
         # Step 2: Compute attention scores for each time step
         batch_size, seq_len, hidden_size = interest_states.shape
-
+        
         # Project candidate to hidden_size if necessary (defined in __init__)
         if self.candidate_proj is not None:
             candidate_for_attention = self.candidate_proj(candidate_emb)
         else:
             candidate_for_attention = candidate_emb
-
+        
         # Compute attention scores for AUGRU
         attention_scores = []
         for t in range(seq_len):
             score = self.attention_layer.attention_net(
-                torch.cat(
-                    [
-                        candidate_for_attention,
-                        interest_states[:, t, :],
-                        candidate_for_attention - interest_states[:, t, :],
-                        candidate_for_attention * interest_states[:, t, :],
-                    ],
-                    dim=-1,
-                )
+                torch.cat([
+                    candidate_for_attention,
+                    interest_states[:, t, :],
+                    candidate_for_attention - interest_states[:, t, :],
+                    candidate_for_attention * interest_states[:, t, :]
+                ], dim=-1)
             )  # [B, 1]
             attention_scores.append(score)
-
-        attention_scores = torch.cat(attention_scores, dim=1).unsqueeze(
-            -1
-        )  # [B, seq_len, 1]
+        
+        attention_scores = torch.cat(attention_scores, dim=1).unsqueeze(-1)  # [B, seq_len, 1]
         attention_scores = torch.sigmoid(attention_scores)  # Normalize to [0, 1]
-
+        
         # Apply mask to attention scores
         attention_scores = attention_scores * mask
-
+        
         # Step 3: Interest Evolution (AUGRU)
         final_states, final_interest = self.interest_evolution(
-            interest_states, attention_scores
+            interest_states,
+            attention_scores
         )  # final_interest: [B, hidden_size]
-
+        
         # Get other features
         other_embeddings = []
         other_embeddings.append(candidate_emb)
         other_embeddings.append(final_interest)
-
+        
         # Other sparse features
         for feat in self.other_sparse_features:
-            feat_emb = self.embedding.embed_dict[feat.embedding_name](
-                x[feat.name].long()
-            )
+            feat_emb = self.embedding.embed_dict[feat.embedding_name](x[feat.name].long())
             other_embeddings.append(feat_emb)
-
+        
         # Dense features
         for feat in self.dense_features_list:
             val = x[feat.name].float()
             if val.dim() == 1:
                 val = val.unsqueeze(1)
             other_embeddings.append(val)
-
+        
         # Concatenate all features
         concat_input = torch.cat(other_embeddings, dim=-1)  # [B, total_dim]
-
+        
         # MLP prediction
         y = self.mlp(concat_input)  # [B, 1]
         return self.prediction_layer(y)

nextrec/models/ranking/din.py

@@ -12,12 +12,7 @@ import torch
 import torch.nn as nn
 
 from nextrec.basic.model import BaseModel
-from nextrec.basic.layers import (
-    EmbeddingLayer,
-    MLP,
-    AttentionPoolingLayer,
-    PredictionLayer,
-)
+from nextrec.basic.layers import EmbeddingLayer, MLP, AttentionPoolingLayer, PredictionLayer
 from nextrec.basic.features import DenseFeature, SparseFeature, SequenceFeature
 
 
@@ -29,28 +24,26 @@ class DIN(BaseModel):
     @property
     def task_type(self):
         return "binary"
-
-    def __init__(
-        self,
-        dense_features: list[DenseFeature],
-        sparse_features: list[SparseFeature],
-        sequence_features: list[SequenceFeature],
-        mlp_params: dict,
-        attention_hidden_units: list[int] = [80, 40],
-        attention_activation: str = "sigmoid",
-        attention_use_softmax: bool = True,
-        target: list[str] = [],
-        optimizer: str = "adam",
-        optimizer_params: dict = {},
-        loss: 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],
+                 mlp_params: dict,
+                 attention_hidden_units: list[int] = [80, 40],
+                 attention_activation: str = 'sigmoid',
+                 attention_use_softmax: bool = True,
+                 target: list[str] = [],
+                 optimizer: str = "adam",
+                 optimizer_params: dict = {},
+                 loss: 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(DIN, self).__init__(
             dense_features=dense_features,
             sparse_features=sparse_features,
@@ -63,54 +56,43 @@ class DIN(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"
-
+        
         # Features classification
         # DIN requires: candidate item + user behavior sequence + other features
         if len(sequence_features) == 0:
-            raise ValueError(
-                "DIN requires at least one sequence feature for user behavior history"
-            )
-
+            raise ValueError("DIN requires at least one sequence feature for user behavior history")
+        
         self.behavior_feature = sequence_features[0]  # User behavior sequence
-        self.candidate_feature = (
-            sparse_features[-1] if sparse_features else None
-        )  # Candidate item
-
+        self.candidate_feature = sparse_features[-1] if sparse_features else None  # Candidate item
+        
         # Other features (excluding behavior sequence in final concatenation)
-        self.other_sparse_features = (
-            sparse_features[:-1] if self.candidate_feature else sparse_features
-        )
+        self.other_sparse_features = sparse_features[:-1] if self.candidate_feature else sparse_features
         self.dense_features_list = dense_features
-
+        
         # All features for embedding
         self.all_features = dense_features + sparse_features + sequence_features
 
         # Embedding layer
         self.embedding = EmbeddingLayer(features=self.all_features)
-
+        
         # Attention layer for behavior sequence
         behavior_emb_dim = self.behavior_feature.embedding_dim
         self.candidate_attention_proj = None
-        if (
-            self.candidate_feature is not None
-            and self.candidate_feature.embedding_dim != behavior_emb_dim
-        ):
-            self.candidate_attention_proj = nn.Linear(
-                self.candidate_feature.embedding_dim, behavior_emb_dim
-            )
+        if self.candidate_feature is not None and self.candidate_feature.embedding_dim != behavior_emb_dim:
+            self.candidate_attention_proj = nn.Linear(self.candidate_feature.embedding_dim, behavior_emb_dim)
         self.attention = AttentionPoolingLayer(
             embedding_dim=behavior_emb_dim,
             hidden_units=attention_hidden_units,
             activation=attention_activation,
-            use_softmax=attention_use_softmax,
+            use_softmax=attention_use_softmax
         )
-
+        
         # Calculate MLP input dimension
         # candidate + attention_pooled_behavior + other_sparse + dense
         mlp_input_dim = 0
@@ -118,88 +100,82 @@ class DIN(BaseModel):
             mlp_input_dim += self.candidate_feature.embedding_dim
         mlp_input_dim += behavior_emb_dim  # attention pooled
         mlp_input_dim += sum([f.embedding_dim for f in self.other_sparse_features])
-        mlp_input_dim += sum(
-            [getattr(f, "embedding_dim", 1) or 1 for f in dense_features]
-        )
-
+        mlp_input_dim += sum([getattr(f, "embedding_dim", 1) or 1 for f in dense_features])
+        
         # MLP for final prediction
         self.mlp = MLP(input_dim=mlp_input_dim, **mlp_params)
         self.prediction_layer = PredictionLayer(task_type=self.task_type)
 
         # Register regularization weights
         self._register_regularization_weights(
-            embedding_attr="embedding",
-            include_modules=["attention", "mlp", "candidate_attention_proj"],
+            embedding_attr='embedding',
+            include_modules=['attention', 'mlp', 'candidate_attention_proj']
         )
 
-        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 candidate item embedding
         if self.candidate_feature:
-            candidate_emb = self.embedding.embed_dict[
-                self.candidate_feature.embedding_name
-            ](
+            candidate_emb = self.embedding.embed_dict[self.candidate_feature.embedding_name](
                 x[self.candidate_feature.name].long()
             )  # [B, emb_dim]
         else:
             candidate_emb = None
-
+        
         # Get behavior sequence embedding
         behavior_seq = x[self.behavior_feature.name].long()  # [B, seq_len]
         behavior_emb = self.embedding.embed_dict[self.behavior_feature.embedding_name](
             behavior_seq
         )  # [B, seq_len, emb_dim]
-
+        
         # Create mask for padding
         if self.behavior_feature.padding_idx is not None:
-            mask = (
-                (behavior_seq != self.behavior_feature.padding_idx)
-                .unsqueeze(-1)
-                .float()
-            )
+            mask = (behavior_seq != self.behavior_feature.padding_idx).unsqueeze(-1).float()
         else:
             mask = (behavior_seq != 0).unsqueeze(-1).float()
-
+        
         # Apply attention pooling
         if candidate_emb is not None:
             candidate_query = candidate_emb
             if self.candidate_attention_proj is not None:
                 candidate_query = self.candidate_attention_proj(candidate_query)
             pooled_behavior = self.attention(
-                query=candidate_query, keys=behavior_emb, mask=mask
+                query=candidate_query,
+                keys=behavior_emb,
+                mask=mask
             )  # [B, emb_dim]
         else:
             # If no candidate, use mean pooling
-            pooled_behavior = torch.sum(behavior_emb * mask, dim=1) / (
-                mask.sum(dim=1) + 1e-9
-            )
-
+            pooled_behavior = torch.sum(behavior_emb * mask, dim=1) / (mask.sum(dim=1) + 1e-9)
+        
         # Get other features
         other_embeddings = []
-
+        
         if candidate_emb is not None:
             other_embeddings.append(candidate_emb)
-
+        
         other_embeddings.append(pooled_behavior)
-
+        
         # Other sparse features
         for feat in self.other_sparse_features:
-            feat_emb = self.embedding.embed_dict[feat.embedding_name](
-                x[feat.name].long()
-            )
+            feat_emb = self.embedding.embed_dict[feat.embedding_name](x[feat.name].long())
             other_embeddings.append(feat_emb)
-
+        
         # Dense features
         for feat in self.dense_features_list:
             val = x[feat.name].float()
             if val.dim() == 1:
                 val = val.unsqueeze(1)
             other_embeddings.append(val)
-
+        
         # Concatenate all features
         concat_input = torch.cat(other_embeddings, dim=-1)  # [B, total_dim]
-
+        
         # MLP prediction
         y = self.mlp(concat_input)  # [B, 1]
         return self.prediction_layer(y)