keras-hub-nightly 0.20.0.dev202503260356__py3-none-any.whl → 0.20.0.dev202503270400__py3-none-any.whl

keras_hub/src/models/gemma3/gemma3_decoder_block.py

@@ -0,0 +1,305 @@
+import keras
+from keras import ops
+
+from keras_hub.src.layers.modeling.transformer_layer_utils import (
+    compute_causal_mask,
+)
+from keras_hub.src.layers.modeling.transformer_layer_utils import (
+    merge_padding_and_attention_mask,
+)
+from keras_hub.src.models.gemma3.gemma3_attention import CachedGemma3Attention
+from keras_hub.src.models.gemma3.rms_normalization import RMSNormalization
+
+
+class Gemma3DecoderBlock(keras.layers.Layer):
+    """Transformer decoder layer for Gemma3.
+
+    This is different from Gemma and Gemma2 in several ways:
+
+    - `use_query_key_norm`: Applies RMS Norm on query, key.
+    - `rope_wavelength`: RoPE wavelength differs from local to global attention
+      layers.
+    - `rope_scaling_factor`: RoPE scaling factor differs from local to global
+      attention layers.
+    - `gate_dim_reduction`: In the gating layers, Gemma and Gemma2 reduce
+      intermediate dimension by 2. For Gemma3, no such reduction happens.
+    - Uses bidirectional attention for images, and causal for everything else.
+    """
+
+    def __init__(
+        self,
+        hidden_dim,
+        intermediate_dim,
+        head_dim,
+        num_query_heads,
+        num_key_value_heads,
+        query_head_dim_normalize=True,
+        use_query_key_norm=False,
+        use_post_ffw_norm=False,
+        use_post_attention_norm=False,
+        gate_dim_reduction=2,
+        logit_soft_cap=None,
+        use_sliding_window_attention=False,
+        sliding_window_size=4096,
+        layer_norm_epsilon=1e-6,
+        rope_wavelength=10_000.0,
+        rope_scaling_factor=1.0,
+        dropout=0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+        self.head_dim = head_dim
+        self.num_query_heads = num_query_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.query_head_dim_normalize = query_head_dim_normalize
+        self.use_query_key_norm = use_query_key_norm
+        self.use_post_ffw_norm = use_post_ffw_norm
+        self.use_post_attention_norm = use_post_attention_norm
+        self.gate_dim_reduction = gate_dim_reduction
+        self.logit_soft_cap = logit_soft_cap
+        self.use_sliding_window_attention = use_sliding_window_attention
+        self.sliding_window_size = sliding_window_size
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.rope_wavelength = rope_wavelength
+        self.rope_scaling_factor = rope_scaling_factor
+        self.dropout = dropout
+
+        self.pre_attention_norm = RMSNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="pre_attention_norm",
+        )
+
+        if use_post_attention_norm:
+            self.post_attention_norm = RMSNormalization(
+                epsilon=self.layer_norm_epsilon,
+                dtype=self.dtype_policy,
+                name="post_attention_norm",
+            )
+
+        self.attention = CachedGemma3Attention(
+            head_dim=head_dim,
+            num_query_heads=num_query_heads,
+            num_key_value_heads=num_key_value_heads,
+            use_query_key_norm=use_query_key_norm,
+            logit_soft_cap=logit_soft_cap,
+            use_sliding_window_attention=use_sliding_window_attention,
+            sliding_window_size=sliding_window_size,
+            query_head_dim_normalize=True,
+            rope_wavelength=rope_wavelength,
+            rope_scaling_factor=rope_scaling_factor,
+            dropout=dropout,
+            dtype=self.dtype_policy,
+            name="attention",
+        )
+
+        if self.dropout > 0:
+            self.attention_dropout = keras.layers.Dropout(rate=dropout)
+            self.feedforward_dropout = keras.layers.Dropout(rate=dropout)
+
+        self.pre_ffw_norm = RMSNormalization(
+            epsilon=self.layer_norm_epsilon,
+            dtype=self.dtype_policy,
+            name="pre_ffw_norm",
+        )
+
+        if use_post_ffw_norm:
+            self.post_ffw_norm = RMSNormalization(
+                epsilon=self.layer_norm_epsilon,
+                dtype=self.dtype_policy,
+                name="post_ffw_norm",
+            )
+
+        self.gating_ffw = keras.layers.EinsumDense(
+            equation="btd,df->btf",
+            output_shape=(None, self.intermediate_dim // gate_dim_reduction),
+            dtype=self.dtype_policy,
+            name="ffw_gating",
+        )
+
+        self.gating_ffw_2 = keras.layers.EinsumDense(
+            equation="btd,df->btf",
+            output_shape=(None, self.intermediate_dim // gate_dim_reduction),
+            dtype=self.dtype_policy,
+            name="ffw_gating_2",
+        )
+
+        self.ffw_linear = keras.layers.EinsumDense(
+            equation="btf,fd->btd",
+            output_shape=(None, self.hidden_dim),
+            dtype=self.dtype_policy,
+            name="ffw_linear",
+        )
+
+    def build(self, input_shape):
+        self.pre_attention_norm.build(input_shape)
+        self.attention.build(input_shape)
+
+        if self.use_post_attention_norm:
+            shape = self.attention.compute_output_shape(input_shape)
+            self.post_attention_norm.build(shape)
+
+        shape = input_shape
+        self.pre_ffw_norm.build(shape)
+        self.gating_ffw.build(shape)
+        self.gating_ffw_2.build(shape)
+
+        shape = self.gating_ffw.compute_output_shape(shape)
+        self.ffw_linear.build(shape)
+
+        if self.use_post_ffw_norm:
+            shape = self.ffw_linear.compute_output_shape(shape)
+            self.post_ffw_norm.build(shape)
+
+        self.built = True
+
+    def compute_output_shape(self, input_shape):
+        # Isometric
+        return input_shape
+
+    def _compute_image_bidirectional_attention_mask(self, text_mask):
+        # text_mask is True for text, False for images. Shape of (bsz, seq_len).
+        bidirectional_mask = ops.logical_not(text_mask)
+
+        # Left pad with 0.
+        padded_mask = ops.cast(
+            ops.pad(bidirectional_mask, [(0, 0), (1, 0)], constant_values=0),
+            dtype="int32",
+        )
+
+        # Assign unique indices to every contiguous span of True.
+        boundary = ops.cast(
+            ops.greater(padded_mask[..., 1:], padded_mask[..., :-1]),
+            dtype="int32",
+        )
+        numbered_boundary = ops.cumsum(boundary, -1)
+        indices = ops.multiply(bidirectional_mask, numbered_boundary)
+
+        indices_expanded_1 = ops.expand_dims(indices, 1)
+        indices_expanded_2 = ops.expand_dims(indices, -1)
+
+        mask = ops.logical_and(
+            ops.equal(
+                indices_expanded_1,
+                indices_expanded_2,
+            ),
+            indices_expanded_2,
+        )
+        return mask
+
+    def _compute_attention_mask(
+        self,
+        x,
+        padding_mask,
+        text_mask,
+        cache,
+        cache_update_index,
+    ):
+        decoder_mask = merge_padding_and_attention_mask(
+            inputs=x, padding_mask=padding_mask, attention_mask=None
+        )
+
+        batch_size = ops.shape(x)[0]
+        input_length = output_length = ops.shape(x)[1]
+        if cache is not None:
+            input_length = ops.shape(cache)[2]
+
+        causal_mask = compute_causal_mask(
+            batch_size=batch_size,
+            input_length=input_length,
+            output_length=output_length,
+            cache_index=cache_update_index,
+        )
+
+        # Compute bidirectional mask (image tokens can attend to each other
+        # in both directions, within the same image).
+        if text_mask is not None:
+            bidirectional_image_mask = (
+                self._compute_image_bidirectional_attention_mask(text_mask)
+            )
+            causal_mask = ops.logical_or(causal_mask, bidirectional_image_mask)
+
+        # Respect the padding mask.
+        if decoder_mask is not None:
+            causal_mask = ops.minimum(decoder_mask, causal_mask)
+
+        return causal_mask
+
+    def call(
+        self,
+        x,
+        padding_mask=None,
+        text_mask=None,
+        cache=None,
+        cache_update_index=0,
+    ):
+        # Note: `text_mask` is used only for Gemma33.
+        normalized_x = self.pre_attention_norm(x)
+        attention_mask = self._compute_attention_mask(
+            normalized_x, padding_mask, text_mask, cache, cache_update_index
+        )
+        if cache is not None:
+            attention, new_cache = self.attention(
+                normalized_x,
+                attention_mask=attention_mask,
+                cache=cache,
+                cache_update_index=cache_update_index,
+            )
+        else:
+            attention = self.attention(
+                normalized_x,
+                attention_mask=attention_mask,
+            )
+
+        if self.use_post_attention_norm:
+            attention = self.post_attention_norm(attention)
+
+        if self.dropout:
+            attention = self.attention_dropout(attention)
+
+        attention_x = x + attention
+        normalized_x = self.pre_ffw_norm(attention_x)
+
+        x1 = self.gating_ffw(normalized_x)
+        x2 = self.gating_ffw_2(normalized_x)
+        x = keras.activations.gelu(x1, approximate=True) * x2
+        x = self.ffw_linear(x)
+
+        if self.use_post_ffw_norm:
+            x = self.post_ffw_norm(x)
+
+        x = x + attention_x
+
+        if cache is not None:
+            return x, new_cache
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "intermediate_dim": self.intermediate_dim,
+                "head_dim": self.head_dim,
+                "num_query_heads": self.num_query_heads,
+                "num_key_value_heads": self.num_key_value_heads,
+                "query_head_dim_normalize": self.query_head_dim_normalize,
+                "use_query_key_norm": self.use_query_key_norm,
+                "use_post_ffw_norm": self.use_post_ffw_norm,
+                "use_post_attention_norm": self.use_post_attention_norm,
+                "gate_dim_reduction": self.gate_dim_reduction,
+                "logit_soft_cap": self.logit_soft_cap,
+                "use_sliding_window_attention": (
+                    self.use_sliding_window_attention
+                ),
+                "sliding_window_size": self.sliding_window_size,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+                "dropout": self.dropout,
+                "rope_wavelength": self.rope_wavelength,
+                "rope_scaling_factor": self.rope_scaling_factor,
+            }
+        )
+        return config

keras_hub/src/models/gemma3/gemma3_image_converter.py

@@ -0,0 +1,8 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.image_converter import ImageConverter
+from keras_hub.src.models.gemma3.gemma3_backbone import Gemma3Backbone
+
+
+@keras_hub_export("keras_hub.layers.Gemma3ImageConverter")
+class Gemma3ImageConverter(ImageConverter):
+    backbone_cls = Gemma3Backbone

keras_hub/src/models/gemma3/gemma3_interleave_embeddings.py

@@ -0,0 +1,79 @@
+import keras
+from keras import ops
+
+
+class Gemma3InterleaveEmbeddings(keras.layers.Layer):
+    """Places image embeddings in the correct position in an embedding sequence.
+
+    Args:
+        num_vision_tokens_per_image: int. Number of soft tokens per image.
+    """
+
+    def __init__(self, num_vision_tokens_per_image, **kwargs):
+        super().__init__(**kwargs)
+
+        self.num_vision_tokens_per_image = num_vision_tokens_per_image
+
+    def call(self, image_embeddings, text_embeddings, vision_indices):
+        """
+        Integrates image embeddings into a text embedding sequence.
+
+        Args:
+            image_embeddings: Tensor of shape
+                `(batch_size * num_images_per_prompt,
+                num_vision_tokens_per_image, embedding_dim)`.
+            text_embeddings: Tensor of shape
+            `(batch_size, seq_length, embedding_dim)`.
+            text_mask: Boolean tensor of shape `(batch_size, seq_length)`.
+
+        Returns:
+            Tensor of shape `(batch_size, seq_length, embedding_dim)`
+            representing the reconstructed embeddings.
+        """
+
+        batch_size, seq_length, embedding_dim = ops.shape(text_embeddings)
+
+        # Flatten text embeddings, text mask and image embeddings.
+        flat_text_embeddings = ops.reshape(
+            text_embeddings, (batch_size * seq_length, embedding_dim)
+        )
+
+        # The image batch size might be different when we pass only text, i.e,
+        # it will be 0 for text-only.
+        image_batch_size = ops.shape(image_embeddings)[0]
+        flat_image_embeddings = ops.reshape(
+            image_embeddings,
+            (
+                image_batch_size * self.num_vision_tokens_per_image,
+                embedding_dim,
+            ),
+        )
+
+        # Reconstruct embeddings.
+        vision_indices_shape = ops.shape(vision_indices)
+        flat_vision_indices = ops.reshape(
+            vision_indices,
+            (vision_indices_shape[0] * vision_indices_shape[1], 1),
+        )
+        indices = ops.cast(flat_vision_indices, "int32")
+        reconstructed_embedding = ops.scatter_update(
+            flat_text_embeddings, indices, flat_image_embeddings
+        )
+
+        # Reshape to original dimensions
+        reconstructed_embedding = ops.reshape(
+            reconstructed_embedding, (batch_size, seq_length, embedding_dim)
+        )
+        return reconstructed_embedding
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_vision_tokens_per_image": self.num_vision_tokens_per_image,
+            }
+        )
+        return config

keras_hub/src/models/gemma3/gemma3_presets.py

@@ -0,0 +1,93 @@
+"""Gemma3 model preset configurations."""
+
+# Metadata for loading pretrained model weights.
+backbone_presets = {
+    "gemma3_1b": {
+        "metadata": {
+            "description": (
+                "1 billion parameter, 26-layer, text-only pretrained "
+                "Gemma3 model."
+            ),
+            "params": 999885952,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_1b/1",
+    },
+    "gemma3_instruct_1b": {
+        "metadata": {
+            "description": (
+                "1 billion parameter, 26-layer, text-only instruction-tuned "
+                "Gemma3 model."
+            ),
+            "params": 999885952,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_instruct_1b/1",
+    },
+    "gemma3_4b_text": {
+        "metadata": {
+            "description": (
+                "4 billion parameter, 34-layer, text-only pretrained "
+                "Gemma3 model."
+            ),
+            "params": 3880099328,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_4b_text/1",
+    },
+    "gemma3_instruct_4b_text": {
+        "metadata": {
+            "description": (
+                "4 billion parameter, 34-layer, text-only instruction-tuned "
+                "Gemma3 model."
+            ),
+            "params": 3880099328,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_instruct_4b_text/2",
+    },
+    "gemma3_12b_text": {
+        "metadata": {
+            "description": (
+                "12 billion parameter, 48-layer, text-only pretrained "
+                "Gemma3 model."
+            ),
+            "params": 11765788416,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_12b_text/1",
+    },
+    "gemma3_instruct_12b_text": {
+        "metadata": {
+            "description": (
+                "12 billion parameter, 48-layer, text-only instruction-tuned "
+                "Gemma3 model."
+            ),
+            "params": 11765788416,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_instruct_12b_text/1",
+    },
+    "gemma3_27b_text": {
+        "metadata": {
+            "description": (
+                "27 billion parameter, 62-layer, text-only pretrained "
+                "Gemma3 model."
+            ),
+            "params": 27009002240,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_27b_text/1",
+    },
+    "gemma3_instruct_27b_text": {
+        "metadata": {
+            "description": (
+                "27 billion parameter, 62-layer, text-only instruction-tuned "
+                "Gemma3 model."
+            ),
+            "params": 27009002240,
+            "path": "gemma3",
+        },
+        "kaggle_handle": "kaggle://keras/gemma3/keras/gemma3_instruct_27b_text/1",
+    },
+}

keras_hub/src/models/gemma3/gemma3_tokenizer.py

@@ -0,0 +1,87 @@
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.gemma3.gemma3_backbone import Gemma3Backbone
+from keras_hub.src.tokenizers.sentence_piece_tokenizer import (
+    SentencePieceTokenizer,
+)
+
+
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.Gemma3Tokenizer",
+        "keras_hub.models.Gemma3Tokenizer",
+    ]
+)
+class Gemma3Tokenizer(SentencePieceTokenizer):
+    """Gemma tokenizer layer based on SentencePiece.
+
+    This tokenizer class will tokenize raw strings into integer sequences and
+    is based on `keras_hub.tokenizers.SentencePieceTokenizer`. Unlike the
+    underlying tokenizer, it will check for all special tokens needed by
+    Gemma models and provides a `from_preset()` method to automatically
+    download a matching vocabulary for a Gemma preset.
+
+    If input is a batch of strings (rank > 0), the layer will output a
+    `tf.RaggedTensor` where the last dimension of the output is ragged.
+
+    If input is a scalar string (rank == 0), the layer will output a dense
+    `tf.Tensor` with static shape `[None]`.
+
+    Args:
+        proto: Either a `string` path to a SentencePiece proto file, or a
+            `bytes` object with a serialized SentencePiece proto. See the
+            [SentencePiece repository](https://github.com/google/sentencepiece)
+            for more details on the format.
+
+    Examples:
+
+    ```python
+    # Unbatched input.
+    tokenizer = keras_hub.models.Gemma3Tokenizer.from_preset(
+        "gemma_instruct_1b"
+    )
+    tokenizer("The quick brown fox jumped.")
+
+    # Batched input.
+    tokenizer(["The quick brown fox jumped.", "The fox slept."])
+
+    # Detokenization.
+    tokenizer.detokenize(tokenizer("The quick brown fox jumped."))
+
+    # Custom vocabulary.
+    bytes_io = io.BytesIO()
+    ds = tf.data.Dataset.from_tensor_slices(["The quick brown fox jumped."])
+    sentencepiece.SentencePieceTrainer.train(
+        sentence_iterator=ds.as_numpy_iterator(),
+        model_writer=bytes_io,
+        vocab_size=8,
+        model_type="WORD",
+        pad_id=0,
+        bos_id=1,
+        eos_id=2,
+        unk_id=3,
+        pad_piece="<pad>",
+        bos_piece="<bos>",
+        eos_piece="<eos>",
+        unk_piece="<unk>",
+    )
+    tokenizer = keras_hub.models.Gemma3Tokenizer(
+        proto=bytes_io.getvalue(),
+    )
+    tokenizer("The quick brown fox jumped.")
+    ```
+    """
+
+    backbone_cls = Gemma3Backbone
+
+    def __init__(self, proto, **kwargs):
+        # Add special tokens.
+
+        # The usual tokens.
+        self._add_special_token("<bos>", "start_token")
+        self._add_special_token("<eos>", "end_token")
+        self._add_special_token("<pad>", "pad_token")
+
+        # Image placeholder token.
+        self._add_special_token("<img>", "image_placeholder")
+
+        super().__init__(proto=proto, **kwargs)