No public description

PiperOrigin-RevId: 595481514

Author: nicolas-aagnes

official/recommendation/uplift/keys.py

@@ -20,9 +20,11 @@ class StrEnum(str, enum.Enum):
   """An Enum represented by a string."""
 
 
-class TwoTowerPredictionKeys(StrEnum):
-  """Keys for prediction tensors."""
+class TwoTowerOutputKeys(StrEnum):
+  """Keys for training and inference output tensors."""
 
-  UPLIFT = "uplift_predictions"
-  CONTROL = "control_predictions"
-  TREATMENT = "treatment_predictions"
+  CONTROL_PREDICTIONS = "control_predictions"
+  TREATMENT_PREDICTIONS = "treatment_predictions"
+  UPLIFT_PREDICTIONS = "uplift_predictions"
+  IS_TREATMENT = "is_treatment"
+  TRUE_LOGITS = "true_logits"

official/recommendation/uplift/models/two_tower_uplift_model.py

@@ -96,9 +96,13 @@ def predict_step(self, data) -> dict[str, tf.Tensor]:
     outputs = super().predict_step(data)
 
     return {
-        keys.TwoTowerPredictionKeys.CONTROL: outputs.control_predictions,
-        keys.TwoTowerPredictionKeys.TREATMENT: outputs.treatment_predictions,
-        keys.TwoTowerPredictionKeys.UPLIFT: outputs.uplift,
+        keys.TwoTowerOutputKeys.CONTROL_PREDICTIONS: (
+            outputs.control_predictions
+        ),
+        keys.TwoTowerOutputKeys.TREATMENT_PREDICTIONS: (
+            outputs.treatment_predictions
+        ),
+        keys.TwoTowerOutputKeys.UPLIFT_PREDICTIONS: outputs.uplift,
     }
 
   def get_config(self) -> Mapping[str, Any]:

official/recommendation/uplift/models/two_tower_uplift_model_test.py

@@ -121,9 +121,9 @@ def test_model_training_and_inference(self):
 
     # Test model inference predictions.
     expected_predictions = {
-        keys.TwoTowerPredictionKeys.CONTROL: tf.zeros((10, 1)),
-        keys.TwoTowerPredictionKeys.TREATMENT: 3 * tf.ones((10, 1)),
-        keys.TwoTowerPredictionKeys.UPLIFT: 3 * tf.ones((10, 1)),
+        keys.TwoTowerOutputKeys.CONTROL_PREDICTIONS: tf.zeros((10, 1)),
+        keys.TwoTowerOutputKeys.TREATMENT_PREDICTIONS: 3 * tf.ones((10, 1)),
+        keys.TwoTowerOutputKeys.UPLIFT_PREDICTIONS: 3 * tf.ones((10, 1)),
     }
     self.assertAllClose(expected_predictions, model.predict(dataset))
 
@@ -132,31 +132,39 @@ def test_model_training_and_inference(self):
           "testcase_name": "identity",
           "inverse_link_fn": tf.identity,
           "expected_predictions": {
-              keys.TwoTowerPredictionKeys.CONTROL: (
+              keys.TwoTowerOutputKeys.CONTROL_PREDICTIONS: (
                   tf.ones((3, 1)) * -1.0
               ),  # 1 - 2 = -1
-              keys.TwoTowerPredictionKeys.TREATMENT: (
+              keys.TwoTowerOutputKeys.TREATMENT_PREDICTIONS: (
                   tf.ones((3, 1)) * 4.0
               ),  # 1 + 3 = 4
-              keys.TwoTowerPredictionKeys.UPLIFT: tf.ones((3, 1)) * 5.0,
+              keys.TwoTowerOutputKeys.UPLIFT_PREDICTIONS: tf.ones((3, 1)) * 5.0,
           },
       },
       {
           "testcase_name": "abs",
           "inverse_link_fn": tf.math.abs,
           "expected_predictions": {
-              keys.TwoTowerPredictionKeys.CONTROL: tf.ones((3, 1)) * 1.0,
-              keys.TwoTowerPredictionKeys.TREATMENT: tf.ones((3, 1)) * 4.0,
-              keys.TwoTowerPredictionKeys.UPLIFT: tf.ones((3, 1)) * 3.0,
+              keys.TwoTowerOutputKeys.CONTROL_PREDICTIONS: (
+                  tf.ones((3, 1)) * 1.0
+              ),
+              keys.TwoTowerOutputKeys.TREATMENT_PREDICTIONS: (
+                  tf.ones((3, 1)) * 4.0
+              ),
+              keys.TwoTowerOutputKeys.UPLIFT_PREDICTIONS: tf.ones((3, 1)) * 3.0,
           },
       },
       {
           "testcase_name": "relu",
           "inverse_link_fn": tf_keras.activations.relu,
           "expected_predictions": {
-              keys.TwoTowerPredictionKeys.CONTROL: tf.ones((3, 1)) * 0.0,
-              keys.TwoTowerPredictionKeys.TREATMENT: tf.ones((3, 1)) * 4.0,
-              keys.TwoTowerPredictionKeys.UPLIFT: tf.ones((3, 1)) * 4.0,
+              keys.TwoTowerOutputKeys.CONTROL_PREDICTIONS: (
+                  tf.ones((3, 1)) * 0.0
+              ),
+              keys.TwoTowerOutputKeys.TREATMENT_PREDICTIONS: (
+                  tf.ones((3, 1)) * 4.0
+              ),
+              keys.TwoTowerOutputKeys.UPLIFT_PREDICTIONS: tf.ones((3, 1)) * 4.0,
           },
       },
   )