



from typing import Optional, Union

import torch
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithPoolingAndNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_mobilenet_va import MobileNetVaConfig


a__   :					List[Any] 						=      logging.get_logger(__name__)


# General docstring
a__   :					List[Any] 						=      "MobileNetV1Config"

# Base docstring
a__   :					Tuple 						=      "google/mobilenet_v1_1.0_224"
a__   :					Tuple 						=      [1, 10_24, 7, 7]

# Image classification docstring
a__   :					Tuple 						=      "google/mobilenet_v1_1.0_224"
a__   :					str 						=      "tabby, tabby cat"


a__   :					Dict 						=      [
    "google/mobilenet_v1_1.0_224",
    "google/mobilenet_v1_0.75_192",
    # See all MobileNetV1 models at https://huggingface.co/models?filter=mobilenet_v1
]

def 	_lowerCAmelCase (    A__    ,	A__    ,	A__=None						):
      lowercase__  =	{}

      if isinstance(A__    ,	A__						):
            lowercase__  =	model.mobilenet_va
      else:
            lowercase__  =	model

      lowercase__  =	'MobilenetV1/Conv2d_0/'
      lowercase__  =	backbone.conv_stem.convolution.weight
      lowercase__  =	backbone.conv_stem.normalization.bias
      lowercase__  =	backbone.conv_stem.normalization.weight
      lowercase__  =	backbone.conv_stem.normalization.running_mean
      lowercase__  =	backbone.conv_stem.normalization.running_var

      for i in range(13						):
            lowercase__  =	i + 1
            lowercase__  =	i * 2

            lowercase__  =	backbone.layer[pt_index]
            lowercase__  =	F'''MobilenetV1/Conv2d_{tf_index}_depthwise/'''
            lowercase__  =	pointer.convolution.weight
            lowercase__  =	pointer.normalization.bias
            lowercase__  =	pointer.normalization.weight
            lowercase__  =	pointer.normalization.running_mean
            lowercase__  =	pointer.normalization.running_var

            lowercase__  =	backbone.layer[pt_index + 1]
            lowercase__  =	F'''MobilenetV1/Conv2d_{tf_index}_pointwise/'''
            lowercase__  =	pointer.convolution.weight
            lowercase__  =	pointer.normalization.bias
            lowercase__  =	pointer.normalization.weight
            lowercase__  =	pointer.normalization.running_mean
            lowercase__  =	pointer.normalization.running_var

      if isinstance(A__    ,	A__						):
            lowercase__  =	'MobilenetV1/Logits/Conv2d_1c_1x1/'
            lowercase__  =	model.classifier.weight
            lowercase__  =	model.classifier.bias

      return tf_to_pt_map

def 	_lowerCAmelCase (    A__    ,	A__    ,	A__						):
      try:
            import numpy as np
            import tensorflow as tf
      except ImportError:
            logger.error(
                'Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see '
                'https://www.tensorflow.org/install/ for installation instructions.'						)
            raise

      # Load weights from TF model
      lowercase__  =	tf.train.list_variables(A__						)
      lowercase__  =	{}
      for name, shape in init_vars:
            logger.info(F'''Loading TF weight {name} with shape {shape}'''						)
            lowercase__  =	tf.train.load_variable(A__    ,	A__						)
            lowercase__  =	array

      # Build TF to PyTorch weights loading map
      lowercase__  =	_build_tf_to_pytorch_map(A__    ,	A__    ,	A__						)

      for name, pointer in tf_to_pt_map.items():
            logger.info(F'''Importing {name}'''						)
            if name not in tf_weights:
                  logger.info(F'''{name} not in tf pre-trained weights, skipping'''						)
                  continue

            lowercase__  =	tf_weights[name]

            if "depthwise_weights" in name:
                  logger.info('Transposing depthwise'						)
                  lowercase__  =	np.transpose(A__    ,	(2, 3, 0, 1)						)
            elif "weights" in name:
                  logger.info('Transposing'						)
                  if len(pointer.shape						) == 2:  # copying into linear layer
                        lowercase__  =	array.squeeze().transpose()
                  else:
                        lowercase__  =	np.transpose(A__    ,	(3, 2, 0, 1)						)

            if pointer.shape != array.shape:
                  raise ValueError(F'''Pointer shape {pointer.shape} and array shape {array.shape} mismatched'''						)

            logger.info(F'''Initialize PyTorch weight {name} {array.shape}'''						)
            lowercase__  =	torch.from_numpy(A__						)

            tf_weights.pop(A__    ,	A__						)
            tf_weights.pop(name + '/RMSProp'    ,	A__						)
            tf_weights.pop(name + '/RMSProp_1'    ,	A__						)
            tf_weights.pop(name + '/ExponentialMovingAverage'    ,	A__						)

      logger.info(F'''Weights not copied to PyTorch model: {', '.join(tf_weights.keys()						)}'''						)
      return model

def 	_lowerCAmelCase (    A__    ,	A__						):
      lowercase__,						lowercase__  =	features.shape[-2:]
      lowercase__,						lowercase__  =	conv_layer.stride
      lowercase__,						lowercase__  =	conv_layer.kernel_size

      if in_height % stride_height == 0:
            lowercase__  =	max(kernel_height - stride_height    ,	0						)
      else:
            lowercase__  =	max(kernel_height - (in_height % stride_height)    ,	0						)

      if in_width % stride_width == 0:
            lowercase__  =	max(kernel_width - stride_width    ,	0						)
      else:
            lowercase__  =	max(kernel_width - (in_width % stride_width)    ,	0						)

      lowercase__  =	pad_along_width // 2
      lowercase__  =	pad_along_width - pad_left
      lowercase__  =	pad_along_height // 2
      lowercase__  =	pad_along_height - pad_top

      lowercase__  =	(pad_left, pad_right, pad_top, pad_bottom)
      return nn.functional.pad(A__    ,	A__    ,	'constant'    ,	0.0						)



class    UpperCAmelCase__(				nn.Module						):

   '''simple docstring'''




   def __init__(    self	:      Optional[int]      ,					lowerCAmelCase	:      MobileNetVaConfig      ,					lowerCAmelCase	:      int      ,					lowerCAmelCase	:      int      ,					lowerCAmelCase	:      int      ,					lowerCAmelCase	:      Optional[int] = 1      ,					lowerCAmelCase	:      Optional[int] = 1      ,					lowerCAmelCase	:      bool = False      ,					lowerCAmelCase	:      Optional[bool] = True      ,					lowerCAmelCase	:      Optional[bool or str] = True      ,					)						->					None:




         """simple docstring"""

         super().__init__()
         lowercase__  =	config

         if in_channels % groups != 0:
               raise ValueError(f'''Input channels ({in_channels}) are not divisible by {groups} groups.''')
         if out_channels % groups != 0:
               raise ValueError(f'''Output channels ({out_channels}) are not divisible by {groups} groups.''')

         lowercase__  =	0 if config.tf_padding else int((kernel_size - 1) / 2)

         lowercase__  =	nn.Convad(
             in_channels=lowerCAmelCase      ,					out_channels=lowerCAmelCase      ,					kernel_size=lowerCAmelCase      ,					stride=lowerCAmelCase      ,					padding=lowerCAmelCase      ,					groups=lowerCAmelCase      ,					bias=lowerCAmelCase      ,					padding_mode='zeros'      ,					)

         if use_normalization:
               lowercase__  =	nn.BatchNormad(
                   num_features=lowerCAmelCase      ,					eps=config.layer_norm_eps      ,					momentum=0.99_97      ,					affine=lowerCAmelCase      ,					track_running_stats=lowerCAmelCase      ,					)
         else:
               lowercase__  =	None

         if use_activation:
               if isinstance(lowerCAmelCase      ,					lowerCAmelCase):
                     lowercase__  =	ACTaFN[use_activation]
               elif isinstance(config.hidden_act      ,					lowerCAmelCase):
                     lowercase__  =	ACTaFN[config.hidden_act]
               else:
                     lowercase__  =	config.hidden_act
         else:
               lowercase__  =	None



   def       UpperCAmelCase		(    self	:      List[Any]      ,					lowerCAmelCase	:      torch.Tensor)						->					torch.Tensor:




         """simple docstring"""

         if self.config.tf_padding:
               lowercase__  =	apply_tf_padding(lowerCAmelCase      ,					self.convolution)
         lowercase__  =	self.convolution(lowerCAmelCase)
         if self.normalization is not None:
               lowercase__  =	self.normalization(lowerCAmelCase)
         if self.activation is not None:
               lowercase__  =	self.activation(lowerCAmelCase)
         return features



class    UpperCAmelCase__(				lowerCamelCase						):

   '''simple docstring'''




   A     :	Tuple												=     MobileNetVaConfig
   A     :	List[Any]												=     load_tf_weights_in_mobilenet_va
   A     :	List[Any]												=     "mobilenet_v1"
   A     :	Any												=     "pixel_values"
   A     :	Union[str, Any]												=     False


   def       UpperCAmelCase		(    self	:      List[Any]      ,					lowerCAmelCase	:      Union[nn.Linear, nn.Convad])						->					None:




         """simple docstring"""

         if isinstance(lowerCAmelCase      ,					(nn.Linear, nn.Convad)):
               module.weight.data.normal_(mean=0.0      ,					std=self.config.initializer_range)
               if module.bias is not None:
                     module.bias.data.zero_()
         elif isinstance(lowerCAmelCase      ,					nn.BatchNormad):
               module.bias.data.zero_()
               module.weight.data.fill_(1.0)
a__   :					List[Any] 						=      R"\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`MobileNetV1Config`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

a__   :					List[str] 						=      R"\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`MobileNetV1ImageProcessor.__call__`] for details.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n"



@add_start_docstrings(
    "The bare MobileNetV1 model outputting raw hidden-states without any specific head on top."  ,				lowerCamelCase  ,				)
class    UpperCAmelCase__(				lowerCamelCase						):

   '''simple docstring'''




   def __init__(    self	:      List[str]      ,					lowerCAmelCase	:      MobileNetVaConfig      ,					lowerCAmelCase	:      bool = True)						->					Tuple:




         """simple docstring"""

         super().__init__(lowerCAmelCase)
         lowercase__  =	config

         lowercase__  =	32
         lowercase__  =	max(int(depth * config.depth_multiplier)      ,					config.min_depth)

         lowercase__  =	MobileNetVaConvLayer(
             lowerCAmelCase      ,					in_channels=config.num_channels      ,					out_channels=lowerCAmelCase      ,					kernel_size=3      ,					stride=2      ,					)

         lowercase__  =	[1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

         lowercase__  =	nn.ModuleList()
         for i in range(13):
               lowercase__  =	out_channels

               if strides[i] == 2 or i == 0:
                     depth *= 2
                     lowercase__  =	max(int(depth * config.depth_multiplier)      ,					config.min_depth)

               self.layer.append(
                   MobileNetVaConvLayer(
                       lowerCAmelCase      ,					in_channels=lowerCAmelCase      ,					out_channels=lowerCAmelCase      ,					kernel_size=3      ,					stride=strides[i]      ,					groups=lowerCAmelCase      ,					))

               self.layer.append(
                   MobileNetVaConvLayer(
                       lowerCAmelCase      ,					in_channels=lowerCAmelCase      ,					out_channels=lowerCAmelCase      ,					kernel_size=1      ,					))

         lowercase__  =	nn.AdaptiveAvgPoolad((1, 1)) if add_pooling_layer else None

         # Initialize weights and apply final processing
         self.post_init()


   def       UpperCAmelCase		(    self	:      Dict      ,					lowerCAmelCase	:      Dict)						->					str:




         """simple docstring"""

         raise NotImplementedError



   @add_start_docstrings_to_model_forward(lowerCAmelCase)
   @add_code_sample_docstrings(
       checkpoint=_CHECKPOINT_FOR_DOC      ,					output_type=lowerCAmelCase      ,					config_class=_CONFIG_FOR_DOC      ,					modality='vision'      ,					expected_output=_EXPECTED_OUTPUT_SHAPE      ,					)
   def       UpperCAmelCase		(    self	:      Optional[int]      ,					lowerCAmelCase	:      Optional[torch.Tensor] = None      ,					lowerCAmelCase	:      Optional[bool] = None      ,					lowerCAmelCase	:      Optional[bool] = None      ,					)						->					Union[tuple, BaseModelOutputWithPoolingAndNoAttention]:




         """simple docstring"""

         lowercase__  =	(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         lowercase__  =	return_dict if return_dict is not None else self.config.use_return_dict

         if pixel_values is None:
               raise ValueError('You have to specify pixel_values')

         lowercase__  =	self.conv_stem(lowerCAmelCase)

         lowercase__  =	() if output_hidden_states else None

         for i, layer_module in enumerate(self.layer):
               lowercase__  =	layer_module(lowerCAmelCase)

               if output_hidden_states:
                     lowercase__  =	all_hidden_states + (hidden_states,)

         lowercase__  =	hidden_states

         if self.pooler is not None:
               lowercase__  =	torch.flatten(self.pooler(lowerCAmelCase)      ,					start_dim=1)
         else:
               lowercase__  =	None

         if not return_dict:
               return tuple(v for v in [last_hidden_state, pooled_output, all_hidden_states] if v is not None)

         return BaseModelOutputWithPoolingAndNoAttention(
             last_hidden_state=lowerCAmelCase      ,					pooler_output=lowerCAmelCase      ,					hidden_states=lowerCAmelCase      ,					)



@add_start_docstrings(
    "\n    MobileNetV1 model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    "  ,				lowerCamelCase  ,				)
class    UpperCAmelCase__(				lowerCamelCase						):

   '''simple docstring'''




   def __init__(    self	:      Tuple      ,					lowerCAmelCase	:      MobileNetVaConfig)						->					None:




         """simple docstring"""

         super().__init__(lowerCAmelCase)

         lowercase__  =	config.num_labels
         lowercase__  =	MobileNetVaModel(lowerCAmelCase)

         lowercase__  =	self.mobilenet_va.layer[-1].convolution.out_channels

         # Classifier head
         lowercase__  =	nn.Dropout(config.classifier_dropout_prob      ,					inplace=lowerCAmelCase)
         lowercase__  =	nn.Linear(lowerCAmelCase      ,					config.num_labels) if config.num_labels > 0 else nn.Identity()

         # Initialize weights and apply final processing
         self.post_init()



   @add_start_docstrings_to_model_forward(lowerCAmelCase)
   @add_code_sample_docstrings(
       checkpoint=_IMAGE_CLASS_CHECKPOINT      ,					output_type=lowerCAmelCase      ,					config_class=_CONFIG_FOR_DOC      ,					expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT      ,					)
   def       UpperCAmelCase		(    self	:      int      ,					lowerCAmelCase	:      Optional[torch.Tensor] = None      ,					lowerCAmelCase	:      Optional[bool] = None      ,					lowerCAmelCase	:      Optional[torch.Tensor] = None      ,					lowerCAmelCase	:      Optional[bool] = None      ,					)						->					Union[tuple, ImageClassifierOutputWithNoAttention]:




         """simple docstring"""

         lowercase__  =	return_dict if return_dict is not None else self.config.use_return_dict

         lowercase__  =	self.mobilenet_va(lowerCAmelCase      ,					output_hidden_states=lowerCAmelCase      ,					return_dict=lowerCAmelCase)

         lowercase__  =	outputs.pooler_output if return_dict else outputs[1]

         lowercase__  =	self.classifier(self.dropout(lowerCAmelCase))

         lowercase__  =	None
         if labels is not None:
               if self.config.problem_type is None:
                     if self.num_labels == 1:
                           lowercase__  =	'regression'
                     elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                           lowercase__  =	'single_label_classification'
                     else:
                           lowercase__  =	'multi_label_classification'

               if self.config.problem_type == "regression":
                     lowercase__  =	MSELoss()
                     if self.num_labels == 1:
                           lowercase__  =	loss_fct(logits.squeeze()      ,					labels.squeeze())
                     else:
                           lowercase__  =	loss_fct(lowerCAmelCase      ,					lowerCAmelCase)
               elif self.config.problem_type == "single_label_classification":
                     lowercase__  =	CrossEntropyLoss()
                     lowercase__  =	loss_fct(logits.view(-1      ,					self.num_labels)      ,					labels.view(-1))
               elif self.config.problem_type == "multi_label_classification":
                     lowercase__  =	BCEWithLogitsLoss()
                     lowercase__  =	loss_fct(lowerCAmelCase      ,					lowerCAmelCase)

         if not return_dict:
               lowercase__  =	(logits,) + outputs[2:]
               return ((loss,) + output) if loss is not None else output

         return ImageClassifierOutputWithNoAttention(
             loss=lowerCAmelCase      ,					logits=lowerCAmelCase      ,					hidden_states=outputs.hidden_states      ,					)





from datasets.utils.patching import _PatchedModuleObj, patch_submodule

from . import _test_patching

def 	_lowerCAmelCase (    ):
      import os as original_os
      from os import path as original_path
      from os import rename as original_rename
      from os.path import dirname as original_dirname
      from os.path import join as original_join

      assert _test_patching.os is original_os
      assert _test_patching.path is original_path
      assert _test_patching.join is original_join

      assert _test_patching.renamed_os is original_os
      assert _test_patching.renamed_path is original_path
      assert _test_patching.renamed_join is original_join

      lowercase__  =	'__test_patch_submodule_mock__'
      with patch_submodule(_test_patching    ,	'os.path.join'    ,	A__						):
            # Every way to access os.path.join must be patched, and the rest must stay untouched

            # check os.path.join
            assert isinstance(_test_patching.os    ,	_PatchedModuleObj						)
            assert isinstance(_test_patching.os.path    ,	_PatchedModuleObj						)
            assert _test_patching.os.path.join is mock

            # check path.join
            assert isinstance(_test_patching.path    ,	_PatchedModuleObj						)
            assert _test_patching.path.join is mock

            # check join
            assert _test_patching.join is mock

            # check that the other attributes are untouched
            assert _test_patching.os.rename is original_rename
            assert _test_patching.path.dirname is original_dirname
            assert _test_patching.os.path.dirname is original_dirname

            # Even renamed modules or objects must be patched

            # check renamed_os.path.join
            assert isinstance(_test_patching.renamed_os    ,	_PatchedModuleObj						)
            assert isinstance(_test_patching.renamed_os.path    ,	_PatchedModuleObj						)
            assert _test_patching.renamed_os.path.join is mock

            # check renamed_path.join
            assert isinstance(_test_patching.renamed_path    ,	_PatchedModuleObj						)
            assert _test_patching.renamed_path.join is mock

            # check renamed_join
            assert _test_patching.renamed_join is mock

            # check that the other attributes are untouched
            assert _test_patching.renamed_os.rename is original_rename
            assert _test_patching.renamed_path.dirname is original_dirname
            assert _test_patching.renamed_os.path.dirname is original_dirname

      # check that everthing is back to normal when the patch is over

      assert _test_patching.os is original_os
      assert _test_patching.path is original_path
      assert _test_patching.join is original_join

      assert _test_patching.renamed_os is original_os
      assert _test_patching.renamed_path is original_path
      assert _test_patching.renamed_join is original_join

def 	_lowerCAmelCase (    ):
      assert _test_patching.open is open

      lowercase__  =	'__test_patch_submodule_builtin_mock__'
      # _test_patching has "open" in its globals
      assert _test_patching.open is open
      with patch_submodule(_test_patching    ,	'open'    ,	A__						):
            assert _test_patching.open is mock

      # check that everthing is back to normal when the patch is over

      assert _test_patching.open is open

def 	_lowerCAmelCase (    ):
      # pandas.read_csv is not present in _test_patching
      lowercase__  =	'__test_patch_submodule_missing_mock__'
      with patch_submodule(_test_patching    ,	'pandas.read_csv'    ,	A__						):
            pass

def 	_lowerCAmelCase (    ):
      # builtin should always be mocked even if they're not in the globals
      # in case they're loaded at one point
      lowercase__  =	'__test_patch_submodule_missing_builtin_mock__'
      # _test_patching doesn't have "len" in its globals
      assert getattr(_test_patching    ,	'len'    ,	A__						) is None
      with patch_submodule(_test_patching    ,	'len'    ,	A__						):
            assert _test_patching.len is mock
      assert _test_patching.len is len

def 	_lowerCAmelCase (    ):
      lowercase__  =	'__test_patch_submodule_start_and_stop_mock__'
      lowercase__  =	patch_submodule(_test_patching    ,	'open'    ,	A__						)
      assert _test_patching.open is open
      patch.start()
      assert _test_patching.open is mock
      patch.stop()
      assert _test_patching.open is open

def 	_lowerCAmelCase (    ):
      from os import rename as original_rename
      from os.path import dirname as original_dirname
      from os.path import join as original_join

      lowercase__  =	'__test_patch_submodule_successive_join__'
      lowercase__  =	'__test_patch_submodule_successive_dirname__'
      lowercase__  =	'__test_patch_submodule_successive_rename__'
      assert _test_patching.os.path.join is original_join
      assert _test_patching.os.path.dirname is original_dirname
      assert _test_patching.os.rename is original_rename

      with patch_submodule(_test_patching    ,	'os.path.join'    ,	A__						):
            with patch_submodule(_test_patching    ,	'os.rename'    ,	A__						):
                  with patch_submodule(_test_patching    ,	'os.path.dirname'    ,	A__						):
                        assert _test_patching.os.path.join is mock_join
                        assert _test_patching.os.path.dirname is mock_dirname
                        assert _test_patching.os.rename is mock_rename

    # try another order
      with patch_submodule(_test_patching    ,	'os.rename'    ,	A__						):
            with patch_submodule(_test_patching    ,	'os.path.join'    ,	A__						):
                  with patch_submodule(_test_patching    ,	'os.path.dirname'    ,	A__						):
                        assert _test_patching.os.path.join is mock_join
                        assert _test_patching.os.path.dirname is mock_dirname
                        assert _test_patching.os.rename is mock_rename

      assert _test_patching.os.path.join is original_join
      assert _test_patching.os.path.dirname is original_dirname
      assert _test_patching.os.rename is original_rename

def 	_lowerCAmelCase (    ):
      lowercase__  =	'__test_patch_submodule_doesnt_exist_mock__'
      with patch_submodule(_test_patching    ,	'__module_that_doesn_exist__.__attribute_that_doesn_exist__'    ,	A__						):
            pass
      with patch_submodule(_test_patching    ,	'os.__attribute_that_doesn_exist__'    ,	A__						):
            pass
