nextrec 0.2.4__py3-none-any.whl → 0.2.5__py3-none-any.whl

nextrec/__version__.py

@@ -1 +1 @@
-__version__ = "0.2.4"
+__version__ = "0.2.5"

nextrec/basic/features.py

@@ -93,6 +93,8 @@ class FeatureSpecMixin:
         dense_features: Sequence[DenseFeature] | None = None,
         sparse_features: Sequence[SparseFeature] | None = None,
         sequence_features: Sequence[SequenceFeature] | None = None,
+        target: str | Sequence[str] | None = None,
+        id_columns: str | Sequence[str] | None = None,
     ) -> None:
         self.dense_features: List[DenseFeature] = list(dense_features) if dense_features else []
         self.sparse_features: List[SparseFeature] = list(sparse_features) if sparse_features else []
@@ -100,8 +102,10 @@ class FeatureSpecMixin:
 
         self.all_features = self.dense_features + self.sparse_features + self.sequence_features
         self.feature_names = [feat.name for feat in self.all_features]
+        self.target_columns = self._normalize_to_list(target)
+        self.id_columns = self._normalize_to_list(id_columns)
 
-    def _set_target_config(
+    def _set_target_id_config(
         self,
         target: str | Sequence[str] | None = None,
         id_columns: str | Sequence[str] | None = None,

nextrec/basic/layers.py

@@ -49,10 +49,6 @@ __all__ = [
 
 
 class PredictionLayer(nn.Module):
-    _CLASSIFICATION_TASKS = {"classification", "binary", "ctr", "ranking", "match", "matching"}
-    _REGRESSION_TASKS = {"regression", "continuous"}
-    _MULTICLASS_TASKS = {"multiclass", "softmax"}
-
     def __init__(
         self,
         task_type: Union[str, Sequence[str]] = "binary",
@@ -131,11 +127,11 @@ class PredictionLayer(nn.Module):
 
     def _get_activation(self, task_type: str):
         task = task_type.lower()
-        if task in self._CLASSIFICATION_TASKS:
+        if task in ['binary','multiclass']:
             return torch.sigmoid
-        if task in self._REGRESSION_TASKS:
+        if task in ['regression']:
             return lambda x: x
-        if task in self._MULTICLASS_TASKS:
+        if task in ['multiclass']:
             return lambda x: torch.softmax(x, dim=-1)
         raise ValueError(f"Unsupported task_type '{task_type}'.")