keras-hub-nightly 0.16.1.dev202410020340__py3-none-any.whl → 0.19.0.dev202501260345__py3-none-any.whl

keras_hub/src/models/vit/vit_presets.py

@@ -0,0 +1,126 @@
+"""ViT model preset configurations."""
+
+# Metadata for loading pretrained model weights.
+backbone_presets = {
+    "vit_base_patch16_224_imagenet": {
+        "metadata": {
+            "description": (
+                "ViT-B16 model pre-trained on the ImageNet 1k dataset with "
+                "image resolution of 224x224 "
+            ),
+            "params": 85798656,
+            "path": "vit",
+        },
+        "kaggle_handle": "kaggle://keras/vit/keras/vit_base_patch16_224_imagenet/2",
+    },
+    "vit_base_patch16_384_imagenet": {
+        "metadata": {
+            "description": (
+                "ViT-B16 model pre-trained on the ImageNet 1k dataset with "
+                "image resolution of 384x384 "
+            ),
+            "params": 86090496,
+            "path": "vit",
+        },
+        "kaggle_handle": "kaggle://keras/vit/keras/vit_base_patch16_384_imagenet/2",
+    },
+    "vit_large_patch16_224_imagenet": {
+        "metadata": {
+            "description": (
+                "ViT-L16 model pre-trained on the ImageNet 1k dataset with "
+                "image resolution of 224x224 "
+            ),
+            "params": 303301632,
+            "path": "vit",
+        },
+        "kaggle_handle": "kaggle://keras/vit/keras/vit_large_patch16_224_imagenet/2",
+    },
+    "vit_large_patch16_384_imagenet": {
+        "metadata": {
+            "description": (
+                "ViT-L16 model pre-trained on the ImageNet 1k dataset with "
+                "image resolution of 384x384 "
+            ),
+            "params": 303690752,
+            "path": "vit",
+        },
+        "kaggle_handle": "kaggle://keras/vit/keras/vit_large_patch16_384_imagenet/2",
+    },
+    "vit_base_patch32_384_imagenet": {
+        "metadata": {
+            "description": (
+                "ViT-B32 model pre-trained on the ImageNet 1k dataset with "
+                "image resolution of 384x384 "
+            ),
+            "params": 87528192,
+            "path": "vit",
+        },
+        "kaggle_handle": "kaggle://keras/vit/keras/vit_base_patch32_384_imagenet/1",
+    },
+    "vit_large_patch32_384_imagenet": {
+        "metadata": {
+            "description": (
+                "ViT-L32 model pre-trained on the ImageNet 1k dataset with "
+                "image resolution of 384x384 "
+            ),
+            "params": 305607680,
+            "path": "vit",
+        },
+        "kaggle_handle": "kaggle://keras/vit/keras/vit_large_patch32_384_imagenet/1",
+    },
+    "vit_base_patch16_224_imagenet21k": {
+        "metadata": {
+            "description": (
+                "ViT-B16 backbone pre-trained on the ImageNet 21k dataset with "
+                "image resolution of 224x224 "
+            ),
+            "params": 85798656,
+            "path": "vit",
+        },
+        "kaggle_handle": "kaggle://keras/vit/keras/vit_base_patch16_224_imagenet21k/1",
+    },
+    "vit_base_patch32_224_imagenet21k": {
+        "metadata": {
+            "description": (
+                "ViT-B32 backbone pre-trained on the ImageNet 21k dataset with "
+                "image resolution of 224x224 "
+            ),
+            "params": 87455232,
+            "path": "vit",
+        },
+        "kaggle_handle": "kaggle://keras/vit/keras/vit_base_patch32_224_imagenet21k/1",
+    },
+    "vit_huge_patch14_224_imagenet21k": {
+        "metadata": {
+            "description": (
+                "ViT-H14 backbone pre-trained on the ImageNet 21k dataset with "
+                "image resolution of 224x224 "
+            ),
+            "params": 630764800,
+            "path": "vit",
+        },
+        "kaggle_handle": "kaggle://keras/vit/keras/vit_huge_patch14_224_imagenet21k/1",
+    },
+    "vit_large_patch16_224_imagenet21k": {
+        "metadata": {
+            "description": (
+                "ViT-L16 backbone pre-trained on the ImageNet 21k dataset with "
+                "image resolution of 224x224 "
+            ),
+            "params": 303301632,
+            "path": "vit",
+        },
+        "kaggle_handle": "kaggle://keras/vit/keras/vit_large_patch16_224_imagenet21k/1",
+    },
+    "vit_large_patch32_224_imagenet21k": {
+        "metadata": {
+            "description": (
+                "ViT-L32 backbone pre-trained on the ImageNet 21k dataset with "
+                "image resolution of 224x224 "
+            ),
+            "params": 305510400,
+            "path": "vit",
+        },
+        "kaggle_handle": "kaggle://keras/vit/keras/vit_large_patch32_224_imagenet21k/1",
+    },
+}

keras_hub/src/models/vit_det/vit_det_backbone.py

@@ -31,7 +31,7 @@ class ViTDetBackbone(Backbone):
         global_attention_layer_indices (list): Indexes for blocks using
             global attention.
         image_shape (tuple[int], optional): The size of the input image in
-            `(H, W, C)` format. Defaults to `(1024, 1024, 3)`.
+            `(H, W, C)` format. Defaults to `(None, None, 3)`.
         patch_size (int, optional): the patch size to be supplied to the
             Patching layer to turn input images into a flattened sequence of
             patches. Defaults to `16`.
@@ -79,7 +79,7 @@ class ViTDetBackbone(Backbone):
         intermediate_dim,
         num_heads,
         global_attention_layer_indices,
-        image_shape=(1024, 1024, 3),
+        image_shape=(None, None, 3),
         patch_size=16,
         num_output_channels=256,
         use_bias=True,
@@ -87,7 +87,7 @@ class ViTDetBackbone(Backbone):
         use_rel_pos=True,
         window_size=14,
         layer_norm_epsilon=1e-6,
-        **kwargs
+        **kwargs,
     ):
         # === Functional model ===
         img_input = keras.layers.Input(shape=image_shape, name="images")
@@ -179,7 +179,9 @@ class ViTDetBackbone(Backbone):
                 "use_abs_pos": self.use_abs_pos,
                 "use_rel_pos": self.use_rel_pos,
                 "window_size": self.window_size,
-                "global_attention_layer_indices": self.global_attention_layer_indices,
+                "global_attention_layer_indices": (
+                    self.global_attention_layer_indices
+                ),
                 "layer_norm_epsilon": self.layer_norm_epsilon,
             }
         )

keras_hub/src/models/vit_det/vit_layers.py

@@ -117,7 +117,7 @@ class AddRelativePositionalEmbedding(keras.layers.Layer):
         """Calculate decomposed Relative Positional Embeddings
 
         The code has been adapted based on
-        https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py  # noqa: E501
+        https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py
 
         Args:
             attention_map (tensor): Attention map.
@@ -193,7 +193,7 @@ class MultiHeadAttentionWithRelativePE(keras.layers.Layer):
         use_bias=True,
         use_rel_pos=False,
         input_size=None,
-        **kwargs
+        **kwargs,
     ):
         super().__init__(**kwargs)
         self.num_heads = num_heads
@@ -378,7 +378,7 @@ class WindowedTransformerEncoder(keras.layers.Layer):
         input_size=None,
         activation="gelu",
         layer_norm_epsilon=1e-6,
-        **kwargs
+        **kwargs,
     ):
         super().__init__(**kwargs)
         self.project_dim = project_dim

keras_hub/src/models/whisper/whisper_audio_converter.py

@@ -39,7 +39,7 @@ class WhisperAudioConverter(AudioConverter):
     audio_tensor = tf.ones((8000,), dtype="float32")
 
     # Compute the log-mel spectrogram.
-    audio_converter = keras_hub.models.WhisperAudioConverter.from_preset(
+    audio_converter = keras_hub.layers.WhisperAudioConverter.from_preset(
         "whisper_base_en",
     )
     audio_converter(audio_tensor)
@@ -172,9 +172,7 @@ class WhisperAudioConverter(AudioConverter):
         )
 
         def tf_log10(x):
-            """
-            Computes log base 10 of input tensor using TensorFlow's natural log operator.
-            """
+            """Computes log base 10 of input tensor using TensorFlow."""
             numerator = tf.math.log(x)
             denominator = tf.math.log(tf.constant(10, dtype=numerator.dtype))
             return numerator / denominator

keras_hub/src/models/whisper/whisper_backbone.py

@@ -30,9 +30,10 @@ class WhisperBackbone(Backbone):
     It includes the embedding lookups and transformer layers, but not the head
     for predicting the next token.
 
-    The default constructor gives a fully customizable, randomly initialized Whisper
-    model with any number of layers, heads, and embedding dimensions. To load
-    preset architectures and weights, use the `from_preset()` constructor.
+    The default constructor gives a fully customizable, randomly initialized
+    Whisper model with any number of layers, heads, and embedding dimensions.
+    To load preset architectures and weights, use the `from_preset()`
+    constructor.
 
     Disclaimer: Pre-trained models are provided on an "as is" basis, without
     warranties or conditions of any kind. The underlying model is provided by a
@@ -53,8 +54,8 @@ class WhisperBackbone(Backbone):
         max_encoder_sequence_length: int. The maximum sequence length that the
             audio encoder can consume. Since the second convolutional layer in
             the encoder reduces the sequence length by half (stride of 2), we
-            use `max_encoder_sequence_length // 2` as the sequence length for the
-            positional embedding layer.
+            use `max_encoder_sequence_length // 2` as the sequence length for
+            the positional embedding layer.
         max_decoder_sequence_length: int. The maximum sequence length that the
             text decoder can consume.
         dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use

keras_hub/src/models/whisper/whisper_decoder.py

@@ -14,11 +14,9 @@ class WhisperDecoder(TransformerDecoder):
     """Whisper decoder.
 
     Inherits from `keras_hub.layers.TransformerDecoder`, and overrides the
-    `build` method to use the
-    `keras_hub.models.whisper.whisper_multi_head_attention.WhisperMultiHeadAttention`
-    layer instead of `keras.layers.MultiHeadAttention` and
-    `keras_hub.models.whisper.whisper_cached_multi_head_attention.WhisperCachedMultiHeadAttention`
-    instead of `keras_hub.layers.cached_multi_head_attention.CachedMultiHeadAttention`.
+    `build` method to use the `WhisperMultiHeadAttention`
+    layer instead of `MultiHeadAttention` and `WhisperCachedMultiHeadAttention`
+    instead of `CachedMultiHeadAttention`.
     """
 
     def build(

keras_hub/src/models/whisper/whisper_presets.py

@@ -7,11 +7,9 @@ backbone_presets = {
                 "English speech data."
             ),
             "params": 37184256,
-            "official_name": "Whisper",
             "path": "whisper",
-            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_tiny_en/3",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_tiny_en/4",
     },
     "whisper_base_en": {
         "metadata": {
@@ -20,11 +18,9 @@ backbone_presets = {
                 "English speech data."
             ),
             "params": 124439808,
-            "official_name": "Whisper",
             "path": "whisper",
-            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_base_en/3",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_base_en/4",
     },
     "whisper_small_en": {
         "metadata": {
@@ -33,11 +29,9 @@ backbone_presets = {
                 "English speech data."
             ),
             "params": 241734144,
-            "official_name": "Whisper",
             "path": "whisper",
-            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_small_en/3",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_small_en/4",
     },
     "whisper_medium_en": {
         "metadata": {
@@ -46,11 +40,9 @@ backbone_presets = {
                 "English speech data."
             ),
             "params": 763856896,
-            "official_name": "Whisper",
             "path": "whisper",
-            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_medium_en/3",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_medium_en/4",
     },
     "whisper_tiny_multi": {
         "metadata": {
@@ -59,11 +51,9 @@ backbone_presets = {
                 "multilingual speech data."
             ),
             "params": 37760640,
-            "official_name": "Whisper",
             "path": "whisper",
-            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_tiny_multi/3",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_tiny_multi/4",
     },
     "whisper_base_multi": {
         "metadata": {
@@ -72,11 +62,9 @@ backbone_presets = {
                 "multilingual speech data."
             ),
             "params": 72593920,
-            "official_name": "Whisper",
             "path": "whisper",
-            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_base_multi/3",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_base_multi/4",
     },
     "whisper_small_multi": {
         "metadata": {
@@ -85,11 +73,9 @@ backbone_presets = {
                 "multilingual speech data."
             ),
             "params": 241734912,
-            "official_name": "Whisper",
             "path": "whisper",
-            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_small_multi/3",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_small_multi/4",
     },
     "whisper_medium_multi": {
         "metadata": {
@@ -98,11 +84,9 @@ backbone_presets = {
                 "multilingual speech data."
             ),
             "params": 763857920,
-            "official_name": "Whisper",
             "path": "whisper",
-            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_medium_multi/3",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_medium_multi/4",
     },
     "whisper_large_multi": {
         "metadata": {
@@ -111,11 +95,9 @@ backbone_presets = {
                 "multilingual speech data."
             ),
             "params": 1543304960,
-            "official_name": "Whisper",
             "path": "whisper",
-            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_large_multi/3",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_large_multi/4",
     },
     "whisper_large_multi_v2": {
         "metadata": {
@@ -125,10 +107,8 @@ backbone_presets = {
                 "of `whisper_large_multi`."
             ),
             "params": 1543304960,
-            "official_name": "Whisper",
             "path": "whisper",
-            "model_card": "https://github.com/openai/whisper/blob/main/model-card.md",
         },
-        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_large_multi_v2/3",
+        "kaggle_handle": "kaggle://keras/whisper/keras/whisper_large_multi_v2/4",
     },
 }

keras_hub/src/models/xlm_roberta/xlm_roberta_masked_lm.py

@@ -9,7 +9,7 @@ from keras_hub.src.models.roberta.roberta_backbone import (
 from keras_hub.src.models.xlm_roberta.xlm_roberta_backbone import (
     XLMRobertaBackbone,
 )
-from keras_hub.src.models.xlm_roberta.xlm_roberta_masked_lm_preprocessor import (
+from keras_hub.src.models.xlm_roberta.xlm_roberta_masked_lm_preprocessor import (  # noqa: E501
     XLMRobertaMaskedLMPreprocessor,
 )
 

keras_hub/src/models/xlm_roberta/xlm_roberta_masked_lm_preprocessor.py

@@ -20,8 +20,8 @@ class XLMRobertaMaskedLMPreprocessor(MaskedLMPreprocessor):
 
     This preprocessing layer will prepare inputs for a masked language modeling
     task. It is primarily intended for use with the
-    `keras_hub.models.XLMRobertaMaskedLM` task model. Preprocessing will occur in
-    multiple steps.
+    `keras_hub.models.XLMRobertaMaskedLM` task model. Preprocessing will occur
+    in multiple steps.
 
     1. Tokenize any number of input segments using the `tokenizer`.
     2. Pack the inputs together with the appropriate `"<s>"`, `"</s>"` and

keras_hub/src/models/xlm_roberta/xlm_roberta_presets.py

@@ -8,11 +8,9 @@ backbone_presets = {
                 "Trained on CommonCrawl in 100 languages."
             ),
             "params": 277450752,
-            "official_name": "XLM-RoBERTa",
             "path": "xlm_roberta",
-            "model_card": "https://github.com/facebookresearch/fairseq/blob/main/examples/xlmr/README.md",
         },
-        "kaggle_handle": "kaggle://keras/xlm_roberta/keras/xlm_roberta_base_multi/2",
+        "kaggle_handle": "kaggle://keras/xlm_roberta/keras/xlm_roberta_base_multi/3",
     },
     "xlm_roberta_large_multi": {
         "metadata": {
@@ -21,10 +19,8 @@ backbone_presets = {
                 "Trained on CommonCrawl in 100 languages."
             ),
             "params": 558837760,
-            "official_name": "XLM-RoBERTa",
             "path": "xlm_roberta",
-            "model_card": "https://github.com/facebookresearch/fairseq/blob/main/examples/xlmr/README.md",
         },
-        "kaggle_handle": "kaggle://keras/xlm_roberta/keras/xlm_roberta_large_multi/2",
+        "kaggle_handle": "kaggle://keras/xlm_roberta/keras/xlm_roberta_large_multi/3",
     },
 }

keras_hub/src/models/xlm_roberta/xlm_roberta_text_classifier.py

@@ -8,7 +8,7 @@ from keras_hub.src.models.text_classifier import TextClassifier
 from keras_hub.src.models.xlm_roberta.xlm_roberta_backbone import (
     XLMRobertaBackbone,
 )
-from keras_hub.src.models.xlm_roberta.xlm_roberta_text_classifier_preprocessor import (
+from keras_hub.src.models.xlm_roberta.xlm_roberta_text_classifier_preprocessor import (  # noqa: E501
     XLMRobertaTextClassifierPreprocessor,
 )
 
@@ -40,9 +40,9 @@ class XLMRobertaTextClassifier(TextClassifier):
     Args:
         backbone: A `keras_hub.models.XLMRobertaBackbone` instance.
         num_classes: int. Number of classes to predict.
-        preprocessor: A `keras_hub.models.XLMRobertaTextClassifierPreprocessor` or `None`. If
-            `None`, this model will not apply preprocessing, and inputs should
-            be preprocessed before calling the model.
+        preprocessor: A `keras_hub.models.XLMRobertaTextClassifierPreprocessor`
+            or `None`. If `None`, this model will not apply preprocessing, and
+            inputs should be preprocessed before calling the model.
         activation: Optional `str` or callable. The activation function to use
             on the model outputs. Set `activation="softmax"` to return output
             probabilities. Defaults to `None`.

keras_hub/src/models/xlm_roberta/xlm_roberta_tokenizer.py

@@ -177,7 +177,8 @@ class XLMRobertaTokenizer(SentencePieceTokenizer):
         # Shift the tokens IDs left by one.
         tokens = tf.subtract(tokens, 1)
 
-        # Correct `unk_token_id`, `end_token_id`, `start_token_id`, respectively.
+        # Correct `unk_token_id`, `end_token_id`, `start_token_id`,
+        # respectively.
         # Note: The `pad_token_id` is taken as 0 (`unk_token_id`) since the
         # proto does not contain `pad_token_id`. This mapping of the pad token
         # is done automatically by the above subtraction.

keras_hub/src/models/xlnet/relative_attention.py

@@ -64,27 +64,28 @@ def _rel_shift(x, klen=-1):
 class TwoStreamRelativeAttention(keras.layers.MultiHeadAttention):
     """Two-stream relative self-attention for XLNet.
 
-    In XLNet, each token has two associated vectors at each self-attention layer,
-    the content stream (h) and the query stream (g). The content stream is the
-    self-attention stream as in Transformer XL and represents the context and
-    content (the token itself). The query stream only has access to contextual
-    information and the position, but not the content.
+    In XLNet, each token has two associated vectors at each self-attention
+    layer, the content stream (h) and the query stream (g). The content stream
+    is the self-attention stream as in Transformer XL and represents the context
+    and content (the token itself). The query stream only has access to
+    contextual information and the position, but not the content.
 
-    This layer shares the same build signature as `keras.layers.MultiHeadAttention`
-    but has different input/output projections.
+    This layer shares the same build signature as
+    `keras.layers.MultiHeadAttention` but has different input/output
+    projections.
 
     We use the notations `B`, `T`, `S`, `M`, `L`, `E`, `P`, `dim`, `num_heads`
-    below, where
-    `B` is the batch dimension, `T` is the target sequence length,
+    below, where `B` is the batch dimension, `T` is the target sequence length,
     `S` in the source sequence length, `M` is the length of the state or memory,
     `L` is the length of relative positional encoding, `E` is the last dimension
-    of query input, `P` is the number of predictions, `dim` is the dimensionality
-    of the encoder layers. and `num_heads` is the number of attention heads.
+    of query input, `P` is the number of predictions, `dim` is the
+    dimensionality of the encoder layers. and `num_heads` is the number of
+    attention heads.
 
     Args:
         content_stream: `Tensor` of shape `[B, T, dim]`.
-        content_attention_bias: Bias `Tensor` for content based attention of shape
-            `[num_heads, dim]`.
+        content_attention_bias: Bias `Tensor` for content based attention of
+            shape `[num_heads, dim]`.
         positional_attention_bias: Bias `Tensor` for position based attention of
             shape `[num_heads, dim]`.
         query_stream: `Tensor` of shape `[B, P, dim]`.
@@ -96,8 +97,8 @@ class TwoStreamRelativeAttention(keras.layers.MultiHeadAttention):
         segment_encoding: Optional `Tensor` representing the segmentation
             encoding as used in XLNet of shape `[2, num_heads, dim]`.
         segment_attention_bias: Optional trainable bias parameter added to the
-            query had when calculating the segment-based attention score used
-            in XLNet of shape `[num_heads, dim]`.
+            query had when calculating the segment-based attention score used in
+            XLNet of shape `[num_heads, dim]`.
         state: Optional `Tensor` of shape `[B, M, E]`.
             If passed, this is also attended over as in Transformer XL.
         content_attention_mask: a boolean mask of shape `[B, T, S]` that
@@ -336,11 +337,11 @@ class TwoStreamRelativeAttention(keras.layers.MultiHeadAttention):
         dimension of query input.
 
         Args:
-            content_stream: The content representation, commonly referred to as h.
-                This serves a similar role to the standard hidden states in
+            content_stream: The content representation, commonly referred to as
+                h. This serves a similar role to the standard hidden states in
                 Transformer-XL.
-            content_attention_bias: A trainable bias parameter added to the query
-                head when calculating the content-based attention score.
+            content_attention_bias: A trainable bias parameter added to the
+                query head when calculating the content-based attention score.
             positional_attention_bias: A trainable bias parameter added to the
                 query head when calculating the position-based attention score.
             query_stream: The query representation, commonly referred to as g.

keras_hub/src/models/xlnet/xlnet_backbone.py

@@ -49,8 +49,8 @@ class XLNetBackbone(Backbone):
             `[batch_size, sequence_length]`.
         segment_ids: Segment token indices to indicate first and second portions
             of the inputs of shape `[batch_size, sequence_length]`.
-        padding_mask: Mask to avoid performing attention on padding token indices
-            of shape `[batch_size, sequence_length]`.
+        padding_mask: Mask to avoid performing attention on padding token
+            indices of shape `[batch_size, sequence_length]`.
 
     Example:
     ```python

keras_hub/src/models/xlnet/xlnet_content_and_query_embedding.py

@@ -3,8 +3,7 @@ from keras import ops
 
 
 class ContentAndQueryEmbedding(keras.layers.Layer):
-    """
-    Content and Query Embedding.
+    """Content and Query Embedding.
 
     This class creates Content and Query Embeddings for XLNet model
     which is later used in XLNet Encoder.
@@ -20,9 +19,8 @@ class ContentAndQueryEmbedding(keras.layers.Layer):
         **kwargs: other keyword arguments.
 
     References:
-     - [XLNet: Generalized Autoregressive Pretraining for Language Understanding]
-     (https://arxiv.org/abs/1906.08237)
-    """
+     - [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://arxiv.org/abs/1906.08237)
+    """  # noqa: E501
 
     def __init__(
         self, vocabulary_size, hidden_dim, dropout, name=None, **kwargs

keras_hub/src/models/xlnet/xlnet_encoder.py

@@ -11,17 +11,16 @@ def xlnet_kernel_initializer(stddev=0.02):
 
 
 class XLNetEncoder(keras.layers.Layer):
-    """
-    XLNet Encoder.
+    """XLNet Encoder.
 
     This class follows the architecture of the transformer encoder layer in the
     paper [Attention is All You Need](https://arxiv.org/abs/1706.03762). Users
     can instantiate multiple instances of this class to stack up an encoder.
 
     Contrary to the single hidden state used in the paper mentioned above, this
-    Encoder uses two hidden states, Content State and Query State. Thus calculates
-    Two Stream Relative Attention using both of the hidden states. To know more
-    please check the reference.
+    Encoder uses two hidden states, Content State and Query State. Thus
+    calculates Two Stream Relative Attention using both of the hidden states.
+    To know more please check the reference.
 
     Args:
         num_heads: int, the number of heads in the
@@ -44,9 +43,8 @@ class XLNetEncoder(keras.layers.Layer):
         **kwargs: other keyword arguments.
 
     References:
-     - [XLNet: Generalized Autoregressive Pretraining for Language Understanding]
-     (https://arxiv.org/abs/1906.08237)
-    """
+     - [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://arxiv.org/abs/1906.08237)
+    """  # noqa: E501
 
     def __init__(
         self,
@@ -60,7 +58,7 @@ class XLNetEncoder(keras.layers.Layer):
         kernel_initializer_range=0.02,
         bias_initializer="zeros",
         name=None,
-        **kwargs
+        **kwargs,
     ):
         super().__init__(name=name, **kwargs)
         self.num_heads = num_heads

keras_hub/src/samplers/contrastive_sampler.py

@@ -150,9 +150,8 @@ class ContrastiveSampler(Sampler):
             # The final score of each candidate token is weighted sum of
             # probability and similarity against previous tokens.
             accumulated_scores = (
-                (1 - self.alpha) * next_token_probabilities
-                - self.alpha * max_similarity_scores
-            )
+                1 - self.alpha
+            ) * next_token_probabilities - self.alpha * max_similarity_scores
             # Unflatten variables to shape [batch_size, self.k, ...] for
             # gather purpose.
             unflat_score = unflatten_beams(accumulated_scores)

keras_hub/src/samplers/sampler.py

@@ -95,7 +95,8 @@ class Sampler:
         def cond(prompt, cache, index):
             if stop_token_ids is None:
                 return True
-            # Stop if all sequences have produced a *new* id from stop_token_ids.
+            # Stop if all sequences have produced a *new* id from
+            # stop_token_ids.
             end_tokens = any_equal(prompt, stop_token_ids, ~mask)
             prompt_done = ops.any(end_tokens, axis=-1)
             return ops.logical_not(ops.all(prompt_done))