Explicitly set TF_USE_LEGACY_KERAS=1 when importing Keras.

Tensorflow Transform is only compatible with Keras 2.

PiperOrigin-RevId: 687967427

Author: tf-transform-team

tensorflow_transform/keras_lib.py

@@ -12,8 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Imports keras 2."""
+import os
+
+from absl import logging
 import tensorflow as tf
 
+if 'TF_USE_LEGACY_KERAS' not in os.environ:
+  # Make sure we are using Keras 2.
+  os.environ['TF_USE_LEGACY_KERAS'] = '1'
+elif os.environ['TF_USE_LEGACY_KERAS'] not in ('true', 'True', '1'):
+  logging.warning(
+      'TF_USE_LEGACY_KERAS is set to %s, which will not use Keras 2. Tensorflow'
+      ' Transform is only compatible with Keras 2. Please set'
+      ' TF_USE_LEGACY_KERAS=1.',
+      os.environ['TF_USE_LEGACY_KERAS'],
+  )
+
 version_fn = getattr(tf.keras, 'version', None)
 if version_fn and version_fn().startswith('3.'):
   # `tf.keras` points to `keras 3`, so use `tf_keras` package