PyPI - keras-hub-nightly - Versions diffs - 0.19.0.dev202502200346__tar.gz → 0.19.0.dev202502220345__tar.gz - Mend

Files changed (424) hide show

{keras_hub_nightly-0.19.0.dev202502200346 → keras_hub_nightly-0.19.0.dev202502220345}/keras_hub/src/models/pali_gemma/pali_gemma_presets.py RENAMED Viewed

@@ -83,6 +83,96 @@ backbone_presets = {
         },
         "kaggle_handle": "kaggle://keras/paligemma2/keras/pali_gemma2_ft_docci_10b_448/2",
     },
+    "pali_gemma2_mix_3b_224": {
+        "metadata": {
+            "description": (
+                "3 billion parameter, image size 224, 27-layer for "
+                "SigLIP-So400m vision encoder and 26-layer Gemma2 2B lanuage "
+                "model. This model has been fine-tuned on a wide range of "
+                "vision-language tasks and domains."
+            ),
+            "params": 3032094960,
+            "official_name": "PaliGemma2",
+            "path": "pali_gemma2",
+            "model_card": "https://www.kaggle.com/models/google/paligemma-2",
+        },
+        "kaggle_handle": "kaggle://keras/paligemma2/keras/pali_gemma2_mix_3b_224/2",
+    },
+    "pali_gemma2_mix_3b_448": {
+        "metadata": {
+            "description": (
+                "3 billion parameter, image size 448, 27-layer for "
+                "SigLIP-So400m vision encoder and 26-layer Gemma2 2B lanuage "
+                "model. This model has been fine-tuned on a wide range of "
+                "vision-language tasks and domains."
+            ),
+            "params": 3032979696,
+            "official_name": "PaliGemma2",
+            "path": "pali_gemma2",
+            "model_card": "https://www.kaggle.com/models/google/paligemma-2",
+        },
+        "kaggle_handle": "kaggle://keras/paligemma2/keras/pali_gemma2_mix_3b_448/2",
+    },
+    "pali_gemma2_mix_10b_224": {
+        "metadata": {
+            "description": (
+                "10 billion parameter, image size 224, 27-layer for "
+                "SigLIP-So400m vision encoder and 42-layer Gemma2 9B lanuage "
+                "model. This model has been fine-tuned on a wide range of "
+                "vision-language tasks and domains."
+            ),
+            "params": 9662409456,
+            "official_name": "PaliGemma2",
+            "path": "pali_gemma2",
+            "model_card": "https://www.kaggle.com/models/google/paligemma-2",
+        },
+        "kaggle_handle": "kaggle://keras/paligemma2/keras/pali_gemma2_mix_10b_224/2",
+    },
+    "pali_gemma2_mix_10b_448": {
+        "metadata": {
+            "description": (
+                "10 billion parameter, image size 448, 27-layer for "
+                "SigLIP-So400m vision encoder and 42-layer Gemma2 9B lanuage "
+                "model. This model has been fine-tuned on a wide range of "
+                "vision-language tasks and domains."
+            ),
+            "params": 9663294192,
+            "official_name": "PaliGemma2",
+            "path": "pali_gemma2",
+            "model_card": "https://www.kaggle.com/models/google/paligemma-2",
+        },
+        "kaggle_handle": "kaggle://keras/paligemma2/keras/pali_gemma2_mix_10b_448/2",
+    },
+    "pali_gemma2_mix_28b_224": {
+        "metadata": {
+            "description": (
+                "28 billion parameter, image size 224, 27-layer for "
+                "SigLIP-So400m vision encoder and 46-layer Gemma2 27B lanuage "
+                "model. This model has been fine-tuned on a wide range of "
+                "vision-language tasks and domains."
+            ),
+            "params": 27650192112,
+            "official_name": "PaliGemma2",
+            "path": "pali_gemma2",
+            "model_card": "https://www.kaggle.com/models/google/paligemma-2",
+        },
+        "kaggle_handle": "kaggle://keras/paligemma2/keras/pali_gemma2_28b_mix_224/2",
+    },
+    "pali_gemma2_mix_28b_448": {
+        "metadata": {
+            "description": (
+                "28 billion parameter, image size 448, 27-layer for "
+                "SigLIP-So400m vision encoder and 46-layer Gemma2 27B lanuage "
+                "model. This model has been fine-tuned on a wide range of "
+                "vision-language tasks and domains."
+            ),
+            "params": 27650192112,
+            "official_name": "PaliGemma2",
+            "path": "pali_gemma2",
+            "model_card": "https://www.kaggle.com/models/google/paligemma-2",
+        },
+        "kaggle_handle": "kaggle://keras/paligemma2/keras/pali_gemma2_28b_mix_448/2",
+    },
     "pali_gemma2_pt_3b_224": {
         "metadata": {
             "description": (
@@ -181,7 +271,7 @@ backbone_presets = {
                 "model. This model has been pre-trained on a mixture of "
                 "datasets."
             ),
-            "params": 9662409456,
+            "params": 27650192112,
             "official_name": "PaliGemma2",
             "path": "pali_gemma2",
             "model_card": "https://www.kaggle.com/models/google/paligemma-2",
@@ -196,7 +286,7 @@ backbone_presets = {
                 "model. This model has been pre-trained on a mixture of "
                 "datasets."
             ),
-            "params": 9663294192,
+            "params": 27650192112,
             "official_name": "PaliGemma2",
             "path": "pali_gemma2",
             "model_card": "https://www.kaggle.com/models/google/paligemma-2",
@@ -211,7 +301,7 @@ backbone_presets = {
                 "model. This model has been pre-trained on a mixture of "
                 "datasets."
             ),
-            "params": 9666833136,
+            "params": 27650192112,
             "official_name": "PaliGemma2",
             "path": "pali_gemma2",
             "model_card": "https://www.kaggle.com/models/google/paligemma-2",

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: keras-hub-nightly
-Version: 0.19.0.dev202502200346
+Version: 0.19.0.dev202502220345
 Summary: Industry-strength Natural Language Processing extensions for Keras.
 Home-page: https://github.com/keras-team/keras-hub
 Author: Keras team

@@ -186,6 +186,14 @@ class Backbone(keras.Model):
         saver = get_preset_saver(preset_dir)
         saver.save_backbone(self)
+    def get_lora_target_names(self):
+        """Returns list of layer names which are to be LoRA-fied.
+        Subclasses can override this method if the names of layers to be
+        LoRa-fied are different.
+        """
+        return ["query_dense", "value_dense", "query", "value"]
     def enable_lora(self, rank):
         """Enable Lora on the backbone.
@@ -193,7 +201,8 @@ class Backbone(keras.Model):
         while enabling Lora on the query & value `EinsumDense` layers
         of the attention layers.
         """
-        target_names = ["query_dense", "value_dense", "query", "value"]
+        target_names = self.get_lora_target_names()
         self.trainable = True
         self._lora_enabled_layers = []
         self._lora_rank = rank

@@ -4,6 +4,7 @@ from keras import ops
 from keras_hub.src.layers.modeling.rotary_embedding import RotaryEmbedding
 from keras_hub.src.utils.keras_utils import clone_initializer
+from keras_hub.src.utils.keras_utils import has_flash_attention_support
 class CachedGemmaAttention(keras.layers.Layer):
@@ -117,6 +118,33 @@ class CachedGemmaAttention(keras.layers.Layer):
             query_normalization = 1 / np.sqrt(
                 self.hidden_dim // self.num_query_heads
             )
+        use_dot_product_attention = not (
+            self.dropout > 0.0 or (len(q.shape) != 4)
+        )
+        if has_flash_attention_support() and use_dot_product_attention:
+            if self.dropout > 0.0:
+                raise ValueError(
+                    "Flash attention does not support dropout. "
+                    "Please set `dropout` to 0.0."
+                )
+            if attention_mask is not None:
+                while len(attention_mask.shape) < 4:
+                    attention_mask = ops.expand_dims(
+                        attention_mask, axis=1
+                    )  # Add dimension for num_heads
+                if attention_mask.shape[1] != self.num_query_heads:
+                    attention_mask = ops.tile(
+                        attention_mask, [1, self.num_query_heads, 1, 1]
+                    )
+            attention_output = ops.dot_product_attention(
+                query=q,
+                key=k,
+                value=v,
+                mask=attention_mask,
+                scale=query_normalization,
+            )
+            return attention_output
         q *= ops.cast(query_normalization, dtype=q.dtype)
         q_shape = ops.shape(q)
@@ -131,8 +159,8 @@ class CachedGemmaAttention(keras.layers.Layer):
         )
         b, q_len, _, _, h = ops.shape(q)
+        # Fallback to standard attention if flash attention is disabled
         attention_logits = ops.einsum("btkgh,bskh->bkgts", q, k)
         if self.logit_soft_cap is not None:
             attention_logits = ops.divide(attention_logits, self.logit_soft_cap)
             attention_logits = ops.multiply(

@@ -274,6 +274,13 @@ class PaliGemmaBackbone(Backbone):
         # Keep the image_sequence_length as a backbone property for easy access.
         self.image_sequence_length = self.vit_encoder.image_sequence_length
+    def get_lora_target_names(self):
+        target_names = super().get_lora_target_names()
+        # Add these for `PaliGemmaVITAttention`.
+        target_names += ["query_proj", "value_proj"]
+        return target_names
     def get_config(self):
         config = super().get_config()
         config.update(

@@ -56,7 +56,19 @@ def standardize_data_format(data_format):
 def has_flash_attention_support():
-    if hasattr(keras.config, "is_flash_attention_enabled"):
+    if (
+        hasattr(keras.config, "is_flash_attention_enabled")
+        and keras.config.backend() == "jax"
+    ):
+        try:
+            from jax.nn import dot_product_attention as dot_product_attention
+        except ImportError:
+            logging.warning(
+                "Flash attention is not supported in your current JAX version. "
+                "Please update it by following the official guide: "
+                "https://jax.readthedocs.io/en/latest/installation.html"
+            )
+            return False
         return True
     else:
         return False

@@ -1,7 +1,7 @@
 from keras_hub.src.api_export import keras_hub_export
 # Unique source of truth for the version number.
-__version__ = "0.19.0.dev202502200346"
+__version__ = "0.19.0.dev202502220345"
 @keras_hub_export("keras_hub.version")

@@ -23,7 +23,7 @@ def get_version(rel_path):
 HERE = pathlib.Path(__file__).parent
 README = (HERE / "README.md").read_text()
-VERSION = "0.19.0.dev202502200346"  # get_version("keras_hub/src/version_utils.py")
+VERSION = "0.19.0.dev202502220345"  # get_version("keras_hub/src/version_utils.py")
 setup(
     name="keras-hub-nightly",  # "keras-hub",

keras-hub-nightly 0.19.0.dev202502200346__tar.gz → 0.19.0.dev202502220345__tar.gz

keras-hub-nightly 0.19.0.dev202502200346tar.gz → 0.19.0.dev202502220345tar.gz