Enable python workflows for ImageConverter, MaskedLMMaskGenerator, MultiSegmentPacker, CausalLMPreprocessor, MaskedLMPreprocessor, TextClassifierPreprocessor (#2769)

* Merge python workflow into `MultiSegmentPacker`.

* Add python workflow to `MaskedLMPreprocessor` and `TextClassifierPreprocessor`.

* Enable python workflow for `MaskedLMMaskGenerator`.

* Remove redundant `preprocessing_function` and update the implementations.

* Enable python workflow to `ImageConverter`.

* Enable python workflow to PaliGemma.

* Update.

* Fix gemini's comments.

* Fix `ESMProteinClassifierPreprocessor`.

* Fix list input preprocessing in ImageConverter._call_python

* Address code review comments from gemini-code-assistant

* Move canonicalize_python_inputs and compute_padding_mask to keras_hub.src.utils.tensor_utils. Simplify the call functions in ESM and RoformerV2.

Author: james77777778

keras_hub/api/layers/__init__.py

@@ -4,7 +4,6 @@
 since your modifications would be overwritten.
 """
 
-from keras_hub.layers import v2 as v2
 from keras_hub.src.layers.modeling.alibi_bias import AlibiBias as AlibiBias
 from keras_hub.src.layers.modeling.anchor_generator import (
     AnchorGenerator as AnchorGenerator,

keras_hub/api/layers/v2/__init__.py

@@ -4,7 +4,6 @@
 since your modifications would be overwritten.
 """
 
-from keras_hub.models import v2 as v2
 from keras_hub.src.models.albert.albert_backbone import (
     AlbertBackbone as AlbertBackbone,
 )

keras_hub/api/models/__init__.py

@@ -1,7 +1,6 @@
 import math
 
 import keras
-import ml_dtypes
 import numpy as np
 from keras import ops
 
@@ -20,95 +19,6 @@
 from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 
-# TODO: Use `keras.layers.Resizing` once `antialias` is configurable.
-# https://github.com/keras-team/keras/pull/20972
-def _saturate_cast(x, dtype, backend_module):
-    def get_dtype_min_max(dtype):
-        if "bool" == dtype:
-            dtype_min = 0
-            dtype_max = 1
-        elif "int" in dtype:
-            dtype_min = ml_dtypes.iinfo(dtype).min
-            dtype_max = ml_dtypes.iinfo(dtype).max
-        else:
-            dtype_min = ml_dtypes.finfo(dtype).min
-            dtype_max = ml_dtypes.finfo(dtype).max
-        return dtype_min, dtype_max
-
-    dtype = keras.backend.standardize_dtype(dtype)
-    in_dtype = keras.backend.standardize_dtype(x.dtype)
-    in_min, in_max = get_dtype_min_max(in_dtype)
-    out_min, out_max = get_dtype_min_max(dtype)
-
-    min_limit = np.maximum(in_min, out_min).astype(in_dtype)
-    if min_limit < out_min:
-        min_limit = np.nextafter(min_limit, 0, dtype=in_dtype)
-    max_limit = np.minimum(in_max, out_max).astype(in_dtype)
-    if max_limit > out_max:
-        max_limit = np.nextafter(max_limit, 0, dtype=in_dtype)
-
-    x = backend_module.numpy.clip(x, min_limit, max_limit)
-    return backend_module.cast(x, dtype)
-
-
-class ResizingAntialiasConfigurable(keras.layers.Resizing):
-    """A preprocessing layer which resizes images.
-
-    This class is the same as `keras.layers.Resizing` but exposes `antialias` as
-    a configurable parameter.
-    """
-
-    def __init__(
-        self,
-        height,
-        width,
-        interpolation="bilinear",
-        antialias=False,
-        crop_to_aspect_ratio=False,
-        pad_to_aspect_ratio=False,
-        fill_mode="constant",
-        fill_value=0.0,
-        data_format=None,
-        **kwargs,
-    ):
-        super().__init__(
-            height=height,
-            width=width,
-            interpolation=interpolation,
-            crop_to_aspect_ratio=crop_to_aspect_ratio,
-            pad_to_aspect_ratio=pad_to_aspect_ratio,
-            fill_mode=fill_mode,
-            fill_value=fill_value,
-            data_format=data_format,
-            **kwargs,
-        )
-        self.antialias = bool(antialias)
-
-    def transform_images(self, images, transformation=None, training=True):
-        size = (self.height, self.width)
-        resized = self.backend.image.resize(
-            images,
-            size=size,
-            interpolation=self.interpolation,
-            antialias=self.antialias,  # Added.
-            data_format=self.data_format,
-            crop_to_aspect_ratio=self.crop_to_aspect_ratio,
-            pad_to_aspect_ratio=self.pad_to_aspect_ratio,
-            fill_mode=self.fill_mode,
-            fill_value=self.fill_value,
-        )
-        if resized.dtype == images.dtype:
-            return resized
-        if keras.backend.is_int_dtype(images.dtype):
-            resized = self.backend.numpy.round(resized)
-        return _saturate_cast(resized, images.dtype, self.backend)
-
-    def get_config(self):
-        config = super().get_config()
-        config.update({"antialias": self.antialias})
-        return config
-
-
 @keras_hub_export("keras_hub.layers.ImageConverter")
 class ImageConverter(PreprocessingLayer):
     """Preprocess raw images into model ready inputs.
@@ -213,7 +123,10 @@ def __init__(
             scale = [scale / s for s in std]
             offset = [-m / s for m, s in zip(kwargs.pop("mean"), std)]
 
-        super().__init__(**kwargs)
+        _allow_python_workflow = kwargs.pop("_allow_python_workflow", True)
+        super().__init__(
+            _allow_python_workflow=_allow_python_workflow, **kwargs
+        )
 
         if crop_to_aspect_ratio and pad_to_aspect_ratio:
             raise ValueError(
@@ -226,7 +139,7 @@ def __init__(
         resizing_kwargs = {}
         if check_bounding_box_support():
             resizing_kwargs["bounding_box_format"] = bounding_box_format
-        self.resizing = ResizingAntialiasConfigurable(
+        self.resizing = keras.layers.Resizing(
             height=image_size[0] if image_size else None,
             width=image_size[1] if image_size else None,
             crop_to_aspect_ratio=crop_to_aspect_ratio,
@@ -262,8 +175,33 @@ def image_size(self, value):
         self.resizing.height = value[0]
         self.resizing.width = value[1]
 
-    @preprocessing_function
-    def call(self, inputs):
+    def _call_python(self, inputs):
+        if isinstance(inputs, list):
+            try:
+                inputs = ops.convert_to_tensor(np.array(inputs))
+            except ValueError:
+                return [self._call_python(x) for x in inputs]
+        elif isinstance(inputs, tuple):
+            try:
+                inputs = ops.convert_to_tensor(np.array(inputs))
+            except ValueError:
+                return tuple(self._call_python(x) for x in inputs)
+
+        def canonicalize(x):
+            if isinstance(x, dict):
+                return {k: canonicalize(v) for k, v in x.items()}
+            if isinstance(x, tuple):
+                return tuple(canonicalize(v) for v in x)
+            if isinstance(x, list):
+                try:
+                    return ops.convert_to_tensor(np.array(x))
+                except ValueError:
+                    return [ops.convert_to_tensor(item) for item in x]
+            if isinstance(x, np.ndarray):
+                return ops.convert_to_tensor(x)
+            return x
+
+        inputs = canonicalize(inputs)
         if self.image_size is not None:
             inputs = self.resizing(inputs)
         # Allow dictionary input for handling bounding boxes.
@@ -287,6 +225,16 @@ def call(self, inputs):
             inputs = x
         return inputs
 
+    @preprocessing_function
+    def _call_tf(self, inputs):
+        return self._call_python(inputs)
+
+    def call(self, inputs):
+        if not self._allow_python_workflow or in_tf_function():
+            return self._call_tf(inputs)
+        else:
+            return self._call_python(inputs)
+
     def _expand_non_channel_dims(self, value, inputs):
         """Expand non channel dims so value is broadcastable with inputs."""
         unbatched = len(ops.shape(inputs)) == 3

keras_hub/api/models/v2/__init__.py

@@ -17,20 +17,38 @@
 
 
 class ImageConverterTest(TestCase):
+    def setUp(self):
+        super().setUp()
+        self._allow_python_workflow = True
+
     def test_resize_simple(self):
-        converter = ImageConverter(height=4, width=4, scale=1 / 255.0)
+        converter = ImageConverter(
+            height=4,
+            width=4,
+            scale=1 / 255.0,
+            _allow_python_workflow=self._allow_python_workflow,
+        )
         inputs = np.ones((10, 10, 3)) * 255.0
         outputs = converter(inputs)
         self.assertAllClose(outputs, ops.ones((4, 4, 3)))
 
     def test_resize_dataset(self):
-        converter = ImageConverter(image_size=(4, 4), scale=1 / 255.0)
+        converter = ImageConverter(
+            image_size=(4, 4),
+            scale=1 / 255.0,
+            _allow_python_workflow=self._allow_python_workflow,
+        )
         ds = tf.data.Dataset.from_tensor_slices(tf.zeros((8, 10, 10, 3)))
         batch = ds.batch(2).map(converter).take(1).get_single_element()
         self.assertAllClose(batch, tf.zeros((2, 4, 4, 3)))
 
     def test_resize_in_model(self):
-        converter = ImageConverter(height=4, width=4, scale=1 / 255.0)
+        converter = ImageConverter(
+            height=4,
+            width=4,
+            scale=1 / 255.0,
+            _allow_python_workflow=self._allow_python_workflow,
+        )
         inputs = keras.Input(shape=(10, 10, 3))
         outputs = converter(inputs)
         model = keras.Model(inputs, outputs)
@@ -42,6 +60,7 @@ def test_unbatched(self):
             image_size=(4, 4),
             scale=(1.0 / 255.0, 0.8 / 255.0, 1.2 / 255.0),
             offset=(0.2, -0.1, 0.25),
+            _allow_python_workflow=self._allow_python_workflow,
         )
         inputs = np.ones((10, 10, 3)) * 128
         outputs = converter(inputs)
@@ -51,15 +70,22 @@ def test_unbatched(self):
         self.assertAllClose(outputs[:, :, 2], np.ones((4, 4)) * 0.852353)
 
     def test_dtypes(self):
-        converter = ImageConverter(image_size=(4, 4), scale=1.0 / 255.0)
+        converter = ImageConverter(
+            image_size=(4, 4),
+            scale=1.0 / 255.0,
+            _allow_python_workflow=self._allow_python_workflow,
+        )
         int_image = ops.ones((10, 10, 3), dtype="uint8") * 255
         float_image = ops.ones((10, 10, 3), dtype="float64") * 255
         self.assertDTypeEqual(converter(int_image), "float32")
         self.assertDTypeEqual(converter(float_image), "float32")
         self.assertAllClose(converter(int_image), np.ones((4, 4, 3)))
         self.assertAllClose(converter(float_image), np.ones((4, 4, 3)))
         converter = ImageConverter(
-            image_size=(4, 4), scale=1.0 / 255.0, dtype="bfloat16"
+            image_size=(4, 4),
+            scale=1.0 / 255.0,
+            dtype="bfloat16",
+            _allow_python_workflow=self._allow_python_workflow,
         )
         self.assertDTypeEqual(converter(int_image), "bfloat16")
         self.assertDTypeEqual(converter(float_image), "bfloat16")
@@ -77,6 +103,7 @@ def test_resize_batch(self, crop_to_aspect_ratio, pad_to_aspect_ratio):
             offset=(0.2, -0.1, 0.25),
             crop_to_aspect_ratio=crop_to_aspect_ratio,
             pad_to_aspect_ratio=pad_to_aspect_ratio,
+            _allow_python_workflow=self._allow_python_workflow,
         )
         inputs = np.ones((2, 10, 10, 3)) * 128
         outputs = converter(inputs)
@@ -92,6 +119,7 @@ def test_pad_and_crop_to_aspect_ratio(self):
                 scale=1 / 255.0,
                 crop_to_aspect_ratio=True,
                 pad_to_aspect_ratio=True,
+                _allow_python_workflow=self._allow_python_workflow,
             )
 
     def test_config(self):
@@ -101,6 +129,7 @@ def test_config(self):
             offset=(0.2, -0.1, 0.25),
             crop_to_aspect_ratio=False,
             interpolation="nearest",
+            _allow_python_workflow=self._allow_python_workflow,
         )
         clone = ImageConverter.from_config(converter.get_config())
         test_batch = np.random.rand(4, 10, 20, 3) * 255
@@ -134,6 +163,7 @@ def test_save_to_preset(self):
         converter = ImageConverter.from_preset(
             "resnet_50_imagenet",
             interpolation="nearest",
+            _allow_python_workflow=self._allow_python_workflow,
         )
         converter.save_to_preset(save_dir)
         # Save a tiny backbone so the preset is valid.
@@ -156,3 +186,31 @@ def test_save_to_preset(self):
         restored = ImageConverter.from_preset(save_dir)
         test_image = np.random.rand(100, 100, 3) * 255
         self.assertAllClose(restored(test_image), converter(test_image))
+
+    def test_inhomogeneous_batch(self):
+        if not self._allow_python_workflow:
+            self.skipTest("TF workflow does not support inhomogeneous batches.")
+        converter = ImageConverter(
+            image_size=(4, 4),
+            scale=1 / 255.0,
+            _allow_python_workflow=self._allow_python_workflow,
+        )
+        inputs = [
+            np.ones((10, 10, 3)) * 255.0,
+            np.ones((20, 20, 3)) * 128.0,
+        ]
+        outputs = converter(inputs)
+        self.assertIsInstance(outputs, list)
+        self.assertEqual(len(outputs), 2)
+        self.assertEqual(ops.shape(outputs[0]), (4, 4, 3))
+        self.assertEqual(ops.shape(outputs[1]), (4, 4, 3))
+        self.assertAllClose(outputs[0], ops.ones((4, 4, 3)))
+        self.assertAllClose(outputs[1], ops.ones((4, 4, 3)) * (128.0 / 255.0))
+
+
+class ImageConverterTFTest(ImageConverterTest):
+    """Set `_allow_python_workflow=False` to test TF execution."""
+
+    def setUp(self):
+        super().setUp()
+        self._allow_python_workflow = False