Add CategoricalOutput

merlin/models/torch/__init__.py

@@ -25,7 +25,12 @@
 from merlin.models.torch.models.base import Model
 from merlin.models.torch.models.ranking import DLRMModel
 from merlin.models.torch.outputs.base import ModelOutput
-from merlin.models.torch.outputs.classification import BinaryOutput
+from merlin.models.torch.outputs.classification import (
+    BinaryOutput,
+    CategoricalOutput,
+    CategoricalTarget,
+    EmbeddingTablePrediction,
+)
 from merlin.models.torch.outputs.regression import RegressionOutput
 from merlin.models.torch.outputs.tabular import TabularOutputBlock
 from merlin.models.torch.router import RouterBlock
@@ -40,6 +45,7 @@
     "Batch",
     "BinaryOutput",
     "Block",
+    "DLRMBlock",
     "MLPBlock",
     "Model",
     "EmbeddingTable",
@@ -60,6 +66,9 @@
     "Concat",
     "Stack",
     "schema",
+    "CategoricalOutput",
+    "CategoricalTarget",
+    "EmbeddingTablePrediction",
     "input_schema",
     "output_schema",
     "feature_schema",

merlin/models/torch/inputs/embedding.py

@@ -374,6 +374,14 @@ def update_feature(self, col_schema: ColumnSchema) -> "EmbeddingTable":
 
         return self
 
+    def feature_weights(self, name: str):
+        if name not in self.domains:
+            raise ValueError(f"{name} not found in table: {self}")
+
+        domain = self.domains[name]
+
+        return self.table.weight[int(domain.min) : int(domain.max)]
+
     def select(self, selection: Selection) -> Selectable:
         selected = select(self.input_schema, selection)
 

merlin/models/torch/outputs/base.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import inspect
 from copy import deepcopy
 from typing import Optional, Sequence
 
@@ -21,7 +22,7 @@
 from torchmetrics import Metric
 
 from merlin.models.torch.block import Block
-from merlin.schema import ColumnSchema, Schema
+from merlin.models.torch.transforms.bias import LogitsTemperatureScaler
 
 
 class ModelOutput(Block):
@@ -47,44 +48,34 @@ class ModelOutput(Block):
 
     Parameters
     ----------
-    schema: Optional[ColumnSchema]
-        The schema defining the column properties.
     loss: nn.Module
         The loss function used for training.
     metrics: Sequence[Metric]
         The metrics used for evaluation.
+    logits_temperature: float, optional
+        Parameter used to reduce model overconfidence, so that logits / T.
+        by default 1.0
     name: Optional[str]
         The name of the model output.
     """
 
     def __init__(
         self,
         *module: nn.Module,
-        schema: Optional[ColumnSchema] = None,
         loss: Optional[nn.Module] = None,
-        metrics: Sequence[Metric] = (),
+        metrics: Optional[Sequence[Metric]] = None,
+        logits_temperature: float = 1.0,
         name: Optional[str] = None,
     ):
         """Initializes a ModelOutput object."""
         super().__init__(*module, name=name)
 
         self.loss = loss
         self.metrics = metrics
-        self.output_schema: Schema = Schema()
 
-        if schema:
-            self.setup_schema(schema)
         self.create_target_buffer()
-
-    def setup_schema(self, schema: Optional[ColumnSchema]):
-        """Set up the schema for the output.
-
-        Parameters
-        ----------
-        schema: ColumnSchema or None
-            The schema defining the column properties.
-        """
-        self.output_schema = Schema([schema])
+        if logits_temperature != 1.0:
+            self.append(LogitsTemperatureScaler(logits_temperature))
 
     def create_target_buffer(self):
         self.register_buffer("target", torch.zeros(1, dtype=torch.float32))
@@ -103,18 +94,24 @@ def eval(self):
         return self.train(False)
 
     def copy(self):
-        metrics = self.metrics
+        metrics = deepcopy(self.metrics)
         self.metrics = []
 
         output = deepcopy(self)
 
         copied_metrics = []
         for metric in metrics:
-            m = metric.__class__()
+            params = inspect.signature(metric.__class__.__init__).parameters
+            kwargs = {}
+            for arg_name, arg_value in params.items():
+                if arg_name in metric.__dict__:
+                    kwargs[arg_name] = metric.__dict__[arg_name]
+            m = metric.__class__(**kwargs)
             m.load_state_dict(metric.state_dict())
             copied_metrics.append(m)
 
         self.metrics = metrics
         output.metrics = copied_metrics
+        output.loss = deepcopy(self.loss)
 
         return output