Feature/pg gemma changes

Co-authored-by: divyashreepathihalli <[email protected]>

keras_nlp/src/models/paligemma/pali_gemma_decoder.py → keras_nlp/src/models/paligemma/pali_gemma_backbone.py

@@ -16,16 +16,18 @@
 from keras_nlp.src.layers.modeling.reversible_embedding import (
     ReversibleEmbedding,
 )
-from keras_nlp.src.models.gemma.gemma_decoder_block import GemmaDecoderBlock
-from keras_nlp.src.models.gemma.gemma_tokenizer import GemmaTokenizer
+from keras_nlp.src.models.backbone import Backbone
 from keras_nlp.src.models.gemma.rms_normalization import RMSNormalization
+from keras_nlp.src.models.paligemma.pali_gemma_decoder_block import (
+    PaliGemmaDecoderBlock,
+)
 
 
-class PaliGemmaDecoder(keras.layers.Layer):
+class PaliGemmaBackbone(Backbone):
     def __init__(
         self,
+        img_sequence_length,
         vocabulary_size,
-        sequence_length,
         num_layers,
         num_query_heads,
         num_key_value_heads,
@@ -34,13 +36,11 @@ def __init__(
         head_dim,
         layer_norm_epsilon=1e-6,
         dropout=0,
-        tokenizer_preset="gemma_2b_en",
-        text_prefix="answer en",
         dtype=None,
         **kwargs,
     ):
+        self.img_sequence_length = img_sequence_length
         self.vocabulary_size = vocabulary_size
-        self.sequence_length = sequence_length
         self.num_layers = num_layers
         self.num_query_heads = num_query_heads
         self.num_key_value_heads = num_key_value_heads
@@ -49,15 +49,10 @@ def __init__(
         self.head_dim = head_dim
         self.layer_norm_epsilon = layer_norm_epsilon
         self.dropout = dropout
-        self.tokenizer_preset = tokenizer_preset
-        self.text_prefix = text_prefix
-        self.dtype = dtype
 
         #
         # Layers
         #
-        self.tokenizer = GemmaTokenizer.from_preset(self.tokenizer_preset)
-
         self.token_embedding = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
@@ -71,11 +66,13 @@ def __init__(
             dtype=dtype,
             name="token_embedding",
         )
+
         self.transformer_layers = []
         for i in range(num_layers):
-            layer = GemmaDecoderBlock(
-                intermediate_dim=intermediate_dim,
+            layer = PaliGemmaDecoderBlock(
+                img_sequence_length=img_sequence_length,
                 hidden_dim=hidden_dim,
+                intermediate_dim=intermediate_dim,
                 num_query_heads=num_query_heads,
                 head_dim=head_dim,
                 num_key_value_heads=num_key_value_heads,
@@ -94,42 +91,50 @@ def __init__(
         # Functional Model
         #
         img_embeddings = keras.Input(
-            shape=(None,), dtype=dtype, name="img_embeddings"
+            shape=(img_sequence_length, hidden_dim),
+            dtype=dtype,
+            name="img_embeddings",
         )
 
-        # TODO: Is there a good data type for text/string input like this?
-        text_in = keras.Input(
-            shape=(sequence_length, vocabulary_size), name="text"
+        token_ids = keras.Input(
+            shape=(None,),
+            dtype="int32",
+            name="token_ids",
         )
 
-        prefixed_text = [self.text_prefix + " " + text for text in text_in]
+        padding_mask = keras.Input(
+            shape=(None,), dtype="float32", name="padding_mask"
+        )
 
-        tokenized_text = self.tokenizer(prefixed_text)
-        text_embeddings = self.token_embedding(tokenized_text)
+        text_embeddings = self.token_embedding(token_ids)
 
         complete_sequence = keras.ops.concatenate(
             (img_embeddings, text_embeddings), axis=1
         )
 
         transformer_out = complete_sequence
         for transformer_layer in self.transformer_layers:
-            transformer_out = transformer_layer(transformer_out)
+            transformer_out = transformer_layer(
+                transformer_out, padding_mask=padding_mask
+            )
 
         text_out = self.layer_norm(transformer_out)
 
         super().__init__(
             inputs={
                 "img_embeddings": img_embeddings,
-                "text": text_in,
+                "token_ids": token_ids,
+                "padding_mask": padding_mask,
             },
-            outputs={"text_out": text_out},
+            outputs=text_out,
             **kwargs,
         )
 
     def get_config(self):
         config = super().get_config()
         config.update(
             {
+                "img_sequence_length": self.img_sequence_length,
                 "vocabulary_size": self.vocabulary_size,
                 "num_layers": self.num_layers,
                 "num_query_heads": self.num_query_heads,
@@ -139,9 +144,6 @@ def get_config(self):
                 "head_dim": self.head_dim,
                 "layer_norm_epsilon": self.layer_norm_epsilon,
                 "dropout": self.dropout,
-                "tokenizer_preset": self.tokenizer_preset,
-                "text_prefix": self.text_prefix,
-                "dtype": self.dtype,
             }
         )
         return config

keras_nlp/src/models/paligemma/pali_gemma_backbone_test.py

@@ -0,0 +1,151 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+
+import keras
+import numpy as np
+import pytest
+
+from keras_nlp.src.models.gemma.gemma_preprocessor import GemmaPreprocessor
+from keras_nlp.src.models.gemma.gemma_preprocessor import GemmaTokenizer
+from keras_nlp.src.models.paligemma.pali_gemma_backbone import PaliGemmaBackbone
+from keras_nlp.src.tests.test_case import TestCase
+
+
+@pytest.mark.keras_3_only
+class PaliGemmaBackboneTest(TestCase):
+    def test_paligemma_preprocessing(self):
+        batch_size = 1
+        text_sequence_length = 8
+        proto = os.path.join(self.get_test_data_dir(), "gemma_test_vocab.spm")
+
+        tokenizer = GemmaTokenizer(proto)
+
+        preprocessor = GemmaPreprocessor(
+            tokenizer, text_sequence_length, False, False
+        )
+
+        dummy_text = ["answer en the quick brown fox"]
+
+        output = preprocessor(dummy_text)
+        self.assertEqual(
+            (
+                batch_size,
+                text_sequence_length,
+            ),
+            output["token_ids"].shape,
+        )
+
+    def test_paligemma_backbone(self):
+        batch_size = 2
+        img_sequence_length = 128
+        vocabulary_size = 256
+        text_sequence_length = 64
+        hidden_dim = 128
+        num_layers = 27
+        num_heads = 16
+        head_dim = 126
+        intermediate_size = 77
+
+        paligemma = PaliGemmaBackbone(
+            img_sequence_length,
+            vocabulary_size,
+            num_layers,
+            num_heads,
+            num_heads,
+            hidden_dim,
+            intermediate_size,
+            head_dim,
+            dtype="float32",
+        )
+
+        dummy_imgs = np.random.rand(batch_size, img_sequence_length, hidden_dim)
+        dummy_text = np.random.rand(batch_size, text_sequence_length)
+
+        output = paligemma(
+            inputs={
+                "token_ids": dummy_text,
+                "img_embeddings": dummy_imgs,
+                "padding_mask": np.ones(
+                    (batch_size, text_sequence_length + img_sequence_length),
+                    dtype="int32",
+                ),
+            }
+        )
+        self.assertEqual(
+            (
+                batch_size,
+                text_sequence_length + img_sequence_length,
+                hidden_dim,
+            ),
+            output.shape,
+        )
+
+    def test_complete_paligemma_backbone(self):
+        batch_size = 2
+        img_sequence_length = 128
+        vocabulary_size = 256
+        text_sequence_length = 64
+        hidden_dim = 128
+        num_layers = 27
+        num_heads = 16
+        head_dim = 126
+        intermediate_size = 77
+        proto = os.path.join(self.get_test_data_dir(), "gemma_test_vocab.spm")
+
+        tokenizer = GemmaTokenizer(proto)
+
+        preprocessor = GemmaPreprocessor(
+            tokenizer, text_sequence_length, False, False
+        )
+
+        paligemma = PaliGemmaBackbone(
+            img_sequence_length,
+            vocabulary_size,
+            num_layers,
+            num_heads,
+            num_heads,
+            hidden_dim,
+            intermediate_size,
+            head_dim,
+            dtype="float32",
+        )
+
+        dummy_imgs = keras.ops.convert_to_tensor(
+            np.random.rand(batch_size, img_sequence_length, hidden_dim)
+        )
+        dummy_text = [
+            "answer en the quick brown fox" for i in range(batch_size)
+        ]
+
+        output = preprocessor(dummy_text)
+
+        output = paligemma(
+            inputs={
+                "token_ids": output["token_ids"],
+                "img_embeddings": dummy_imgs,
+                "padding_mask": keras.ops.ones(
+                    (batch_size, text_sequence_length + img_sequence_length),
+                    dtype="int32",
+                ),
+            }
+        )
+        self.assertEqual(
+            (
+                batch_size,
+                text_sequence_length + img_sequence_length,
+                hidden_dim,
+            ),
+            output.shape,
+        )

keras_nlp/src/models/paligemma/pali_gemma_decoder_block.py

@@ -0,0 +1,106 @@
+# Copyright 2023 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from keras_nlp.src.backend import ops
+from keras_nlp.src.layers.modeling.transformer_layer_utils import (
+    compute_causal_mask,
+)
+from keras_nlp.src.layers.modeling.transformer_layer_utils import (
+    merge_padding_and_attention_mask,
+)
+from keras_nlp.src.models.gemma.gemma_decoder_block import GemmaDecoderBlock
+
+
+class PaliGemmaDecoderBlock(GemmaDecoderBlock):
+    def __init__(
+        self,
+        img_sequence_length,
+        hidden_dim,
+        intermediate_dim,
+        head_dim,
+        num_query_heads,
+        num_key_value_heads,
+        layer_norm_epsilon=1e-6,
+        dropout=0,
+        **kwargs,
+    ):
+        super().__init__(
+            hidden_dim,
+            intermediate_dim,
+            head_dim,
+            num_query_heads,
+            num_key_value_heads,
+            layer_norm_epsilon,
+            dropout,
+            **kwargs,
+        )
+
+        self.img_sequence_length = img_sequence_length
+
+    def _compute_attention_mask(
+        self, x, padding_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,
+        )
+
+        # Image Sequence Embeddings should be fully self-attended without causality
+        img_causal_mask = ops.concatenate(
+            [
+                ops.ones((batch_size, output_length, self.img_sequence_length)),
+                ops.zeros(
+                    (
+                        batch_size,
+                        output_length,
+                        input_length - self.img_sequence_length,
+                    )
+                ),
+            ],
+            axis=-1,
+        )
+
+        causal_mask = ops.maximum(causal_mask, img_causal_mask)
+
+        return (
+            ops.minimum(decoder_mask, causal_mask)
+            if decoder_mask is not None
+            else causal_mask
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "img_sequence_length": self.img_sequence_length,
+                "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,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+                "dropout": self.dropout,
+            }
+        )
+        return config