Resolved failing test cases.

Author: Sohaib-Ahmed21

keras_hub/api/layers/__init__.py

@@ -48,6 +48,7 @@
 from keras_hub.src.models.deeplab_v3.deeplab_v3_image_converter import (
     DeepLabV3ImageConverter,
 )
+from keras_hub.src.models.deit.deit_image_converter import DeiTImageConverter
 from keras_hub.src.models.densenet.densenet_image_converter import (
     DenseNetImageConverter,
 )
@@ -71,9 +72,6 @@
     RetinaNetImageConverter,
 )
 from keras_hub.src.models.sam.sam_image_converter import SAMImageConverter
-from keras_hub.src.models.deit.deit_image_converter import (
-    DeiTImageConverter,
-)
 from keras_hub.src.models.sam.sam_mask_decoder import SAMMaskDecoder
 from keras_hub.src.models.sam.sam_prompt_encoder import SAMPromptEncoder
 from keras_hub.src.models.segformer.segformer_image_converter import (

keras_hub/api/models/__init__.py

@@ -101,20 +101,18 @@
 from keras_hub.src.models.deeplab_v3.deeplab_v3_segmenter import (
     DeepLabV3ImageSegmenter,
 )
+from keras_hub.src.models.deit.deit_backbone import DeiTBackbone
+from keras_hub.src.models.deit.deit_image_classifier import DeiTImageClassifier
+from keras_hub.src.models.deit.deit_image_classifier_preprocessor import (
+    DeiTImageClassifierPreprocessor,
+)
 from keras_hub.src.models.densenet.densenet_backbone import DenseNetBackbone
 from keras_hub.src.models.densenet.densenet_image_classifier import (
     DenseNetImageClassifier,
 )
 from keras_hub.src.models.densenet.densenet_image_classifier_preprocessor import (
     DenseNetImageClassifierPreprocessor,
 )
-from keras_hub.src.models.deit.deit_backbone import(
-    DeiTBackbone,
-)
-from keras_hub.src.models.deit.deit_image_classifier import DeiTImageClassifier
-from keras_hub.src.models.deit.deit_image_classifier_preprocessor import (
-    DeiTImageClassifierPreprocessor,
-)
 from keras_hub.src.models.distil_bert.distil_bert_backbone import (
     DistilBertBackbone,
 )

keras_hub/src/models/deit/deit_backbone.py

@@ -2,38 +2,46 @@
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.backbone import Backbone
-from keras_hub.src.utils.keras_utils import standardize_data_format
-from keras_hub.src.models.deit.deit_layers import DeiTEncoder
 from keras_hub.src.models.deit.deit_layers import DeiTEmbeddings
+from keras_hub.src.models.deit.deit_layers import DeiTEncoder
+from keras_hub.src.utils.keras_utils import standardize_data_format
+
 
 @keras_hub_export("keras_hub.models.DeiTBackbone")
 class DeiTBackbone(Backbone):
     """DeiT backbone.
-    
-    This backbone implements the Data-efficient Image Transformer (DeiT) architecture 
-    as described in [Training data-efficient image transformers & distillation through 
-    attention](https://arxiv.org/abs/2012.12877).
-    
+
+    This backbone implements the Data-efficient Image Transformer (DeiT)
+    architecture as described in [Training data-efficient image
+    transformers & distillation through attention]
+    (https://arxiv.org/abs/2012.12877).
+
     Args:
         image_shape: A tuple or list of 3 integers representing the shape of the
             input image `(height, width, channels)`, `height` and `width` must
             be equal.
-        patch_size: int. The size of each image patch, the input image
-            will be divided into patches of shape `(patch_size_h, patch_size_w)`.
+        patch_size: int. The size of each image patch, the input image will
+            be divided into patches of shape `(patch_size_h, patch_size_w)`.
         num_layers: int. The number of transformer encoder layers.
-        num_heads: int. The number of attention heads in each Transformer encoder layer.
+        num_heads: int. The number of attention heads in each Transformer
+            encoder layer.
         hidden_dim: int. The dimensionality of the hidden representations.
-        intermediate_dim: int. The dimensionality of the intermediate MLP layer in
-            each Transformer encoder layer.
-        dropout_rate: float. The dropout rate for the Transformer encoder layers.
+        intermediate_dim: int. The dimensionality of the intermediate MLP layer
+            in each Transformer encoder layer.
+        dropout_rate: float. The dropout rate for the Transformer encoder
+            layers.
         attention_dropout: float. The dropout rate for the attention mechanism
             in each Transformer encoder layer.
-        layer_norm_epsilon: float. Value used for numerical stability in layer normalization.
-        use_mha_bias: bool. Whether to use bias in the multi-head attention layers.
+        layer_norm_epsilon: float. Value used for numerical stability in layer
+            normalization.
+        use_mha_bias: bool. Whether to use bias in the multi-head attention
+            layers.
         data_format: str. `"channels_last"` or `"channels_first"`, specifying
-            the data format for the input image. If `None`, defaults to `"channels_last"`.
+            the data format for the input image. If `None`, defaults to
+            `"channels_last"`.
         dtype: The dtype of the layer weights. Defaults to None.
-        **kwargs: Additional keyword arguments to be passed to the parent `Backbone` class.
+        **kwargs: Additional keyword arguments to be passed to the parent
+            `Backbone` class.
     """
 
     def __init__(
@@ -108,8 +116,8 @@ def __init__(
             attention_dropout=attention_dropout,
             layer_norm_epsilon=layer_norm_epsilon,
             dtype=dtype,
-            name="deit_encoder"
-          )(x)
+            name="deit_encoder",
+        )(x)
 
         super().__init__(
             inputs=inputs,
@@ -130,7 +138,7 @@ def __init__(
         self.layer_norm_epsilon = layer_norm_epsilon
         self.use_mha_bias = use_mha_bias
         self.data_format = data_format
-    
+
     def get_config(self):
         config = super().get_config()
         config.update(
@@ -147,4 +155,4 @@ def get_config(self):
                 "use_mha_bias": self.use_mha_bias,
             }
         )
-        return config
+        return config

keras_hub/src/models/deit/deit_backbone_test.py

@@ -25,7 +25,7 @@ def test_backbone_basics(self):
             init_kwargs={**self.init_kwargs},
             input_data=self.input_data,
             # 49+2 positions(49 patches, cls and distillation token)
-            expected_output_shape=(2, 51, 48)
+            expected_output_shape=(2, 51, 48),
             run_quantization_check=False,
         )
 

keras_hub/src/models/deit/deit_image_classifier.py

@@ -2,12 +2,12 @@
 from keras import ops
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.models.image_classifier import ImageClassifier
-from keras_hub.src.models.task import Task
 from keras_hub.src.models.deit.deit_backbone import DeiTBackbone
 from keras_hub.src.models.deit.deit_image_classifier_preprocessor import (
     DeiTImageClassifierPreprocessor,
 )
+from keras_hub.src.models.image_classifier import ImageClassifier
+from keras_hub.src.models.task import Task
 
 
 @keras_hub_export("keras_hub.models.DeiTImageClassifier")

keras_hub/src/models/deit/deit_image_classifier_preprocessor.py

@@ -1,9 +1,9 @@
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.deit.deit_backbone import DeiTBackbone
+from keras_hub.src.models.deit.deit_image_converter import DeiTImageConverter
 from keras_hub.src.models.image_classifier_preprocessor import (
     ImageClassifierPreprocessor,
 )
-from keras_hub.src.models.deit.deit_backbone import DeiTBackbone
-from keras_hub.src.models.deit.deit_image_converter import DeiTImageConverter
 
 
 @keras_hub_export("keras_hub.models.ViTImageClassifierPreprocessor")

keras_hub/src/models/deit/deit_image_converter.py

@@ -1,7 +1,7 @@
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.utils.tensor_utils import preprocessing_function
 from keras_hub.src.layers.preprocessing.image_converter import ImageConverter
 from keras_hub.src.models.deit.deit_backbone import DeiTBackbone
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 
 @keras_hub_export("keras_hub.layers.DeiTImageConverter")
@@ -20,7 +20,9 @@ class DeiTImageConverter(ImageConverter):
     ```python
     import keras
     import numpy as np
-    from keras_hub.src.models.deit.deit_image_converter import DeiTImageConverter
+    from keras_hub.src.models.deit.deit_image_converter import (
+        DeiTImageConverter
+    )
     # Example image (replace with your actual image data)
     image = np.random.rand(1, 384, 384, 3)  # Example: (B, H, W, C)
     # Create a DeiTImageConverter instance
@@ -60,4 +62,4 @@ def get_config(self):
                 "norm_std": self.norm_std,
             }
         )
-        return config
+        return config