keras-hub-nightly 0.19.0.dev202501280343__tar.gz → 0.19.0.dev202501300343__tar.gz

{keras_hub_nightly-0.19.0.dev202501280343 → keras_hub_nightly-0.19.0.dev202501300343}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: keras-hub-nightly
-Version: 0.19.0.dev202501280343
+Version: 0.19.0.dev202501300343
 Summary: Industry-strength Natural Language Processing extensions for Keras.
 Home-page: https://github.com/keras-team/keras-hub
 Author: Keras team
@@ -123,7 +123,7 @@ print(keras_hub.utils.decode_imagenet_predictions(preds))
 Load a Bert model and fine-tune it on IMDb movie reviews:
 
 ```python
-classifier = keras_hub.models.BertClassifier.from_preset(
+classifier = keras_hub.models.TextClassifier.from_preset(
     "bert_base_en_uncased",
     activation="softmax",
     num_classes=2,

{keras_hub_nightly-0.19.0.dev202501280343 → keras_hub_nightly-0.19.0.dev202501300343}/README.md

@@ -78,7 +78,7 @@ print(keras_hub.utils.decode_imagenet_predictions(preds))
 Load a Bert model and fine-tune it on IMDb movie reviews:
 
 ```python
-classifier = keras_hub.models.BertClassifier.from_preset(
+classifier = keras_hub.models.TextClassifier.from_preset(
     "bert_base_en_uncased",
     activation="softmax",
     num_classes=2,

{keras_hub_nightly-0.19.0.dev202501280343 → keras_hub_nightly-0.19.0.dev202501300343}/keras_hub/src/models/falcon/falcon_attention.py

@@ -110,9 +110,11 @@ class FalconAttention(keras.layers.Layer):
 
         attention_scores = ops.einsum("bqnh,bknh->bnqk", query, key)
         attention_scores = ops.add(attention_scores, alibi)
-        attention_scores = (
-            attention_scores * self.inv_norm_factor
-        )  # [batch_size, num_heads, query_length, kv_length]
+        # [batch_size, num_heads, query_length, kv_length]
+        attention_scores = ops.multiply(
+            attention_scores,
+            ops.cast(self.inv_norm_factor, self.compute_dtype),
+        )
         attention_scores = self.softmax(
             attention_scores, ops.expand_dims(attention_mask, 1)
         )
@@ -120,6 +122,7 @@ class FalconAttention(keras.layers.Layer):
         attention_output = ops.einsum(
             "bnqk,bknh->bqnh", attention_scores, value
         )
+
         attention_output = ops.reshape(
             attention_output,
             [batch_size, seq_length, self.num_heads * self.head_dim],

{keras_hub_nightly-0.19.0.dev202501280343 → keras_hub_nightly-0.19.0.dev202501300343}/keras_hub/src/models/gpt_neo_x/gpt_neo_x_attention.py

@@ -1,8 +1,11 @@
+import math
+
 import keras
 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 GPTNeoXAttention(keras.layers.Layer):
@@ -58,6 +61,8 @@ class GPTNeoXAttention(keras.layers.Layer):
         self.bias_initializer = keras.initializers.get(bias_initializer)
         self.max_sequence_length = max_sequence_length
 
+        self._inv_norm_factor = 1.0 / math.sqrt(self.attn_head_size)
+
     def build(self, input_shape):
         self._qkv_dense = keras.layers.EinsumDense(
             equation="abc,cde->abde",
@@ -120,14 +125,26 @@ class GPTNeoXAttention(keras.layers.Layer):
     def _compute_attention(
         self, query, key, value, attention_mask=None, training=None
     ):
-        attention_scores = ops.einsum("aecd,abcd->acbe", key, query)
+        if has_flash_attention_support() and self.dropout == 0:
+            # Use `dot_product_attention` with Flash Attention support if
+            # available.
+            if attention_mask is not None:
+                attention_mask = ops.expand_dims(attention_mask, axis=1)
+                attention_mask = ops.cast(attention_mask, dtype="bool")
+            attention_output = ops.dot_product_attention(
+                query,
+                key,
+                value,
+                mask=attention_mask,
+                scale=self._inv_norm_factor,
+            )
+            return attention_output
 
-        norm_factor = ops.sqrt(
-            ops.convert_to_tensor(self.attn_head_size, self.compute_dtype)
+        attention_scores = ops.einsum("aecd,abcd->acbe", key, query)
+        attention_scores = ops.multiply(
+            attention_scores,
+            ops.cast(self._inv_norm_factor, self.compute_dtype),
         )
-
-        attention_scores /= norm_factor
-
         attention_scores = self._masked_softmax(
             attention_scores, attention_mask
         )

{keras_hub_nightly-0.19.0.dev202501280343 → keras_hub_nightly-0.19.0.dev202501300343}/keras_hub/src/models/llama/llama_attention.py

@@ -1,8 +1,11 @@
+import math
+
 import keras
 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 LlamaAttention(keras.layers.Layer):
@@ -43,7 +46,7 @@ class LlamaAttention(keras.layers.Layer):
         # h = head dim
         hidden_dim = inputs_shape[-1]
         head_dim = hidden_dim // self.num_query_heads
-        self._norm_factor = ops.sqrt(ops.cast(head_dim, self.compute_dtype))
+        self._inv_norm_factor = 1.0 / math.sqrt(head_dim)
 
         self._query_dense = keras.layers.EinsumDense(
             equation="bqm,muh->bquh",
@@ -182,9 +185,27 @@ class LlamaAttention(keras.layers.Layer):
         return self._softmax(attention_scores)
 
     def _compute_attention(self, query, key, value, attention_mask=None):
+        if has_flash_attention_support():
+            # Use `dot_product_attention` with Flash Attention support if
+            # available.
+            if attention_mask is not None:
+                attention_mask = ops.expand_dims(attention_mask, axis=1)
+                attention_mask = ops.cast(attention_mask, dtype="bool")
+            attention_output = ops.dot_product_attention(
+                query,
+                key,
+                value,
+                mask=attention_mask,
+                scale=self._inv_norm_factor,
+            )
+            return attention_output
+
         attention_scores = ops.einsum(self._dot_product_equation, query, key)
 
-        attention_scores = attention_scores / self._norm_factor
+        attention_scores = ops.multiply(
+            attention_scores,
+            ops.cast(self._inv_norm_factor, self.compute_dtype),
+        )
         attention_scores = self._masked_softmax(
             attention_scores, attention_mask
         )

{keras_hub_nightly-0.19.0.dev202501280343 → keras_hub_nightly-0.19.0.dev202501300343}/keras_hub/src/models/mistral/mistral_attention.py

@@ -1,8 +1,11 @@
+import math
+
 import keras
 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
 
 
 # This is just a self-attention layer in Mistral. But it can be generalized
@@ -52,6 +55,7 @@ class CachedMistralAttention(keras.layers.Layer):
         # h = head dim
         self._hidden_dim = inputs_shape[-1]
         self._head_dim = self._hidden_dim // self._num_query_heads
+        self._inv_norm_factor = 1.0 / math.sqrt(self._head_dim)
 
         self._query_dense = keras.layers.EinsumDense(
             equation="bqm,muh->bquh",
@@ -192,11 +196,26 @@ class CachedMistralAttention(keras.layers.Layer):
         return self._softmax(attention_scores)
 
     def _compute_attention(self, query, key, value, attention_mask=None):
-        attention_scores = ops.einsum(self._dot_product_equation, query, key)
-
-        norm_factor = ops.sqrt(ops.cast(self._head_dim, self.compute_dtype))
+        if has_flash_attention_support():
+            # Use `dot_product_attention` with Flash Attention support if
+            # available.
+            if attention_mask is not None:
+                attention_mask = ops.expand_dims(attention_mask, axis=1)
+                attention_mask = ops.cast(attention_mask, dtype="bool")
+            attention_output = ops.dot_product_attention(
+                query,
+                key,
+                value,
+                mask=attention_mask,
+                scale=self._inv_norm_factor,
+            )
+            return attention_output
 
-        attention_scores = attention_scores / norm_factor
+        attention_scores = ops.einsum(self._dot_product_equation, query, key)
+        attention_scores = ops.multiply(
+            attention_scores,
+            ops.cast(self._inv_norm_factor, self.compute_dtype),
+        )
         attention_scores = self._masked_softmax(
             attention_scores, attention_mask
         )

{keras_hub_nightly-0.19.0.dev202501280343 → keras_hub_nightly-0.19.0.dev202501300343}/keras_hub/src/models/phi3/phi3_attention.py

@@ -1,3 +1,5 @@
+import math
+
 import keras
 from keras import ops
 
@@ -6,6 +8,7 @@ from keras_hub.src.models.phi3.phi3_rotary_embedding import (
     Phi3SuScaledRotaryEmbedding,
 )
 from keras_hub.src.utils.keras_utils import clone_initializer
+from keras_hub.src.utils.keras_utils import has_flash_attention_support
 
 
 class Phi3Attention(keras.layers.Layer):
@@ -53,7 +56,7 @@ class Phi3Attention(keras.layers.Layer):
         # h = head dim
         hidden_dim = inputs_shape[-1]
         head_dim = hidden_dim // self.num_query_heads
-        self._norm_factor = ops.sqrt(ops.cast(head_dim, self.compute_dtype))
+        self._inv_norm_factor = 1.0 / math.sqrt(head_dim)
 
         self.query_dense = keras.layers.EinsumDense(
             equation="bqm,muh->bquh",
@@ -214,8 +217,26 @@ class Phi3Attention(keras.layers.Layer):
         return self.softmax(attention_scores)
 
     def _compute_attention(self, query, key, value, attention_mask=None):
+        if has_flash_attention_support():
+            # Use `dot_product_attention` with Flash Attention support if
+            # available.
+            if attention_mask is not None:
+                attention_mask = ops.expand_dims(attention_mask, axis=1)
+                attention_mask = ops.cast(attention_mask, dtype="bool")
+            attention_output = ops.dot_product_attention(
+                query,
+                key,
+                value,
+                mask=attention_mask,
+                scale=self._inv_norm_factor,
+            )
+            return attention_output
+
         attention_scores = ops.einsum("bquh,bkuh->buqk", query, key)
-        attention_scores = attention_scores / self._norm_factor
+        attention_scores = ops.multiply(
+            attention_scores,
+            ops.cast(self._inv_norm_factor, self.compute_dtype),
+        )
         attention_scores = self._masked_softmax(
             attention_scores, attention_mask
         )

{keras_hub_nightly-0.19.0.dev202501280343 → keras_hub_nightly-0.19.0.dev202501300343}/keras_hub/src/models/stable_diffusion_3/mmdit.py

@@ -7,6 +7,7 @@ from keras import ops
 from keras_hub.src.layers.modeling.position_embedding import PositionEmbedding
 from keras_hub.src.models.backbone import Backbone
 from keras_hub.src.utils.keras_utils import gelu_approximate
+from keras_hub.src.utils.keras_utils import has_flash_attention_support
 from keras_hub.src.utils.keras_utils import standardize_data_format
 
 
@@ -770,17 +771,14 @@ class MMDiTBlock(layers.Layer):
     def _compute_attention(self, query, key, value):
         batch_size = ops.shape(query)[0]
 
-        # Use the fast path when `ops.dot_product_attention` and flash attention
-        # are available.
-        if hasattr(ops, "dot_product_attention") and hasattr(
-            keras.config, "is_flash_attention_enabled"
-        ):
+        if has_flash_attention_support():
+            # Use `dot_product_attention` with Flash Attention support if
+            # available.
             encoded = ops.dot_product_attention(
                 query,
                 key,
                 value,
                 scale=self._inverse_sqrt_key_dim,
-                flash_attention=keras.config.is_flash_attention_enabled(),
             )
             return ops.reshape(
                 encoded, (batch_size, -1, self.num_heads * self.head_dim)
@@ -793,10 +791,9 @@ class MMDiTBlock(layers.Layer):
         probs = self.softmax(logits)
         probs = ops.cast(probs, self.compute_dtype)
         encoded = ops.einsum("BNTS,BSNH->BTNH", probs, value)
-        encoded = ops.reshape(
+        return ops.reshape(
             encoded, (batch_size, -1, self.num_heads * self.head_dim)
         )
-        return encoded
 
     def call(self, inputs, context, timestep_embedding, training=None):
         # Compute pre-attention.

{keras_hub_nightly-0.19.0.dev202501280343 → keras_hub_nightly-0.19.0.dev202501300343}/keras_hub/src/utils/keras_utils.py

@@ -53,3 +53,10 @@ def standardize_data_format(data_format):
             f"Received: data_format={data_format}"
         )
     return data_format
+
+
+def has_flash_attention_support():
+    if hasattr(keras.config, "is_flash_attention_enabled"):
+        return True
+    else:
+        return False

{keras_hub_nightly-0.19.0.dev202501280343 → keras_hub_nightly-0.19.0.dev202501300343}/keras_hub/src/utils/preset_utils.py

@@ -240,7 +240,7 @@ def tf_copy_gfile_to_cache(preset, path):
         try:
             import tensorflow as tf
 
-            os.make_dirs(os.path.dirname(local_path), exist_ok=True)
+            os.makedirs(os.path.dirname(local_path), exist_ok=True)
             tf.io.gfile.copy(url, local_path)
         except Exception as e:
             # gfile.copy will leave an empty file after an error.

{keras_hub_nightly-0.19.0.dev202501280343 → keras_hub_nightly-0.19.0.dev202501300343}/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.19.0.dev202501280343"
+__version__ = "0.19.0.dev202501300343"
 
 
 @keras_hub_export("keras_hub.version")

{keras_hub_nightly-0.19.0.dev202501280343 → keras_hub_nightly-0.19.0.dev202501300343}/keras_hub_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: keras-hub-nightly
-Version: 0.19.0.dev202501280343
+Version: 0.19.0.dev202501300343
 Summary: Industry-strength Natural Language Processing extensions for Keras.
 Home-page: https://github.com/keras-team/keras-hub
 Author: Keras team
@@ -123,7 +123,7 @@ print(keras_hub.utils.decode_imagenet_predictions(preds))
 Load a Bert model and fine-tune it on IMDb movie reviews:
 
 ```python
-classifier = keras_hub.models.BertClassifier.from_preset(
+classifier = keras_hub.models.TextClassifier.from_preset(
     "bert_base_en_uncased",
     activation="softmax",
     num_classes=2,

{keras_hub_nightly-0.19.0.dev202501280343 → keras_hub_nightly-0.19.0.dev202501300343}/setup.py

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