keras-hub-nightly 0.20.0.dev202504030357__py3-none-any.whl → 0.21.0.dev202504050402__py3-none-any.whl

keras_hub/src/models/gemma3/{gemma3_vit.py → gemma3_vision_encoder.py}

@@ -1,11 +1,150 @@
 import keras
 from keras import ops
 
+from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.gemma.rms_normalization import RMSNormalization
 from keras_hub.src.utils.keras_utils import clone_initializer
 
 
-class Gemma3VitEmbeddings(keras.layers.Layer):
+@keras_hub_export("keras_hub.models.Gemma3VisionEncoder")
+class Gemma3VisionEncoder(keras.Model):
+    """Vision Transformer (ViT) model for Gemma3.
+
+    Args:
+        image_size: int. The height/width of the image. Both height and width is
+            expected to be the same.
+        patch_size: int. The size of each square patch in the input image.
+        num_heads: int. The number of attention heads for the vision(image)
+            transformer encoder.
+        hidden_dim: int. The size of the transformer hidden state at the end
+            of each vision transformer layer.
+        num_layers: int. The number of transformer layers.
+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for transformer.
+        output_dim: int. The odimension of the output returned by the model.
+        pool_size: int. Factors by which to downscale `(dim1, dim2)` in the
+            average pooling layer. The same value is used for `"strides"`.
+            Defaults to 14.
+        layer_norm_epsilon: float. The epsilon value user for every layer norm
+            in all transformer blocks. Defaults to `1e-6`.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for the models computations and weights. Note that some
+            computations, such as softmax and layer normalization will always
+            be done a float32 precision regardless of dtype.
+
+    Example:
+    ```python
+    image = np.random.rand(224, 224, 3)
+    vit_model = Gemma3VisionEncoder(image_size=224)
+    # The output will be of shape:
+    # [batch_size, num_vision_tokens_per_image, hidden_dim]
+    output = vit_model([image])
+    ```
+    """
+
+    def __init__(
+        self,
+        image_size,
+        patch_size,
+        num_heads,
+        hidden_dim,
+        num_layers,
+        intermediate_dim,
+        output_dim,
+        pool_size=14,
+        layer_norm_epsilon=1e-6,
+        dtype=None,
+        **kwargs,
+    ):
+        # If the passed dtype is `bfloat16`, use `float32` to maintain parity
+        # with other framework implementations.
+        if dtype == "bfloat16":
+            dtype = "float32"
+
+        # === Functional Model ===
+        image_input = keras.Input(
+            shape=(None, image_size, image_size, 3),
+            name="images",
+        )
+        x = image_input  # Intermediate result.
+        x = Gemma3VisionEncoderBlock(
+            hidden_dim=hidden_dim,
+            num_layers=num_layers,
+            num_heads=num_heads,
+            intermediate_dim=intermediate_dim,
+            patch_size=patch_size,
+            image_size=image_size,
+            dtype=dtype,
+            name="image_encoder",
+        )(x)
+
+        x = Gemma3VisionAveragePooling(
+            image_size=image_size,
+            patch_size=patch_size,
+            pool_size=pool_size,
+            dtype=dtype,
+            name="pooling",
+        )(x)
+
+        x = Gemma3VisionOutput(
+            output_dim=output_dim,
+            layer_norm_epsilon=layer_norm_epsilon,
+            kernel_initializer=keras.initializers.RandomNormal(
+                mean=0.0, stddev=0.01
+            ),
+            dtype=dtype,
+            name="vision_output_encoder",
+        )(x)
+
+        outputs = x
+        super().__init__(
+            inputs=image_input,
+            outputs=outputs,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        self.intermediate_dim = intermediate_dim
+        self.output_dim = output_dim
+        self.pool_size = pool_size
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.num_vision_tokens_per_image = (
+            (image_size // patch_size) ** 2
+        ) // (pool_size**2)
+
+        # Before Keras 3.2, there is no `keras.dtype_policies.get`.
+        if hasattr(keras.dtype_policies, "get"):
+            self.dtype_policy = keras.dtype_policies.get(dtype)
+        else:
+            if isinstance(dtype, keras.dtype_policies.DTypePolicy):
+                dtype = dtype.name
+            dtype = dtype or keras.config.dtype_policy().name
+            self.dtype_policy = keras.dtype_policies.DTypePolicy(dtype)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_heads": self.num_heads,
+                "hidden_dim": self.hidden_dim,
+                "num_layers": self.num_layers,
+                "intermediate_dim": self.intermediate_dim,
+                "output_dim": self.output_dim,
+                "pool_size": self.pool_size,
+                "image_size": self.image_size,
+                "patch_size": self.patch_size,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+            }
+        )
+        return config
+
+
+class Gemma3VisionEmbedding(keras.layers.Layer):
     def __init__(
         self,
         image_size,
@@ -62,7 +201,7 @@ class Gemma3VitEmbeddings(keras.layers.Layer):
         )
 
 
-class Gemma3VitAttention(keras.layers.Layer):
+class Gemma3VisionAttention(keras.layers.Layer):
     """
     Adapted from https://github.com/huggingface/transformers/blob/main/src/transformers/models/clip/modeling_clip.py
     """
@@ -197,7 +336,7 @@ class Gemma3VitAttention(keras.layers.Layer):
         return config
 
 
-class Gemma3VitEncoderBlock(keras.layers.Layer):
+class Gemma3VisionEncoderLayer(keras.layers.Layer):
     def __init__(
         self,
         num_heads,
@@ -217,7 +356,7 @@ class Gemma3VitEncoderBlock(keras.layers.Layer):
 
     def build(self, input_shape):
         hidden_dim = input_shape[-1]
-        self.attn = Gemma3VitAttention(
+        self.attn = Gemma3VisionAttention(
             hidden_dim,
             self.num_heads,
             dtype=self.dtype_policy,
@@ -277,7 +416,7 @@ class Gemma3VitEncoderBlock(keras.layers.Layer):
         return config
 
 
-class Gemma3VitEncoder(keras.layers.Layer):
+class Gemma3VisionEncoderBlock(keras.layers.Layer):
     def __init__(
         self,
         patch_size,
@@ -303,7 +442,7 @@ class Gemma3VitEncoder(keras.layers.Layer):
             dtype=dtype,
             name="encoder_layer_norm",
         )
-        self.vision_embeddings = Gemma3VitEmbeddings(
+        self.vision_embeddings = Gemma3VisionEmbedding(
             hidden_dim=hidden_dim,
             patch_size=patch_size,
             image_size=image_size,
@@ -311,7 +450,7 @@ class Gemma3VitEncoder(keras.layers.Layer):
             name="encoder_embeddings",
         )
         self.resblocks = [
-            Gemma3VitEncoderBlock(
+            Gemma3VisionEncoderLayer(
                 self.num_heads,
                 self.intermediate_dim,
                 dtype=dtype,
@@ -321,7 +460,7 @@ class Gemma3VitEncoder(keras.layers.Layer):
         ]
 
     def build(self, inputs_shape):
-        # Collapse `batch_size`, dummy axis, `image_max_length` into one.
+        # Collapse `batch_size`, dummy axis, `max_images_per_prompt` into one.
         inputs_shape = [None] + list(inputs_shape[2:])
         self.vision_embeddings.build(inputs_shape)
         for block in self.resblocks:
@@ -332,7 +471,7 @@ class Gemma3VitEncoder(keras.layers.Layer):
     def call(self, inputs, mask=None):
         inputs_shape = ops.shape(inputs)
 
-        # Collapse `batch_size`, dummy axis, `image_max_length` into one.
+        # Collapse `batch_size`, dummy axis, `max_images_per_prompt` into one.
         inputs = ops.reshape(
             inputs,
             [inputs_shape[0] * inputs_shape[1]] + list(inputs_shape[2:]),
@@ -372,7 +511,7 @@ class Gemma3VitEncoder(keras.layers.Layer):
         return config
 
 
-class AveragePooling(keras.layers.Layer):
+class Gemma3VisionAveragePooling(keras.layers.Layer):
     def __init__(self, image_size, patch_size, pool_size, **kwargs):
         super().__init__(**kwargs)
 
@@ -425,7 +564,7 @@ class AveragePooling(keras.layers.Layer):
         return config
 
 
-class Gemma3VisionOutputEncoder(keras.layers.Layer):
+class Gemma3VisionOutput(keras.layers.Layer):
     def __init__(
         self,
         output_dim,
@@ -478,131 +617,3 @@ class Gemma3VisionOutputEncoder(keras.layers.Layer):
 
     def compute_output_shape(self, input_shape):
         return input_shape[:-1] + (self.output_dim,)
-
-
-class Gemma3Vit(keras.Model):
-    """Vision Transformer (ViT) model for Gemma3.
-
-    Args:
-        image_size: int. The height/width of the image. Both height and width is
-            expected to be the same.
-        patch_size: int. The size of each square patch in the input image.
-        num_heads: int. The number of attention heads for the vision(image)
-            transformer encoder.
-        hidden_dim: int. The size of the transformer hidden state at the end
-            of each vision transformer layer.
-        num_layers: int. The number of transformer layers.
-        intermediate_dim: int. The output dimension of the first Dense layer in
-            a two-layer feedforward network for transformer.
-        pool_size: int. Factors by which to downscale `(dim1, dim2)` in the
-            average pooling layer. The same value is used for `"strides"`.
-        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
-            for the models computations and weights. Note that some
-            computations, such as softmax and layer normalization will always
-            be done a float32 precision regardless of dtype.
-
-    Example:
-    ```python
-    image = np.random.rand(224, 224, 3)
-    vit_model = Gemma3Vit(image_size=224)
-    # The output will be of shape:
-    # [batch_size, num_vision_tokens_per_image, hidden_dim]
-    output = vit_model([image])
-    ```
-    """
-
-    def __init__(
-        self,
-        image_size,
-        patch_size,
-        num_heads,
-        hidden_dim,
-        num_layers,
-        intermediate_dim,
-        output_dim,
-        pool_size=14,
-        layer_norm_epsilon=1e-6,
-        dtype=None,
-        **kwargs,
-    ):
-        # === Functional Model ===
-        image_input = keras.Input(
-            shape=(None, image_size, image_size, 3),
-            name="images",
-        )
-        x = image_input  # Intermediate result.
-        x = Gemma3VitEncoder(
-            hidden_dim=hidden_dim,
-            num_layers=num_layers,
-            num_heads=num_heads,
-            intermediate_dim=intermediate_dim,
-            patch_size=patch_size,
-            image_size=image_size,
-            dtype=dtype,
-            name="image_encoder",
-        )(x)
-
-        x = AveragePooling(
-            image_size=image_size,
-            patch_size=patch_size,
-            pool_size=pool_size,
-            dtype=dtype,
-            name="pooling",
-        )(x)
-
-        x = Gemma3VisionOutputEncoder(
-            output_dim=output_dim,
-            layer_norm_epsilon=layer_norm_epsilon,
-            kernel_initializer=keras.initializers.RandomNormal(
-                mean=0.0, stddev=0.01
-            ),
-            dtype=dtype,
-            name="vision_output_encoder",
-        )(x)
-
-        outputs = x
-        super().__init__(
-            inputs=image_input,
-            outputs=outputs,
-            **kwargs,
-        )
-
-        # === Config ===
-        self.image_size = image_size
-        self.patch_size = patch_size
-        self.num_heads = num_heads
-        self.hidden_dim = hidden_dim
-        self.num_layers = num_layers
-        self.intermediate_dim = intermediate_dim
-        self.output_dim = output_dim
-        self.pool_size = pool_size
-        self.layer_norm_epsilon = layer_norm_epsilon
-        self.num_vision_tokens_per_image = (
-            (image_size // patch_size) ** 2
-        ) // (pool_size**2)
-
-        # Before Keras 3.2, there is no `keras.dtype_policies.get`.
-        if hasattr(keras.dtype_policies, "get"):
-            self.dtype_policy = keras.dtype_policies.get(dtype)
-        else:
-            if isinstance(dtype, keras.dtype_policies.DTypePolicy):
-                dtype = dtype.name
-            dtype = dtype or keras.config.dtype_policy().name
-            self.dtype_policy = keras.dtype_policies.DTypePolicy(dtype)
-
-    def get_config(self):
-        config = super().get_config()
-        config.update(
-            {
-                "num_heads": self.num_heads,
-                "hidden_dim": self.hidden_dim,
-                "num_layers": self.num_layers,
-                "intermediate_dim": self.intermediate_dim,
-                "output_dim": self.output_dim,
-                "pool_size": self.pool_size,
-                "image_size": self.image_size,
-                "patch_size": self.patch_size,
-                "layer_norm_epsilon": self.layer_norm_epsilon,
-            }
-        )
-        return config

keras_hub/src/models/qwen/qwen_backbone.py

@@ -1,7 +1,6 @@
 import keras
 from keras import ops
 
-from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.layers.modeling.reversible_embedding import (
     ReversibleEmbedding,
 )
@@ -14,12 +13,6 @@ def _qwen_kernel_initializer(stddev=0.02):
     return keras.initializers.RandomNormal(stddev=stddev)
 
 
-@keras_hub_export(
-    [
-        "keras_hub.models.QwenBackbone",
-        "keras_hub.models.Qwen2Backbone",
-    ]
-)
 class QwenBackbone(Backbone):
     """
     The Qwen Transformer core architecture with hyperparameters.

keras_hub/src/models/qwen/qwen_causal_lm.py

@@ -1,7 +1,6 @@
 import keras
 from keras import ops
 
-from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.causal_lm import CausalLM
 from keras_hub.src.models.qwen.qwen_backbone import QwenBackbone
 from keras_hub.src.models.qwen.qwen_causal_lm_preprocessor import (
@@ -10,12 +9,6 @@ from keras_hub.src.models.qwen.qwen_causal_lm_preprocessor import (
 from keras_hub.src.utils.tensor_utils import any_equal
 
 
-@keras_hub_export(
-    [
-        "keras_hub.models.QwenCausalLM",
-        "keras_hub.models.Qwen2CausalLM",
-    ]
-)
 class QwenCausalLM(CausalLM):
     backbone_cls = QwenBackbone
     preprocessor_cls = QwenCausalLMPreprocessor

keras_hub/src/models/qwen/qwen_causal_lm_preprocessor.py

@@ -1,15 +1,8 @@
-from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.causal_lm_preprocessor import CausalLMPreprocessor
 from keras_hub.src.models.qwen.qwen_backbone import QwenBackbone
 from keras_hub.src.models.qwen.qwen_tokenizer import QwenTokenizer
 
 
-@keras_hub_export(
-    [
-        "keras_hub.models.QwenCausalLMPreprocessor",
-        "keras_hub.models.Qwen2CausalLMPreprocessor",
-    ]
-)
 class QwenCausalLMPreprocessor(CausalLMPreprocessor):
     backbone_cls = QwenBackbone
     tokenizer_cls = QwenTokenizer

keras_hub/src/models/qwen/qwen_tokenizer.py

@@ -1,16 +1,7 @@
-from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.qwen.qwen_backbone import QwenBackbone
 from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
 
 
-@keras_hub_export(
-    [
-        "keras_hub.tokenizers.QwenTokenizer",
-        "keras_hub.tokenizers.Qwen2Tokenizer",
-        "keras_hub.models.QwenTokenizer",
-        "keras_hub.models.Qwen2Tokenizer",
-    ]
-)
 class QwenTokenizer(BytePairTokenizer):
     """Tokenizer for Qwen models.
 

keras_hub/src/models/roformer_v2/roformer_v2_backbone.py

@@ -16,7 +16,7 @@ def roformer_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
-@keras_hub_export("keras_hub.models.RorformerV2Backbone")
+@keras_hub_export("keras_hub.models.RoformerV2Backbone")
 class RoformerV2Backbone(Backbone):
     """A RoformerV2 encoder network.
 

keras_hub/src/models/roformer_v2/roformer_v2_text_classifier.py

@@ -10,8 +10,8 @@ from keras_hub.src.models.roformer_v2.roformer_v2_text_classifier_preprocessor i
 )
 
 
-@keras_hub_export("keras_hub.models.RorformerV2TextClassifier")
-class RorformerV2TextClassifier(RobertaTextClassifier):
+@keras_hub_export("keras_hub.models.RoformerV2TextClassifier")
+class RoformerV2TextClassifier(RobertaTextClassifier):
     """An end-to-end RoformerV2 model for classification tasks.
 
     This model attaches a classification head to

keras_hub/src/models/vit/vit_image_converter.py

@@ -53,12 +53,17 @@ class ViTImageConverter(ImageConverter):
 
     @preprocessing_function
     def call(self, inputs):
+        # TODO: Remove this whole function. Why can just use scale and offset
+        # in the base class.
         x = super().call(inputs)
-        # By default normalize using imagenet mean and std
         if self.norm_mean:
-            x = x - self._expand_non_channel_dims(self.norm_mean, x)
+            norm_mean = self._expand_non_channel_dims(self.norm_mean, x)
+            x, norm_mean = self._convert_types(x, norm_mean, self.compute_dtype)
+            x = x - norm_mean
         if self.norm_std:
-            x = x / self._expand_non_channel_dims(self.norm_std, x)
+            norm_std = self._expand_non_channel_dims(self.norm_std, x)
+            x, norm_std = self._convert_types(x, norm_std, x.dtype)
+            x = x / norm_std
 
         return x
 

keras_hub/src/tests/test_case.py

@@ -197,6 +197,7 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         input_data,
         expected_output=None,
         expected_detokenize_output=None,
+        return_output=False,
     ):
         """Run basic tests for a preprocessing layer."""
         layer = cls(**init_kwargs)
@@ -230,6 +231,9 @@ class TestCase(tf.test.TestCase, parameterized.TestCase):
         if expected_output:
             self.assertAllClose(output, expected_output)
 
+        if return_output:
+            return output
+
     def run_preprocessor_test(
         self,
         cls,

keras_hub/src/utils/tensor_utils.py

@@ -28,6 +28,12 @@ def no_convert_scope():
         NO_CONVERT_COUNTER.count = getattr(NO_CONVERT_COUNTER, "count", 0) - 1
 
 
+def in_tf_function():
+    if tf is None:
+        return False
+    return not tf.executing_eagerly()
+
+
 def in_no_convert_scope():
     return getattr(NO_CONVERT_COUNTER, "count", 0) > 0
 

keras_hub/src/version_utils.py

@@ -1,7 +1,7 @@
 from keras_hub.src.api_export import keras_hub_export
 
 # Unique source of truth for the version number.
-__version__ = "0.20.0.dev202504030357"
+__version__ = "0.21.0.dev202504050402"
 
 
 @keras_hub_export("keras_hub.version")

{keras_hub_nightly-0.20.0.dev202504030357.dist-info → keras_hub_nightly-0.21.0.dev202504050402.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-hub-nightly
-Version: 0.20.0.dev202504030357
+Version: 0.21.0.dev202504050402
 Summary: Industry-strength Natural Language Processing extensions for Keras.
 Home-page: https://github.com/keras-team/keras-hub
 Author: Keras team