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

nextrec/models/match/mind.py

@@ -6,7 +6,6 @@ Reference:
     [1] Li C, Liu Z, Wu M, et al. Multi-interest network with dynamic routing for recommendation at Tmall[C]
         //Proceedings of the 28th ACM international conference on information and knowledge management. 2019: 2615-2623.
 """
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -21,41 +20,39 @@ class MIND(BaseMatchModel):
     @property
     def model_name(self) -> str:
         return "MIND"
-
+    
     @property
     def support_training_modes(self) -> list[str]:
         """MIND only supports pointwise training mode"""
-        return ["pointwise"]
-
-    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,
-        embedding_dim: int = 64,
-        num_interests: int = 4,
-        capsule_bilinear_type: int = 2,
-        routing_times: int = 3,
-        relu_layer: bool = False,
-        item_dnn_hidden_units: list[int] = [256, 128],
-        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 = "mind",
-    ):
-
+        return ['pointwise']
+    
+    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,
+                 embedding_dim: int = 64,
+                 num_interests: int = 4,
+                 capsule_bilinear_type: int = 2,
+                 routing_times: int = 3,
+                 relu_layer: bool = False,
+                 item_dnn_hidden_units: list[int] = [256, 128],
+                 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 = 'mind'):
+        
         super(MIND, self).__init__(
             user_dense_features=user_dense_features,
             user_sparse_features=user_sparse_features,
@@ -73,9 +70,9 @@ class MIND(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.num_interests = num_interests
         self.item_dnn_hidden_units = item_dnn_hidden_units
@@ -87,20 +84,16 @@ class MIND(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)
-
+            
             if not user_sequence_features or len(user_sequence_features) == 0:
                 raise ValueError("MIND requires at least one user sequence feature")
-
-            seq_max_len = (
-                user_sequence_features[0].max_len
-                if user_sequence_features[0].max_len
-                else 50
-            )
+            
+            seq_max_len = user_sequence_features[0].max_len if user_sequence_features[0].max_len else 50
             seq_embedding_dim = user_sequence_features[0].embedding_dim
-
+            
             # Capsule Network for multi-interest extraction
             self.capsule_network = CapsuleNetwork(
                 embedding_dim=seq_embedding_dim,
@@ -108,17 +101,15 @@ class MIND(BaseMatchModel):
                 bilinear_type=capsule_bilinear_type,
                 interest_num=num_interests,
                 routing_times=routing_times,
-                relu_layer=relu_layer,
+                relu_layer=relu_layer
             )
-
+            
             if seq_embedding_dim != embedding_dim:
-                self.interest_projection = nn.Linear(
-                    seq_embedding_dim, embedding_dim, bias=False
-                )
+                self.interest_projection = nn.Linear(seq_embedding_dim, embedding_dim, bias=False)
                 nn.init.xavier_uniform_(self.interest_projection.weight)
             else:
                 self.interest_projection = None
-
+        
         # Item tower
         item_features = []
         if item_dense_features:
@@ -127,10 +118,10 @@ class MIND(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
@@ -138,7 +129,7 @@ class MIND(BaseMatchModel):
                 item_input_dim += feat.embedding_dim
             for feat in item_sequence_features or []:
                 item_input_dim += feat.embedding_dim
-
+            
             # Item DNN
             if len(item_dnn_hidden_units) > 0:
                 item_dnn_units = item_dnn_hidden_units + [embedding_dim]
@@ -147,25 +138,26 @@ class MIND(BaseMatchModel):
                     dims=item_dnn_units,
                     output_layer=False,
                     dropout=dnn_dropout,
-                    activation=dnn_activation,
+                    activation=dnn_activation
                 )
             else:
                 self.item_dnn = None
-
+        
         self._register_regularization_weights(
-            embedding_attr="user_embedding", include_modules=["capsule_network"]
+            embedding_attr='user_embedding',
+            include_modules=['capsule_network']
         )
         self._register_regularization_weights(
-            embedding_attr="item_embedding",
-            include_modules=["item_dnn"] if self.item_dnn else [],
+            embedding_attr='item_embedding',
+            include_modules=['item_dnn'] if self.item_dnn else []
         )
-
+        
         self.to(device)
-
+    
     def user_tower(self, user_input: dict) -> torch.Tensor:
         """
         User tower with multi-interest extraction
-
+        
         Returns:
             user_interests: [batch_size, num_interests, embedding_dim]
         """
@@ -176,53 +168,43 @@ class MIND(BaseMatchModel):
         seq_emb = embed(seq_input.long())  # [batch_size, seq_len, embedding_dim]
 
         mask = (seq_input != seq_feature.padding_idx).float()  # [batch_size, seq_len]
-
-        multi_interests = self.capsule_network(
-            seq_emb, mask
-        )  # [batch_size, num_interests, seq_embedding_dim]
-
+        
+        multi_interests = self.capsule_network(seq_emb, mask)  # [batch_size, num_interests, seq_embedding_dim]
+        
         if self.interest_projection is not None:
-            multi_interests = self.interest_projection(
-                multi_interests
-            )  # [batch_size, num_interests, embedding_dim]
-
+            multi_interests = self.interest_projection(multi_interests)  # [batch_size, num_interests, embedding_dim]
+        
         # L2 normalization
         multi_interests = F.normalize(multi_interests, p=2, dim=-1)
-
+        
         return multi_interests
-
+    
     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)
-
+        
         if self.item_dnn is not None:
             item_emb = self.item_dnn(item_emb)
-
+        
         # L2 normalization
         item_emb = F.normalize(item_emb, p=2, dim=1)
-
+        
         return item_emb
-
-    def compute_similarity(
-        self, user_emb: torch.Tensor, item_emb: torch.Tensor
-    ) -> torch.Tensor:
+    
+    def compute_similarity(self, user_emb: torch.Tensor, item_emb: torch.Tensor) -> torch.Tensor:
         item_emb_expanded = item_emb.unsqueeze(1)
-
-        if self.similarity_metric == "dot":
+        
+        if self.similarity_metric == 'dot':
             similarities = torch.sum(user_emb * item_emb_expanded, dim=-1)
-        elif self.similarity_metric == "cosine":
+        elif self.similarity_metric == 'cosine':
             similarities = F.cosine_similarity(user_emb, item_emb_expanded, dim=-1)
-        elif self.similarity_metric == "euclidean":
+        elif self.similarity_metric == 'euclidean':
             similarities = -torch.sum((user_emb - item_emb_expanded) ** 2, dim=-1)
         else:
             raise ValueError(f"Unknown similarity metric: {self.similarity_metric}")
 
         max_similarity, _ = torch.max(similarities, dim=1)  # [batch_size]
         max_similarity = max_similarity / self.temperature
-
+        
         return max_similarity

nextrec/models/match/sdm.py

@@ -6,7 +6,6 @@ Reference:
     [1] Ying H, Zhuang F, Zhang F, et al. Sequential recommender system based on hierarchical attention networks[C]
         //IJCAI. 2018: 3926-3932.
 """
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -21,42 +20,40 @@ class SDM(BaseMatchModel):
     @property
     def model_name(self) -> str:
         return "SDM"
-
+    
     @property
     def support_training_modes(self) -> list[str]:
-        return ["pointwise"]
-
-    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,
-        embedding_dim: int = 64,
-        rnn_type: Literal["GRU", "LSTM"] = "GRU",
-        rnn_hidden_size: int = 64,
-        rnn_num_layers: int = 1,
-        rnn_dropout: float = 0.0,
-        use_short_term: bool = True,
-        use_long_term: bool = True,
-        item_dnn_hidden_units: list[int] = [256, 128],
-        dnn_activation: str = "relu",
-        dnn_dropout: float = 0.0,
-        training_mode: Literal["pointwise", "pairwise", "listwise"] = "pointwise",
-        num_negative_samples: int = 4,
-        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 = "sdm",
-    ):
-
+        return ['pointwise']
+    
+    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,
+                 embedding_dim: int = 64,
+                 rnn_type: Literal['GRU', 'LSTM'] = 'GRU',
+                 rnn_hidden_size: int = 64,
+                 rnn_num_layers: int = 1,
+                 rnn_dropout: float = 0.0,
+                 use_short_term: bool = True,
+                 use_long_term: bool = True,
+                 item_dnn_hidden_units: list[int] = [256, 128],
+                 dnn_activation: str = 'relu',
+                 dnn_dropout: float = 0.0,
+                 training_mode: Literal['pointwise', 'pairwise', 'listwise'] = 'pointwise',
+                 num_negative_samples: int = 4,
+                 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 = 'sdm'):
+        
         super(SDM, self).__init__(
             user_dense_features=user_dense_features,
             user_sparse_features=user_sparse_features,
@@ -74,16 +71,16 @@ class SDM(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.rnn_type = rnn_type
         self.rnn_hidden_size = rnn_hidden_size
         self.use_short_term = use_short_term
         self.use_long_term = use_long_term
         self.item_dnn_hidden_units = item_dnn_hidden_units
-
+        
         # User tower
         user_features = []
         if user_dense_features:
@@ -92,54 +89,54 @@ class SDM(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)
-
+            
             if not user_sequence_features or len(user_sequence_features) == 0:
                 raise ValueError("SDM requires at least one user sequence feature")
-
+            
             seq_emb_dim = user_sequence_features[0].embedding_dim
-
-            if rnn_type == "GRU":
+            
+            if rnn_type == 'GRU':
                 self.rnn = nn.GRU(
                     input_size=seq_emb_dim,
                     hidden_size=rnn_hidden_size,
                     num_layers=rnn_num_layers,
                     batch_first=True,
-                    dropout=rnn_dropout if rnn_num_layers > 1 else 0.0,
+                    dropout=rnn_dropout if rnn_num_layers > 1 else 0.0
                 )
-            elif rnn_type == "LSTM":
+            elif rnn_type == 'LSTM':
                 self.rnn = nn.LSTM(
                     input_size=seq_emb_dim,
                     hidden_size=rnn_hidden_size,
                     num_layers=rnn_num_layers,
                     batch_first=True,
-                    dropout=rnn_dropout if rnn_num_layers > 1 else 0.0,
+                    dropout=rnn_dropout if rnn_num_layers > 1 else 0.0
                 )
             else:
                 raise ValueError(f"Unknown RNN type: {rnn_type}")
-
+            
             user_final_dim = 0
             if use_long_term:
-                user_final_dim += rnn_hidden_size
+                user_final_dim += rnn_hidden_size  
             if use_short_term:
-                user_final_dim += seq_emb_dim
-
+                user_final_dim += seq_emb_dim  
+            
             for feat in user_dense_features or []:
                 user_final_dim += 1
             for feat in user_sparse_features or []:
                 user_final_dim += feat.embedding_dim
-
+            
             # User DNN to final embedding
             self.user_dnn = MLP(
                 input_dim=user_final_dim,
                 dims=[rnn_hidden_size * 2, embedding_dim],
                 output_layer=False,
                 dropout=dnn_dropout,
-                activation=dnn_activation,
+                activation=dnn_activation
             )
-
+        
         # Item tower
         item_features = []
         if item_dense_features:
@@ -148,10 +145,10 @@ class SDM(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
@@ -159,7 +156,7 @@ class SDM(BaseMatchModel):
                 item_input_dim += feat.embedding_dim
             for feat in item_sequence_features or []:
                 item_input_dim += feat.embedding_dim
-
+            
             # Item DNN
             if len(item_dnn_hidden_units) > 0:
                 item_dnn_units = item_dnn_hidden_units + [embedding_dim]
@@ -168,58 +165,53 @@ class SDM(BaseMatchModel):
                     dims=item_dnn_units,
                     output_layer=False,
                     dropout=dnn_dropout,
-                    activation=dnn_activation,
+                    activation=dnn_activation
                 )
             else:
                 self.item_dnn = None
-
+        
         self._register_regularization_weights(
-            embedding_attr="user_embedding", include_modules=["rnn", "user_dnn"]
+            embedding_attr='user_embedding',
+            include_modules=['rnn', 'user_dnn']
         )
         self._register_regularization_weights(
-            embedding_attr="item_embedding",
-            include_modules=["item_dnn"] if self.item_dnn else [],
+            embedding_attr='item_embedding',
+            include_modules=['item_dnn'] if self.item_dnn else []
         )
-
+        
         self.to(device)
-
+    
     def user_tower(self, user_input: dict) -> torch.Tensor:
         seq_feature = self.user_sequence_features[0]
         seq_input = user_input[seq_feature.name]
-
+        
         embed = self.user_embedding.embed_dict[seq_feature.embedding_name]
         seq_emb = embed(seq_input.long())  # [batch_size, seq_len, seq_emb_dim]
-
-        if self.rnn_type == "GRU":
-            rnn_output, hidden = self.rnn(
-                seq_emb
-            )  # hidden: [num_layers, batch, hidden_size]
-        elif self.rnn_type == "LSTM":
+        
+        if self.rnn_type == 'GRU':
+            rnn_output, hidden = self.rnn(seq_emb)  # hidden: [num_layers, batch, hidden_size]
+        elif self.rnn_type == 'LSTM':
             rnn_output, (hidden, cell) = self.rnn(seq_emb)
-
+        
         features_list = []
-
+        
         if self.use_long_term:
             if self.rnn.num_layers > 1:
                 long_term = hidden[-1, :, :]  # [batch_size, hidden_size]
             else:
                 long_term = hidden.squeeze(0)  # [batch_size, hidden_size]
             features_list.append(long_term)
-
+        
         if self.use_short_term:
-            mask = (
-                seq_input != seq_feature.padding_idx
-            ).float()  # [batch_size, seq_len]
+            mask = (seq_input != seq_feature.padding_idx).float()  # [batch_size, seq_len]
             seq_lengths = mask.sum(dim=1).long() - 1  # [batch_size]
             seq_lengths = torch.clamp(seq_lengths, min=0)
-
+            
             batch_size = seq_emb.size(0)
             batch_indices = torch.arange(batch_size, device=seq_emb.device)
-            short_term = seq_emb[
-                batch_indices, seq_lengths, :
-            ]  # [batch_size, seq_emb_dim]
+            short_term = seq_emb[batch_indices, seq_lengths, :]  # [batch_size, seq_emb_dim]
             features_list.append(short_term)
-
+        
         if self.user_dense_features:
             dense_features = []
             for feat in self.user_dense_features:
@@ -230,7 +222,7 @@ class SDM(BaseMatchModel):
                     dense_features.append(val)
             if dense_features:
                 features_list.append(torch.cat(dense_features, dim=1))
-
+        
         if self.user_sparse_features:
             sparse_features = []
             for feat in self.user_sparse_features:
@@ -240,26 +232,22 @@ class SDM(BaseMatchModel):
                     sparse_features.append(sparse_emb)
             if sparse_features:
                 features_list.append(torch.cat(sparse_features, dim=1))
-
+        
         user_features = torch.cat(features_list, dim=1)
         user_emb = self.user_dnn(user_features)
         user_emb = F.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)
-
+        
         if self.item_dnn is not None:
             item_emb = self.item_dnn(item_emb)
-
+        
         # L2 normalization
         item_emb = F.normalize(item_emb, p=2, dim=1)
-
+        
         return item_emb