Add tf_keras imports to prepare for Keras 3 (#28588)

* Port core files + ESM (because ESM code is odd)

* Search-replace in modelling code

* Fix up transfo_xl as well

* Fix other core files + tests (still need to add correct import to tests)

* Fix cookiecutter

* make fixup, fix imports in some more core files

* Auto-add imports to tests

* Cleanup, add imports to sagemaker tests

* Use correct exception for importing tf_keras

* Fixes in modeling_tf_utils

* make fixup

* Correct version parsing code

* Ensure the pipeline tests correctly revert to float32 after each test

* Ensure the pipeline tests correctly revert to float32 after each test

* More tf.keras -> keras

* Add dtype cast

* Better imports of tf_keras

* Add a cast for tf.assign, just in case

* Fix callback imports

examples/tensorflow/image-classification/run_image_classification.py

@@ -47,6 +47,7 @@ from transformers import (
     set_seed,
 )
 from transformers.keras_callbacks import KerasMetricCallback
+from transformers.modeling_tf_utils import keras
 from transformers.trainer_utils import get_last_checkpoint, is_main_process
 from transformers.utils import check_min_version, send_example_telemetry
 from transformers.utils.versions import require_version
@@ -363,7 +364,7 @@ def main():
 
     def _train_transforms(image):
         img_size = image_size
-        image = tf.keras.utils.img_to_array(image)
+        image = keras.utils.img_to_array(image)
         image = random_resized_crop(image, size=img_size)
         image = tf.image.random_flip_left_right(image)
         image /= 255.0
@@ -372,7 +373,7 @@ def main():
         return image
 
     def _val_transforms(image):
-        image = tf.keras.utils.img_to_array(image)
+        image = keras.utils.img_to_array(image)
         image = tf.image.resize(image, size=image_size)
         # image = np.array(image) # FIXME - use tf.image function
         image = center_crop(image, size=image_size)

examples/tensorflow/language-modeling-tpu/run_mlm.py

@@ -22,6 +22,7 @@ import os
 import re
 
 import tensorflow as tf
+from packaging.version import parse
 
 from transformers import (
     AutoConfig,
@@ -33,6 +34,19 @@ from transformers import (
 )
 
 
+try:
+    import tf_keras as keras
+except (ModuleNotFoundError, ImportError):
+    import keras
+
+    if parse(keras.__version__).major > 2:
+        raise ValueError(
+            "Your currently installed version of Keras is Keras 3, but this is not yet supported in "
+            "Transformers. Please install the backwards-compatible tf-keras package with "
+            "`pip install tf-keras`."
+        )
+
+
 logger = logging.getLogger(__name__)
 
 AUTO = tf.data.AUTOTUNE
@@ -209,7 +223,7 @@ def main(args):
         strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")
 
     if args.bfloat16:
-        tf.keras.mixed_precision.set_global_policy("mixed_bfloat16")
+        keras.mixed_precision.set_global_policy("mixed_bfloat16")
 
     tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
     config = AutoConfig.from_pretrained(args.pretrained_model_config)

examples/tensorflow/question-answering/run_qa.py

@@ -30,6 +30,7 @@ from typing import Optional
 import evaluate
 import tensorflow as tf
 from datasets import load_dataset
+from packaging.version import parse
 from utils_qa import postprocess_qa_predictions
 
 import transformers
@@ -48,6 +49,19 @@ from transformers import (
 from transformers.utils import CONFIG_NAME, TF2_WEIGHTS_NAME, check_min_version, send_example_telemetry
 
 
+try:
+    import tf_keras as keras
+except (ModuleNotFoundError, ImportError):
+    import keras
+
+    if parse(keras.__version__).major > 2:
+        raise ValueError(
+            "Your currently installed version of Keras is Keras 3, but this is not yet supported in "
+            "Transformers. Please install the backwards-compatible tf-keras package with "
+            "`pip install tf-keras`."
+        )
+
+
 # Will error if the minimal version of Transformers is not installed. Remove at your own risks.
 check_min_version("4.38.0.dev0")
 
@@ -233,7 +247,7 @@ class DataTrainingArguments:
 
 
 # region Helper classes
-class SavePretrainedCallback(tf.keras.callbacks.Callback):
+class SavePretrainedCallback(keras.callbacks.Callback):
     # Hugging Face models have a save_pretrained() method that saves both the weights and the necessary
     # metadata to allow them to be loaded as a pretrained model in future. This is a simple Keras callback
     # that saves the model with this method after each epoch.

examples/tensorflow/test_tensorflow_examples.py

@@ -23,6 +23,20 @@ from unittest import skip
 from unittest.mock import patch
 
 import tensorflow as tf
+from packaging.version import parse
+
+
+try:
+    import tf_keras as keras
+except (ModuleNotFoundError, ImportError):
+    import keras
+
+    if parse(keras.__version__).major > 2:
+        raise ValueError(
+            "Your currently installed version of Keras is Keras 3, but this is not yet supported in "
+            "Transformers. Please install the backwards-compatible tf-keras package with "
+            "`pip install tf-keras`."
+        )
 
 from transformers.testing_utils import TestCasePlus, get_gpu_count, slow
 
@@ -115,7 +129,7 @@ class ExamplesTests(TestCasePlus):
         with patch.object(sys, "argv", testargs):
             run_text_classification.main()
             # Reset the mixed precision policy so we don't break other tests
-            tf.keras.mixed_precision.set_global_policy("float32")
+            keras.mixed_precision.set_global_policy("float32")
             result = get_results(tmp_dir)
             self.assertGreaterEqual(result["eval_accuracy"], 0.75)
 

examples/tensorflow/text-classification/run_text_classification.py

@@ -27,6 +27,7 @@ from typing import Optional
 
 import numpy as np
 from datasets import load_dataset
+from packaging.version import parse
 
 from transformers import (
     AutoConfig,
@@ -46,11 +47,24 @@ os.environ["TF_CPP_MIN_LOG_LEVEL"] = "1"  # Reduce the amount of console output
 import tensorflow as tf  # noqa: E402
 
 
+try:
+    import tf_keras as keras
+except (ModuleNotFoundError, ImportError):
+    import keras
+
+    if parse(keras.__version__).major > 2:
+        raise ValueError(
+            "Your currently installed version of Keras is Keras 3, but this is not yet supported in "
+            "Transformers. Please install the backwards-compatible tf-keras package with "
+            "`pip install tf-keras`."
+        )
+
+
 logger = logging.getLogger(__name__)
 
 
 # region Helper classes
-class SavePretrainedCallback(tf.keras.callbacks.Callback):
+class SavePretrainedCallback(keras.callbacks.Callback):
     # Hugging Face models have a save_pretrained() method that saves both the weights and the necessary
     # metadata to allow them to be loaded as a pretrained model in future. This is a simple Keras callback
     # that saves the model with this method after each epoch.