tf-models-nightly 2.17.0.dev20240410__py2.py3-none-any.whl → 2.17.0.dev20240412__py2.py3-none-any.whl

official/recommendation/uplift/keys.py

@@ -17,7 +17,7 @@
 import enum
 
 
-class TwoTowerOutputKeys(enum.StrEnum):
+class TwoTowerOutputKeys(str, enum.Enum):
   """Keys for training and inference output tensors."""
 
   CONTROL_PREDICTIONS = "control_predictions"

official/recommendation/uplift/layers/heads/two_tower_logits_head.py

@@ -24,7 +24,7 @@ import tensorflow as tf, tf_keras
 
 
 @enum.unique
-class LayeringMethod(enum.StrEnum):
+class LayeringMethod(str, enum.Enum):
   """Layering method between the control and treatment towers."""
 
   # No layering.

official/recommendation/uplift/metrics/loss_metric.py

@@ -19,6 +19,7 @@ from __future__ import annotations
 import inspect
 from typing import Any, Callable
 
+import numpy as np
 import tensorflow as tf, tf_keras
 
 from official.recommendation.uplift import types
@@ -29,9 +30,6 @@ from official.recommendation.uplift.metrics import treatment_sliced_metric
 class LossMetric(tf_keras.metrics.Metric):
   """Computes a loss sliced by treatment group.
 
-  Note that the prediction tensor is expected to be of type
-  `TwoTowerTrainingOutputs`.
-
   Example standalone usage:
 
   >>> sliced_loss = LossMetric(tf_keras.losses.mean_squared_error)
@@ -74,9 +72,10 @@ class LossMetric(tf_keras.metrics.Metric):
         `__call__(y_true: tf,Tensor, y_pred: tf.Tensor, **loss_fn_kwargs)`. Note
         that the `loss_fn_kwargs` will not be passed to the `__call__` method if
         `loss_fn` is a Keras metric.
-      from_logits: Specifies whether the true logits or true predictions should
-        be used from the model outputs to compute the loss. Defaults to using
-        the true logits.
+      from_logits: When `y_pred` is of type `TwoTowerTrainingOutputs`, specifies
+        whether the true logits or true predictions should be used to compute
+        the loss (defaults to using the true logits). Othwerwise, this argument
+        will be ignored if `y_pred` is of type `tf.Tensor`.
       slice_by_treatment: Specifies whether the loss should be sliced by the
         treatment indicator tensor. If `True`, `loss_fn` will be wrapped in a
         `TreatmentSlicedMetric` to report the loss values sliced by the
@@ -129,16 +128,16 @@ class LossMetric(tf_keras.metrics.Metric):
   def update_state(
       self,
       y_true: tf.Tensor,
-      y_pred: types.TwoTowerTrainingOutputs,
+      y_pred: types.TwoTowerTrainingOutputs | tf.Tensor | np.ndarray,
       sample_weight: tf.Tensor | None = None,
   ):
     """Updates the overall, control and treatment losses.
 
     Args:
       y_true: A `tf.Tensor` with the targets.
-      y_pred: Two tower training outputs. The treatment indicator tensor is used
-        to slice the true logits or true predictions into control and treatment
-        losses.
+      y_pred: Model outputs. If of type `TwoTowerTrainingOutputs`, the treatment
+        indicator tensor is used to slice the true logits or true predictions
+        into control and treatment losses.
       sample_weight: Optional sample weight to compute weighted losses. If
         given, the sample weight will also be sliced by the treatment indicator
         tensor to compute the weighted control and treatment losses.
@@ -146,14 +145,23 @@ class LossMetric(tf_keras.metrics.Metric):
     Raises:
       TypeError: if `y_pred` is not of type `TwoTowerTrainingOutputs`.
     """
-    if not isinstance(y_pred, types.TwoTowerTrainingOutputs):
+    if isinstance(y_pred, (tf.Tensor, np.ndarray)):
+      if self._slice_by_treatment:
+        raise ValueError(
+            "`slice_by_treatment` must be False when y_pred is a `tf.Tensor` or"
+            " `np.ndarray`."
+        )
+      pred = y_pred
+    elif isinstance(y_pred, types.TwoTowerTrainingOutputs):
+      pred = (
+          y_pred.true_logits if self._from_logits else y_pred.true_predictions
+      )
+    else:
       raise TypeError(
-          "y_pred must be of type `TwoTowerTrainingOutputs` but got type"
-          f" {type(y_pred)} instead."
+          "y_pred must be of type `TwoTowerTrainingOutputs`, `tf.Tensor` or"
+          f" `np.ndarray` but got type {type(y_pred)} instead."
       )
 
-    pred = y_pred.true_logits if self._from_logits else y_pred.true_predictions
-
     is_treatment = {}
     if self._slice_by_treatment:
       is_treatment["is_treatment"] = y_pred.is_treatment

official/recommendation/uplift/metrics/loss_metric_test.py

@@ -288,6 +288,19 @@ class LossMetricTest(keras_test_case.KerasTestCase, parameterized.TestCase):
     metric(y_true, outputs, sample_weight=sample_weight)
     self.assertEqual(expected_losses, metric.result())
 
+  def test_metric_with_y_pred_tensor(self):
+    y_true = tf.constant([[0], [0], [2], [7]])
+    y_pred = tf.constant([[1], [2], [3], [4]])
+    sample_weight = tf.constant([[0.5], [0.5], [0.7], [1.8]])
+
+    metric = loss_metric.LossMetric(
+        loss_fn=tf_keras.metrics.mae, slice_by_treatment=False
+    )
+    metric(y_true, y_pred, sample_weight)
+
+    expected_loss = np.average([1, 2, 1, 3], weights=[0.5, 0.5, 0.7, 1.8])
+    self.assertAllClose(expected_loss, metric.result())
+
   def test_multiple_update_batches_returns_aggregated_sliced_losses(self):
     metric = loss_metric.LossMetric(
         loss_fn=tf_keras.losses.mean_absolute_error,
@@ -397,8 +410,26 @@ class LossMetricTest(keras_test_case.KerasTestCase, parameterized.TestCase):
     metric = loss_metric.LossMetric(
         tf_keras.metrics.mean_absolute_percentage_error
     )
+    y_true = tf.ones((3, 1))
+    y_pred = types.TwoTowerNetworkOutputs(
+        shared_embedding=tf.ones((3, 5)),
+        control_logits=tf.ones((3, 1)),
+        treatment_logits=tf.ones((3, 1)),
+    )
+    with self.assertRaisesRegex(
+        TypeError,
+        "y_pred must be of type `TwoTowerTrainingOutputs`, `tf.Tensor` or"
+        " `np.ndarray`",
+    ):
+      metric.update_state(y_true=y_true, y_pred=y_pred)
+
+  def test_slice_by_treatment_with_y_pred_tensor_raises_error(self):
+    metric = loss_metric.LossMetric(
+        tf_keras.metrics.mae, slice_by_treatment=True
+    )
     with self.assertRaisesRegex(
-        TypeError, "y_pred must be of type `TwoTowerTrainingOutputs`"
+        ValueError,
+        "`slice_by_treatment` must be False when y_pred is a `tf.Tensor`.",
     ):
       metric.update_state(y_true=tf.ones((3, 1)), y_pred=tf.ones((3, 1)))
 

tensorflow_models/__init__.py

@@ -15,6 +15,7 @@
 """TensorFlow Models Libraries."""
 # pylint: disable=wildcard-import
 from tensorflow_models import nlp
+from tensorflow_models import uplift
 from tensorflow_models import vision
 
 from official import core

tensorflow_models/tensorflow_models_test.py

@@ -35,6 +35,17 @@ class TensorflowModelsTest(tf.test.TestCase):
     _ = tfm.optimization.LinearWarmup(
         after_warmup_lr_sched=0.0, warmup_steps=10, warmup_learning_rate=0.1)
 
+  def testUpliftImports(self):
+    _ = tfm.uplift.keys.TwoTowerOutputKeys.CONTROL_PREDICTIONS
+    _ = tfm.uplift.types.TwoTowerNetworkOutputs(
+        shared_embedding=tf.ones((10, 10)),
+        control_logits=tf.ones((10, 1)),
+        treatment_logits=tf.ones((10, 1)),
+    )
+    _ = tfm.uplift.layers.encoders.concat_features.ConcatFeatures(['feature'])
+    _ = tfm.uplift.metrics.treatment_fraction.TreatmentFraction()
+    _ = tfm.uplift.losses.true_logits_loss.TrueLogitsLoss(tf_keras.losses.mse)
+
 
 if __name__ == '__main__':
   tf.test.main()

tensorflow_models/uplift/__init__.py

@@ -0,0 +1,23 @@
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""TensorFlow Models Uplift Libraries."""
+
+from official.recommendation.uplift import keys
+from official.recommendation.uplift import layers
+from official.recommendation.uplift import losses
+from official.recommendation.uplift import metrics
+from official.recommendation.uplift import models
+from official.recommendation.uplift import types
+from official.recommendation.uplift import utils

{tf_models_nightly-2.17.0.dev20240410.dist-info → tf_models_nightly-2.17.0.dev20240412.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tf-models-nightly
-Version: 2.17.0.dev20240410
+Version: 2.17.0.dev20240412
 Summary: TensorFlow Official Models
 Home-page: https://github.com/tensorflow/models
 Author: Google Inc.

{tf_models_nightly-2.17.0.dev20240410.dist-info → tf_models_nightly-2.17.0.dev20240412.dist-info}/RECORD

@@ -887,7 +887,7 @@ official/recommendation/ranking/data/data_pipeline_multi_hot_test.py,sha256=arLj
 official/recommendation/ranking/data/data_pipeline_test.py,sha256=VRYo7WqURRkM3lbmfctvSZxyH1EzUfqwxR3sy2sZxdc,2345
 official/recommendation/uplift/__init__.py,sha256=_jZilTPWKu-MfMaz1IgBjEW6wqkK3FNZ1QAP4a8my3I,990
 official/recommendation/uplift/keras_test_case.py,sha256=gF5Z2FzXlKAvhuJDdj7PmFj3jsW_ZmUAv_F9Xokvs2M,6156
-official/recommendation/uplift/keys.py,sha256=jXrgCwSl1Ma4sThUeh8GdeQAig4lzdSbgrOwBcjvQSE,1030
+official/recommendation/uplift/keys.py,sha256=7zkxkPIcXceIN5hWm4ATai4h8ymwCj85dU2r8-3XifM,1032
 official/recommendation/uplift/types.py,sha256=OBpMvU4uAOKXGwN1UIPZw6KWJXBXeuZGT3P3VyDkyWc,4769
 official/recommendation/uplift/utils.py,sha256=ZSrzoFJosSdmu7P3yYDo-TER8UQKFimy9ntuXqV4L4c,3501
 official/recommendation/uplift/utils_test.py,sha256=K-pzTe3MoOEjdy5tLNSRpIQoWWH12MjFTPkhEIHfS3A,4610
@@ -896,7 +896,7 @@ official/recommendation/uplift/layers/encoders/__init__.py,sha256=QaQk0BMHJNE8Pc
 official/recommendation/uplift/layers/encoders/concat_features.py,sha256=UNjDXN2C-GBXJaiitMYmKuhIsOiR1ROfJ73vHkuGkoA,4210
 official/recommendation/uplift/layers/encoders/concat_features_test.py,sha256=sIjdoXqQHfQQ5P0yOz_PCuGQHl6j-wKNzaN42PL1jKI,9368
 official/recommendation/uplift/layers/heads/__init__.py,sha256=nbtNxVbIGZq4GLEw9APPwCX3Zhatv5qwc0ddvAaelvU,728
-official/recommendation/uplift/layers/heads/two_tower_logits_head.py,sha256=1HkR2LP3nBbAQrYcbGHsQPO36UuVQ8Pb8CAjWFT0c1Q,6613
+official/recommendation/uplift/layers/heads/two_tower_logits_head.py,sha256=EyeVdiRNe4qk1aEvYEp1tEeGdllUFqtt8saCrK_QI1M,6615
 official/recommendation/uplift/layers/heads/two_tower_logits_head_test.py,sha256=GKk8kYgq_gXcheUfcG1u_i5I5nVh8jRdam9KxvEHsXU,7615
 official/recommendation/uplift/layers/uplift_networks/__init__.py,sha256=eaM75bIO4WZRABzDjTl5bIsR5kKiICpzw-QpgXh7K0A,917
 official/recommendation/uplift/layers/uplift_networks/base_uplift_networks.py,sha256=pGVmdNXyDZR0TA9Htxa07KoUOzE0uEl5IyHEAVRgTk4,1329
@@ -912,8 +912,8 @@ official/recommendation/uplift/metrics/label_mean.py,sha256=ECaes7FZmsksnwySn7jf
 official/recommendation/uplift/metrics/label_mean_test.py,sha256=b_d3lNlpkDm2xKLUkxfiXeQg7pjL8HNx7y9NaYarpV0,7083
 official/recommendation/uplift/metrics/label_variance.py,sha256=9DCl42BJkehxfWD3pSbZnRNvwfhVM6VyHwivGdaU72s,3610
 official/recommendation/uplift/metrics/label_variance_test.py,sha256=k0mdEU1WU53-HIEO5HGtfp1MleifD-h4bZNKtTvM3Ws,7681
-official/recommendation/uplift/metrics/loss_metric.py,sha256=qWMcLk1JmLkru5T9qcqO2UbB_AJQKjQ8uUSAxjR1l58,6883
-official/recommendation/uplift/metrics/loss_metric_test.py,sha256=QIUujZvUud-A2itI2vs2KpzWfLwbq2XDd9H9CFq5D4Y,14929
+official/recommendation/uplift/metrics/loss_metric.py,sha256=owN7A98TCc_UhOURvGfccaoVGOthdHdx1_fawEUGnmw,7289
+official/recommendation/uplift/metrics/loss_metric_test.py,sha256=48rQG8bKFdy0xBFjoOLXKRUlYpCEyAzSmPOFoF7FX94,16021
 official/recommendation/uplift/metrics/metric_configs.py,sha256=Z-r79orE4EycQ5TJ7xdI5LhjOHT3wzChYyDxcxGqLXk,1670
 official/recommendation/uplift/metrics/sliced_metric.py,sha256=O2I2apZK6IfOQK9Q_mgSiTUCnGokczp4e14zrrYNeRU,8564
 official/recommendation/uplift/metrics/sliced_metric_test.py,sha256=dhY41X8lqT_WW04XLjyjDerZwujEBGeTXtxf4NkYThw,11359
@@ -1199,13 +1199,14 @@ orbit/utils/summary_manager.py,sha256=oTdTQYDjH315rbJ__NFALldruHX5lLrwY34IU1KJwN
 orbit/utils/summary_manager_interface.py,sha256=MXXPd3aWqushByK-gh-Jk9A42VC5ETEEm7XY5DCB-9w,2253
 orbit/utils/tpu_summaries.py,sha256=vjfLJlNifFJayLo_nkdRD85XoisVBV2VDuzgZWZud6U,5641
 orbit/utils/tpu_summaries_test.py,sha256=ZRYguds3BcXn1gOXBHVJ5rGRZdeXQLa5w3IbwEBwGS4,4523
-tensorflow_models/__init__.py,sha256=etxw45SHxuwFCRX5qGxGMP83II0JfJulzNl5GSNJvhw,909
-tensorflow_models/tensorflow_models_test.py,sha256=AxUYUdiQn416UR7jg0h6rmv688esvlKDfpyDCIQkF18,1395
+tensorflow_models/__init__.py,sha256=n7s5ZGPDM-RqOQNldPyipjL8aEB4DdAQQUENYhgf_Q8,946
+tensorflow_models/tensorflow_models_test.py,sha256=nc6A9K53OGqF25xN5St8EiWvdVbdanoDB12hb32dzPc,1897
 tensorflow_models/nlp/__init__.py,sha256=4tA5Pf4qaFwT-fIFOpX7x7FHJpnyJT-5UgOeFYTyMlc,807
+tensorflow_models/uplift/__init__.py,sha256=mqfa55gweOdpKoaQyid4A_4u7xw__FcQeSIF0k_pYmI,999
 tensorflow_models/vision/__init__.py,sha256=zBorY_v5xva1uI-qxhZO3Qh-Dii-Suq6wEYh6hKHDfc,833
-tf_models_nightly-2.17.0.dev20240410.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.17.0.dev20240410.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.17.0.dev20240410.dist-info/METADATA,sha256=Bh56e1__Pn_yvpm9QMq4h12nKucvnKAftuk0ElQqn_w,1432
-tf_models_nightly-2.17.0.dev20240410.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.17.0.dev20240410.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.17.0.dev20240410.dist-info/RECORD,,
+tf_models_nightly-2.17.0.dev20240412.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.17.0.dev20240412.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.17.0.dev20240412.dist-info/METADATA,sha256=y_U6M920Hgob94pOpRJWfqistgFzPqBH0Ds0HOZOpoo,1432
+tf_models_nightly-2.17.0.dev20240412.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.17.0.dev20240412.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.17.0.dev20240412.dist-info/RECORD,,