nextrec 0.3.5__py3-none-any.whl → 0.4.1__py3-none-any.whl

nextrec/__init__.py

@@ -1,33 +1,3 @@
-"""
-NextRec - A Unified Deep Learning Framework for Recommender Systems
-===================================================================
-
-NextRec provides a comprehensive suite of recommendation models including:
-- Ranking models (CTR prediction)
-- Matching models (retrieval)
-- Multi-task learning models
-- Generative recommendation models
-
-Quick Start
------------
->>> from nextrec.basic.features import DenseFeature, SparseFeature
->>> from nextrec.models.ranking.deepfm import DeepFM
->>> 
->>> # Define features
->>> dense_features = [DenseFeature('age')]
->>> sparse_features = [SparseFeature('category', vocab_size=100, embedding_dim=16)]
->>> 
->>> # Build model
->>> model = DeepFM(
-...     dense_features=dense_features,
-...     sparse_features=sparse_features,
-...     targets=['label']
-... )
->>> 
->>> # Train model
->>> model.fit(train_data=df_train, valid_data=df_valid)
-"""
-
 from nextrec.__version__ import __version__
 
 __all__ = [

nextrec/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.5"
+__version__ = "0.4.1"

nextrec/basic/layers.py

@@ -80,9 +80,7 @@ class PredictionLayer(nn.Module):
             else:
                 raise ValueError(f"[PredictionLayer Error]: Unsupported task_type '{task_type}'.")
             outputs.append(activation(task_logits))
-        result = torch.cat(outputs, dim=-1)
-        if result.shape[-1] == 1:
-            result = result.squeeze(-1)
+        result = torch.cat(outputs, dim=-1)  # single: (N,1), multi-task/multi-class: (N,total_dim)
         return result
 
 class EmbeddingLayer(nn.Module):
@@ -235,14 +233,28 @@ class InputMask(nn.Module):
         super().__init__()
 
     def forward(self, x: dict[str, torch.Tensor], feature: SequenceFeature, seq_tensor: torch.Tensor | None = None):
-        values = seq_tensor if seq_tensor is not None else x[feature.name]
-        if feature.padding_idx is not None:
-            mask = (values.long() != feature.padding_idx)
+        if seq_tensor is not None:
+            values = seq_tensor
         else:
-            mask = (values.long() != 0)
+            values = x[feature.name]
+        values = values.long()
+        padding_idx = feature.padding_idx if feature.padding_idx is not None else 0
+        mask = (values != padding_idx) 
+
         if mask.dim() == 1:
-            mask = mask.unsqueeze(-1)
-        return mask.unsqueeze(1).float()
+            # [B] -> [B, 1, 1]
+            mask = mask.unsqueeze(1).unsqueeze(2)
+        elif mask.dim() == 2:
+            # [B, L] -> [B, 1, L]
+            mask = mask.unsqueeze(1)
+        elif mask.dim() == 3:
+            # [B, 1, L]
+            # [B, L, 1]  -> [B, L] -> [B, 1, L]
+            if mask.size(1) != 1 and mask.size(2) == 1:
+                mask = mask.squeeze(-1).unsqueeze(1)
+        else:
+            raise ValueError(f"InputMask only supports 1D/2D/3D tensors, got shape {values.shape}")
+        return mask.float()
 
 class LR(nn.Module):
     def __init__(
@@ -299,20 +311,25 @@ class MLP(nn.Module):
         super().__init__()
         if dims is None:
             dims = []
-        layers = list()
+        layers = []
+        current_dim = input_dim
+
         for i_dim in dims:
-            layers.append(nn.Linear(input_dim, i_dim))
+            layers.append(nn.Linear(current_dim, i_dim))
             layers.append(nn.BatchNorm1d(i_dim))
             layers.append(activation_layer(activation))
             layers.append(nn.Dropout(p=dropout))
-            input_dim = i_dim
+            current_dim = i_dim
+            
         if output_layer:
-            layers.append(nn.Linear(input_dim, 1))
+            layers.append(nn.Linear(current_dim, 1))
+            self.output_dim = 1
+        else:
+            self.output_dim = current_dim  
         self.mlp = nn.Sequential(*layers)
-
     def forward(self, x):
         return self.mlp(x)
-
+    
 class FM(nn.Module):
     def __init__(self, reduce_sum: bool = True):
         super().__init__()

nextrec/basic/loggers.py

@@ -99,7 +99,7 @@ def setup_logger(session_id: str | os.PathLike | None = None):
     session = create_session(str(session_id) if session_id is not None else None)
     log_dir = session.logs_dir
     log_dir.mkdir(parents=True, exist_ok=True)
-    log_file = log_dir / f"{session.experiment_id}.log"
+    log_file = log_dir / f"{session.log_basename}.log"
 
     console_format = '%(message)s'
     file_format = '%(asctime)s - %(levelname)s - %(message)s'