Adding ContrastiveOutput

merlin/models/torch/outputs/classification.py

@@ -22,6 +22,7 @@
 
 import merlin.dtypes as md
 from merlin.models.torch import schema
+from merlin.models.torch.batch import Batch
 from merlin.models.torch.inputs.embedding import EmbeddingTable
 from merlin.models.torch.outputs.base import ModelOutput
 from merlin.schema import ColumnSchema, Schema, Tags
@@ -288,6 +289,12 @@ def embeddings(self) -> nn.Parameter:
  """
  return self.linear.weight.t()
 
+ def should_apply_contrastive(self, batch: Optional[Batch]) -> bool:
+ if batch is not None and batch.targets and self.training:
+ return True
+
+ return False
+
 
 class EmbeddingTablePrediction(nn.Module):
  """Prediction of a categorical feature using weight-sharing [1] with an embedding table.
@@ -318,6 +325,7 @@ def __init__(self, table: EmbeddingTable, selection: Optional[schema.Selection]
  self.num_classes = table.num_embeddings
  self.col_schema = table.input_schema.first
  self.col_name = self.col_schema.name
+ self.target_name = self.col_name
  self.bias = nn.Parameter(
  torch.zeros(self.num_classes, dtype=torch.float32, device=self.embeddings().device)
  )
@@ -368,6 +376,7 @@ def add_selection(self, selection: schema.Selection):
  self.col_name = self.col_schema.name
  self.num_classes = self.col_schema.int_domain.max + 1
  self.output_schema = categorical_output_schema(self.col_schema, self.num_classes)
+ self.target_name = self.col_name
 
  return self
 
@@ -399,6 +408,12 @@ def embedding_lookup(self, inputs: torch.Tensor) -> torch.Tensor:
  """
  return self.table({self.col_name: inputs})[self.col_name]
 
+ def should_apply_contrastive(self, batch: Optional[Batch]) -> bool:
+ if batch is not None and batch.targets and self.training:
+ return True
+
+ return False
+
 
 def categorical_output_schema(target: ColumnSchema, num_classes: int) -> Schema:
  """Return the output schema given the target column schema."""