Merge 'upstream/keras-hub' into segmentation

keras_nlp/api/models/__init__.py

@@ -65,7 +65,7 @@
 )
 from keras_nlp.src.models.bloom.bloom_tokenizer import BloomTokenizer
 from keras_nlp.src.models.causal_lm import CausalLM
-from keras_nlp.src.models.classifier import Classifier
+from keras_nlp.src.models.causal_lm_preprocessor import CausalLMPreprocessor
 from keras_nlp.src.models.csp_darknet.csp_darknet_backbone import (
     CSPDarkNetBackbone,
 )
@@ -96,16 +96,16 @@
 from keras_nlp.src.models.deberta_v3.deberta_v3_tokenizer import (
     DebertaV3Tokenizer,
 )
+from keras_nlp.src.models.densenet.densenet_backbone import DenseNetBackbone
+from keras_nlp.src.models.densenet.densenet_image_classifier import (
+    DenseNetImageClassifier,
+)
 from keras_nlp.src.models.deeplab_v3.deeplab_v3_backbone import (
     DeepLabV3Backbone,
 )
 from keras_nlp.src.models.deeplab_v3.deeplab_v3_segmenter import (
     DeepLabV3ImageSegmenter,
 )
-from keras_nlp.src.models.densenet.densenet_backbone import DenseNetBackbone
-from keras_nlp.src.models.densenet.densenet_image_classifier import (
-    DenseNetImageClassifier,
-)
 from keras_nlp.src.models.distil_bert.distil_bert_backbone import (
     DistilBertBackbone,
 )
@@ -134,10 +134,6 @@
     EfficientNetBackbone,
 )
 from keras_nlp.src.models.electra.electra_backbone import ElectraBackbone
-from keras_nlp.src.models.electra.electra_preprocessor import (
-    ElectraPreprocessor,
-)
-from keras_nlp.src.models.electra.electra_backbone import ElectraBackbone
 from keras_nlp.src.models.electra.electra_tokenizer import ElectraTokenizer
 from keras_nlp.src.models.f_net.f_net_backbone import FNetBackbone
 from keras_nlp.src.models.f_net.f_net_masked_lm import FNetMaskedLM
@@ -266,11 +262,10 @@
 from keras_nlp.src.models.t5.t5_backbone import T5Backbone
 from keras_nlp.src.models.t5.t5_tokenizer import T5Tokenizer
 from keras_nlp.src.models.task import Task
-from keras_nlp.src.models.vgg.vgg_backbone import VGGBackbone
-from keras_nlp.src.models.vgg.vgg_image_classifier import VGGImageClassifier
-from keras_nlp.src.models.vit_det.vit_det_backbone import ViTDetBackbone
-from keras_nlp.src.models.whisper.whisper_audio_feature_extractor import (
-    WhisperAudioFeatureExtractor,
+from keras_nlp.src.models.text_classifier import TextClassifier
+from keras_nlp.src.models.text_classifier import TextClassifier as Classifier
+from keras_nlp.src.models.text_classifier_preprocessor import (
+    TextClassifierPreprocessor,
 )
 from keras_nlp.src.models.vgg.vgg_backbone import VGGBackbone
 from keras_nlp.src.models.vgg.vgg_image_classifier import VGGImageClassifier

keras_nlp/src/bounding_box/utils_test.py

@@ -32,15 +32,19 @@ def test_clip_to_image_standard(self):
             bounding_boxes, bounding_box_format="xyxy", images=image
         )
         boxes = bounding_boxes["boxes"]
-        self.assertAllGreaterEqual(boxes, 0)
+        self.assertAllGreaterEqual(ops.convert_to_numpy(boxes), 0)
         (
             x1,
             y1,
             x2,
             y2,
         ) = ops.split(boxes, 4, axis=1)
-        self.assertAllLessEqual(ops.concatenate([x1, x2], axis=1), width)
-        self.assertAllLessEqual(ops.concatenate([y1, y2], axis=1), height)
+        self.assertAllLessEqual(
+            ops.convert_to_numpy(ops.concatenate([x1, x2], axis=1)), width
+        )
+        self.assertAllLessEqual(
+            ops.convert_to_numpy(ops.concatenate([y1, y2], axis=1)), height
+        )
         # Test relative format batched
         image = ops.ones(shape=(1, height, width, 3))
 
@@ -51,7 +55,8 @@ def test_clip_to_image_standard(self):
         bounding_boxes = utils.clip_to_image(
             bounding_boxes, bounding_box_format="rel_xyxy", images=image
         )
-        self.assertAllLessEqual(bounding_boxes["boxes"], 1)
+        boxes = bounding_boxes["boxes"]
+        self.assertAllLessEqual(ops.convert_to_numpy(boxes), 1)
 
     def test_clip_to_image_filters_fully_out_bounding_boxes(self):
         # Test xyxy format unbatched

keras_nlp/src/models/backbone.py

@@ -25,8 +25,6 @@
 from keras_nlp.src.utils.preset_utils import save_metadata
 from keras_nlp.src.utils.preset_utils import save_serialized_object
 from keras_nlp.src.utils.python_utils import classproperty
-from keras_nlp.src.utils.timm.convert import load_timm_backbone
-from keras_nlp.src.utils.transformers.convert import load_transformers_backbone
 
 
 @keras_nlp_export("keras_nlp.models.Backbone")
@@ -193,15 +191,9 @@ class like `keras_nlp.models.Backbone.from_preset()`, or from
         )
         ```
         """
-        format = check_format(preset)
-
-        if format == "transformers":
-            return load_transformers_backbone(cls, preset, load_weights)
-        elif format == "timm":
-            return load_timm_backbone(cls, preset, load_weights, **kwargs)
-
-        preset_cls = check_config_class(preset)
-        if not issubclass(preset_cls, cls):
+        loader = get_preset_loader(preset)
+        backbone_cls = loader.check_backbone_class()
+        if not issubclass(backbone_cls, cls):
             raise ValueError(
                 f"Saved preset has type `{backbone_cls.__name__}` which is not "
                 f"a subclass of calling class `{cls.__name__}`. Call "

keras_nlp/src/models/stable_diffusion_v3/clip_preprocessor.py

@@ -18,9 +18,7 @@
 from keras_nlp.src.models.stable_diffusion_v3.clip_tokenizer import (
     CLIPTokenizer,
 )
-from keras_nlp.src.utils.keras_utils import (
-    convert_inputs_to_list_of_tensor_segments,
-)
+from keras_nlp.src.utils.tensor_utils import preprocessing_function
 
 try:
     import tensorflow as tf
@@ -65,22 +63,13 @@ def build(self, input_shape):
         )
         self.built = True
 
-    # TODO: Use `@tf_preprocessing_function` after rebasing.
+    @preprocessing_function
     def call(self, x, y=None, sample_weight=None, sequence_length=None):
-        x = convert_inputs_to_list_of_tensor_segments(x)
-        if len(x) != 1:
-            raise ValueError(
-                "T5XXL requires each input feature to contain only "
-                f"one segment, but received {len(x)}. If you are using T5XXL"
-                " for a multi-segment classification task, please refer to "
-                "classification models like BERT or RoBERTa."
-            )
         if self.to_lower:
             x = tf.strings.lower(x)
-        sequence_length = sequence_length or self.sequence_length
         token_ids, padding_mask = self.packer(
-            self.tokenizer(x[0]),
-            sequence_length=sequence_length,
+            self.tokenizer(x),
+            sequence_length=sequence_length or self.sequence_length,
             add_start_value=self.add_start_token,
             add_end_value=self.add_end_token,
         )

keras_nlp/src/models/stable_diffusion_v3/t5_xxl_preprocessor.py

@@ -16,9 +16,7 @@
 from keras_nlp.src.layers.preprocessing.start_end_packer import StartEndPacker
 from keras_nlp.src.models.preprocessor import Preprocessor
 from keras_nlp.src.models.t5.t5_tokenizer import T5Tokenizer
-from keras_nlp.src.utils.keras_utils import (
-    convert_inputs_to_list_of_tensor_segments,
-)
+from keras_nlp.src.utils.tensor_utils import preprocessing_function
 
 
 class T5XXLPreprocessor(Preprocessor):
@@ -50,19 +48,11 @@ def build(self, input_shape):
         )
         self.built = True
 
+    @preprocessing_function
     def call(self, x, y=None, sample_weight=None, sequence_length=None):
-        x = convert_inputs_to_list_of_tensor_segments(x)
-        if len(x) != 1:
-            raise ValueError(
-                "T5XXL requires each input feature to contain only "
-                f"one segment, but received {len(x)}. If you are using T5XXL"
-                " for a multi-segment classification task, please refer to "
-                "classification models like BERT or RoBERTa."
-            )
-        sequence_length = sequence_length or self.sequence_length
         token_ids, padding_mask = self.packer(
-            self.tokenizer(x[0]),
-            sequence_length=sequence_length,
+            self.tokenizer(x),
+            sequence_length=sequence_length or self.sequence_length,
             add_start_value=self.add_start_token,
             add_end_value=self.add_end_token,
         )

keras_nlp/src/utils/preset_utils.py

@@ -571,35 +571,7 @@ def load_json(preset, config_file=CONFIG_FILE):
     return config
 
 
-def check_format(preset):
-    if check_file_exists(preset, SAFETENSOR_FILE) or check_file_exists(
-        preset, SAFETENSOR_CONFIG_FILE
-    ):
-        # Determine the format by parsing the config file.
-        config = load_config(preset, HF_CONFIG_FILE)
-        if "hf://timm" in preset or "architecture" in config:
-            return "timm"
-        return "transformers"
-
-    if not check_file_exists(preset, METADATA_FILE):
-        raise FileNotFoundError(
-            f"The preset directory `{preset}` doesn't have a file named `{METADATA_FILE}`, "
-            "or you do not have access to it. This file is required to load a Keras model "
-            "preset. Please verify that the model you are trying to load is a Keras model."
-        )
-    metadata = load_config(preset, METADATA_FILE)
-    if "keras_version" not in metadata:
-        raise ValueError(
-            f"`{METADATA_FILE}` in the preset directory `{preset}` doesn't have `keras_version`. "
-            "Please verify that the model you are trying to load is a Keras model."
-        )
-    return "keras"
-
-
-def load_serialized_object(preset, config_file=CONFIG_FILE, **kwargs):
-    kwargs = kwargs or {}
-    config = load_config(preset, config_file)
-
+def load_serialized_object(config, **kwargs):
     # `dtype` in config might be a serialized `DTypePolicy` or `DTypePolicyMap`.
     # Ensure that `dtype` is properly configured.
     dtype = kwargs.pop("dtype", None)

keras_nlp/src/utils/timm/convert_resnet.py

@@ -13,10 +13,9 @@
 # limitations under the License.
 import numpy as np
 
-from keras_nlp.src.utils.preset_utils import HF_CONFIG_FILE
-from keras_nlp.src.utils.preset_utils import jax_memory_cleanup
-from keras_nlp.src.utils.preset_utils import load_config
-from keras_nlp.src.utils.transformers.safetensor_utils import SafetensorLoader
+from keras_nlp.src.models.resnet.resnet_backbone import ResNetBackbone
+
+backbone_cls = ResNetBackbone
 
 
 def convert_backbone_config(timm_config):
@@ -56,6 +55,8 @@ def convert_backbone_config(timm_config):
         stackwise_num_strides=[1, 2, 2, 2],
         block_type=block_type,
         use_pre_activation=use_pre_activation,
+        input_conv_filters=[64],
+        input_conv_kernel_sizes=[7],
     )
 
 
@@ -100,10 +101,10 @@ def port_batch_normalization(keras_layer_name, hf_weight_prefix):
     for stack_index in range(num_stacks):
         for block_idx in range(backbone.stackwise_num_blocks[stack_index]):
             if version == "v1":
-                keras_name = f"v1_stack{stack_index}_block{block_idx}"
+                keras_name = f"stack{stack_index}_block{block_idx}"
                 hf_name = f"layer{stack_index+1}.{block_idx}"
             else:
-                keras_name = f"v2_stack{stack_index}_block{block_idx}"
+                keras_name = f"stack{stack_index}_block{block_idx}"
                 hf_name = f"stages.{stack_index}.blocks.{block_idx}"
 
             if version == "v1":
@@ -157,15 +158,3 @@ def port_batch_normalization(keras_layer_name, hf_weight_prefix):
     normalization_layer.input_mean = mean
     normalization_layer.input_variance = [s**2 for s in std]
     normalization_layer.build(normalization_layer._build_input_shape)
-
-
-def load_resnet_backbone(cls, preset, load_weights, **kwargs):
-    timm_config = load_config(preset, HF_CONFIG_FILE)
-    keras_config = convert_backbone_config(timm_config)
-    backbone = cls(**keras_config, **kwargs)
-    if load_weights:
-        jax_memory_cleanup(backbone)
-        # Use prefix="" to avoid using `get_prefixed_key`.
-        with SafetensorLoader(preset, prefix="") as loader:
-            convert_weights(backbone, loader, timm_config)
-    return backbone