

'''simple docstring'''






import argparse
import os

import gluonnlp as nlp
import mxnet as mx
import numpy as np
import torch
from gluonnlp.base import get_home_dir
from gluonnlp.model.bert import BERTEncoder
from gluonnlp.model.utils import _load_vocab
from gluonnlp.vocab import Vocab
from packaging import version
from torch import nn

from transformers import BertConfig, BertForMaskedLM, BertModel, RobertaTokenizer
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.utils import logging


if version.parse(nlp.__version__) != version.parse('0.8.3'):
 raise Exception('requires gluonnlp == 0.8.3')

if version.parse(mx.__version__) != version.parse('1.5.0'):
 raise Exception('requires mxnet == 1.5.0')

logging.set_verbosity_info()
__SCREAMING_SNAKE_CASE     =						logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE     =						'The Nymphenburg Palace is a beautiful palace in Munich!'



def 			__a		(       lowerCAmelCase__ :    List[Any]     ,  lowerCAmelCase__ :    Dict      ):
     a__   :      int     =						{
         '''attention_cell''': '''multi_head''',
         '''num_layers''': 4,
         '''units''': 1024,
         '''hidden_size''': 768,
         '''max_length''': 512,
         '''num_heads''': 8,
         '''scaled''': True,
         '''dropout''': 0.1,
         '''use_residual''': True,
         '''embed_size''': 1024,
         '''embed_dropout''': 0.1,
         '''word_embed''': None,
         '''layer_norm_eps''': 1E-5,
         '''token_type_vocab_size''': 2,
     }

     a__   :      Optional[Any]     =						bort_4_8_768_1024_hparams

     # Let's construct the original Bort model here
     # Taken from official BERT implementation, see:
     # https://github.com/alexa/bort/blob/master/bort/bort.py
     a__   :      List[Any]     =						BERTEncoder(
         attention_cell=predefined_args['''attention_cell''']     ,  num_layers=predefined_args['''num_layers''']     ,  units=predefined_args['''units''']     ,  hidden_size=predefined_args['''hidden_size''']     ,  max_length=predefined_args['''max_length''']     ,  num_heads=predefined_args['''num_heads''']     ,  scaled=predefined_args['''scaled''']     ,  dropout=predefined_args['''dropout''']     ,  output_attention=UpperCAmelCase_     ,  output_all_encodings=UpperCAmelCase_     ,  use_residual=predefined_args['''use_residual''']     ,  activation=predefined_args.get('''activation'''     ,  '''gelu'''      )     ,  layer_norm_eps=predefined_args.get('''layer_norm_eps'''     ,  UpperCAmelCase_      )     ,  )

     # Vocab information needs to be fetched first
     # It's the same as RoBERTa, so RobertaTokenizer can be used later
     a__   :      Dict     =						'''openwebtext_ccnews_stories_books_cased'''

     # Specify download folder to Gluonnlp's vocab
     a__   :      int     =						os.path.join(get_home_dir()     ,  '''models'''      )
     a__   :      int     =						_load_vocab(UpperCAmelCase_     ,  UpperCAmelCase_     ,  UpperCAmelCase_     ,  cls=UpperCAmelCase_      )

     a__   :      Any     =						nlp.model.BERTModel(
         UpperCAmelCase_     ,  len(UpperCAmelCase_      )     ,  units=predefined_args['''units''']     ,  embed_size=predefined_args['''embed_size''']     ,  embed_dropout=predefined_args['''embed_dropout''']     ,  word_embed=predefined_args['''word_embed''']     ,  use_pooler=UpperCAmelCase_     ,  use_token_type_embed=UpperCAmelCase_     ,  token_type_vocab_size=predefined_args['''token_type_vocab_size''']     ,  use_classifier=UpperCAmelCase_     ,  use_decoder=UpperCAmelCase_     ,  )

     original_bort.load_parameters(UpperCAmelCase_     ,  cast_dtype=UpperCAmelCase_     ,  ignore_extra=UpperCAmelCase_      )
     a__   :      Optional[int]     =						original_bort._collect_params_with_prefix()

     # Build our config 🤗
     a__   :      Optional[Any]     =						{
         '''architectures''': ['''BertForMaskedLM'''],
         '''attention_probs_dropout_prob''': predefined_args['''dropout'''],
         '''hidden_act''': '''gelu''',
         '''hidden_dropout_prob''': predefined_args['''dropout'''],
         '''hidden_size''': predefined_args['''embed_size'''],
         '''initializer_range''': 0.02,
         '''intermediate_size''': predefined_args['''hidden_size'''],
         '''layer_norm_eps''': predefined_args['''layer_norm_eps'''],
         '''max_position_embeddings''': predefined_args['''max_length'''],
         '''model_type''': '''bort''',
         '''num_attention_heads''': predefined_args['''num_heads'''],
         '''num_hidden_layers''': predefined_args['''num_layers'''],
         '''pad_token_id''': 1,  # 2 = BERT, 1 = RoBERTa
         '''type_vocab_size''': 1,  # 2 = BERT, 1 = RoBERTa
         '''vocab_size''': len(UpperCAmelCase_      ),
     }

     a__   :      Optional[int]     =						BertConfig.from_dict(UpperCAmelCase_      )
     a__   :      Any     =						BertForMaskedLM(UpperCAmelCase_      )
     hf_bort_model.eval()

     # Parameter mapping table (Gluonnlp to Transformers)
     # * denotes layer index
     #
     # | Gluon Parameter                                                | Transformers Parameter
     # | -------------------------------------------------------------- | ----------------------
     # | `encoder.layer_norm.beta`                                      | `bert.embeddings.LayerNorm.bias`
     # | `encoder.layer_norm.gamma`                                     | `bert.embeddings.LayerNorm.weight`
     # | `encoder.position_weight`                                      | `bert.embeddings.position_embeddings.weight`
     # | `word_embed.0.weight`                                          | `bert.embeddings.word_embeddings.weight`
     # | `encoder.transformer_cells.*.attention_cell.proj_key.bias`     | `bert.encoder.layer.*.attention.self.key.bias`
     # | `encoder.transformer_cells.*.attention_cell.proj_key.weight`   | `bert.encoder.layer.*.attention.self.key.weight`
     # | `encoder.transformer_cells.*.attention_cell.proj_query.bias`   | `bert.encoder.layer.*.attention.self.query.bias`
     # | `encoder.transformer_cells.*.attention_cell.proj_query.weight` | `bert.encoder.layer.*.attention.self.query.weight`
     # | `encoder.transformer_cells.*.attention_cell.proj_value.bias`   | `bert.encoder.layer.*.attention.self.value.bias`
     # | `encoder.transformer_cells.*.attention_cell.proj_value.weight` | `bert.encoder.layer.*.attention.self.value.weight`
     # | `encoder.transformer_cells.*.ffn.ffn_2.bias`                   | `bert.encoder.layer.*.attention.output.dense.bias`
     # | `encoder.transformer_cells.*.ffn.ffn_2.weight`                 | `bert.encoder.layer.*.attention.output.dense.weight`
     # | `encoder.transformer_cells.*.layer_norm.beta`                  | `bert.encoder.layer.*.attention.output.LayerNorm.bias`
     # | `encoder.transformer_cells.*.layer_norm.gamma`                 | `bert.encoder.layer.*.attention.output.LayerNorm.weight`
     # | `encoder.transformer_cells.*.ffn.ffn_1.bias`                   | `bert.encoder.layer.*.intermediate.dense.bias`
     # | `encoder.transformer_cells.*.ffn.ffn_1.weight`                 | `bert.encoder.layer.*.intermediate.dense.weight`
     # | `encoder.transformer_cells.*.ffn.layer_norm.beta`              | `bert.encoder.layer.*.output.LayerNorm.bias`
     # | `encoder.transformer_cells.*.ffn.layer_norm.gamma`             | `bert.encoder.layer.*.output.LayerNorm.weight`
     # | `encoder.transformer_cells.*.proj.bias`                        | `bert.encoder.layer.*.output.dense.bias`
     # | `encoder.transformer_cells.*.proj.weight`                      | `bert.encoder.layer.*.output.dense.weight`

     # Helper function to convert MXNET Arrays to PyTorch
     def to_torch(lowerCAmelCase__ :    Union[str, Any]      ) -> nn.Parameter:
          return nn.Parameter(torch.FloatTensor(mx_array.data().asnumpy()      )      )

     # Check param shapes and map new HF param back
     def check_and_map_params(lowerCAmelCase__ :    Tuple     ,  lowerCAmelCase__ :    Optional[Any]      ):
          a__   :      Tuple     =						hf_param.shape

          a__   :      List[str]     =						to_torch(params[gluon_param]      )
          a__   :      List[Any]     =						gluon_param.shape

          assert (
              shape_hf == shape_gluon
          ), F'The gluon parameter {gluon_param} has shape {shape_gluon}, but expects shape {shape_hf} for Transformers'

          return gluon_param

     a__   :      int     =						check_and_map_params(
         hf_bort_model.bert.embeddings.word_embeddings.weight     ,  '''word_embed.0.weight'''      )
     a__   :      Dict     =						check_and_map_params(
         hf_bort_model.bert.embeddings.position_embeddings.weight     ,  '''encoder.position_weight'''      )
     a__   :      List[Any]     =						check_and_map_params(
         hf_bort_model.bert.embeddings.LayerNorm.bias     ,  '''encoder.layer_norm.beta'''      )
     a__   :      List[Any]     =						check_and_map_params(
         hf_bort_model.bert.embeddings.LayerNorm.weight     ,  '''encoder.layer_norm.gamma'''      )

     # Inspired by RoBERTa conversion script, we just zero them out (Bort does not use them)
     a__   :      Tuple     =						torch.zeros_like(
         hf_bort_model.bert.embeddings.token_type_embeddings.weight.data      )

     for i in range(hf_bort_config.num_hidden_layers      ):
          a__   :      Dict     =						hf_bort_model.bert.encoder.layer[i]

          # self attention
          a__   :      List[Any]     =						layer.attention.self

          a__   :      Any     =						check_and_map_params(
              self_attn.key.bias.data     ,  F'encoder.transformer_cells.{i}.attention_cell.proj_key.bias'      )

          a__   :      Dict     =						check_and_map_params(
              self_attn.key.weight.data     ,  F'encoder.transformer_cells.{i}.attention_cell.proj_key.weight'      )
          a__   :      int     =						check_and_map_params(
              self_attn.query.bias.data     ,  F'encoder.transformer_cells.{i}.attention_cell.proj_query.bias'      )
          a__   :      Optional[Any]     =						check_and_map_params(
              self_attn.query.weight.data     ,  F'encoder.transformer_cells.{i}.attention_cell.proj_query.weight'      )
          a__   :      Tuple     =						check_and_map_params(
              self_attn.value.bias.data     ,  F'encoder.transformer_cells.{i}.attention_cell.proj_value.bias'      )
          a__   :      Dict     =						check_and_map_params(
              self_attn.value.weight.data     ,  F'encoder.transformer_cells.{i}.attention_cell.proj_value.weight'      )

          # self attention output
          a__   :      List[str]     =						layer.attention.output

          a__   :      int     =						check_and_map_params(
              self_output.dense.bias     ,  F'encoder.transformer_cells.{i}.proj.bias'      )
          a__   :      Dict     =						check_and_map_params(
              self_output.dense.weight     ,  F'encoder.transformer_cells.{i}.proj.weight'      )
          a__   :      Optional[Any]     =						check_and_map_params(
              self_output.LayerNorm.bias     ,  F'encoder.transformer_cells.{i}.layer_norm.beta'      )
          a__   :      Any     =						check_and_map_params(
              self_output.LayerNorm.weight     ,  F'encoder.transformer_cells.{i}.layer_norm.gamma'      )

          # intermediate
          a__   :      List[Any]     =						layer.intermediate

          a__   :      Optional[Any]     =						check_and_map_params(
              intermediate.dense.bias     ,  F'encoder.transformer_cells.{i}.ffn.ffn_1.bias'      )
          a__   :      Dict     =						check_and_map_params(
              intermediate.dense.weight     ,  F'encoder.transformer_cells.{i}.ffn.ffn_1.weight'      )

          # output
          a__   :      List[str]     =						layer.output

          a__   :      str     =						check_and_map_params(
              bert_output.dense.bias     ,  F'encoder.transformer_cells.{i}.ffn.ffn_2.bias'      )
          a__   :      Optional[Any]     =						check_and_map_params(
              bert_output.dense.weight     ,  F'encoder.transformer_cells.{i}.ffn.ffn_2.weight'      )
          a__   :      Dict     =						check_and_map_params(
              bert_output.LayerNorm.bias     ,  F'encoder.transformer_cells.{i}.ffn.layer_norm.beta'      )
          a__   :      Any     =						check_and_map_params(
              bert_output.LayerNorm.weight     ,  F'encoder.transformer_cells.{i}.ffn.layer_norm.gamma'      )

     # Save space and energy 🎄
     hf_bort_model.half()

     # Compare output of both models
     a__   :      Optional[int]     =						RobertaTokenizer.from_pretrained('''roberta-base'''      )

     a__   :      Optional[int]     =						tokenizer.encode_plus(UpperCAmelCase_      )['''input_ids''']

     # Get gluon output
     a__   :      str     =						mx.nd.array([input_ids]      )
     a__   :      Optional[Any]     =						original_bort(inputs=UpperCAmelCase_     ,  token_types=[]      )

     # Get Transformer output (save and reload model again)
     hf_bort_model.save_pretrained(UpperCAmelCase_      )
     a__   :      Optional[Any]     =						BertModel.from_pretrained(UpperCAmelCase_      )
     hf_bort_model.eval()

     a__   :      int     =						tokenizer.encode_plus(UpperCAmelCase_     ,  return_tensors='''pt'''      )
     a__   :      Union[str, Any]     =						hf_bort_model(**UpperCAmelCase_      )[0]

     a__   :      List[Any]     =						output_gluon[0].asnumpy()
     a__   :      int     =						output_hf[0].detach().numpy()

     a__   :      List[Any]     =						np.max(np.abs(hf_layer - gluon_layer      )      ).item()
     a__   :      Optional[int]     =						np.allclose(UpperCAmelCase_     ,  UpperCAmelCase_     ,  atol=1E-3      )

     if success:
          print('''✔️ Both model do output the same tensors'''      )
     else:
          print('''❌ Both model do **NOT** output the same tensors'''      )
          print('''Absolute difference is:'''     ,  UpperCAmelCase_      )


if __name__ == "__main__":
 __SCREAMING_SNAKE_CASE     =						argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '--bort_checkpoint_path', default=None, type=str, required=True, help='Path the official Bort params file.'
 )
 parser.add_argument(
     '--pytorch_dump_folder_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
 )
 __SCREAMING_SNAKE_CASE     =						parser.parse_args()
 convert_bort_checkpoint_to_pytorch(args.bort_checkpoint_path, args.pytorch_dump_folder_path)


"""simple docstring"""


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


SCREAMING_SNAKE_CASE			  =  logging.get_logger(__name__)


# General docstring
SCREAMING_SNAKE_CASE			  =  """MobileNetV1Config"""

# Base docstring
SCREAMING_SNAKE_CASE			  =  """google/mobilenet_v1_1.0_224"""
SCREAMING_SNAKE_CASE			  =  [1, 1_024, 7, 7]

# Image classification docstring
SCREAMING_SNAKE_CASE			  =  """google/mobilenet_v1_1.0_224"""
SCREAMING_SNAKE_CASE			  =  """tabby, tabby cat"""


SCREAMING_SNAKE_CASE			  =  [
    """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__      (				UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_=None					)->				int:







       """simple docstring"""
       UpperCamelCase       			=		{}

       if isinstance(UpperCAmelCase_		,   UpperCAmelCase_					):
              UpperCamelCase       			=		model.mobilenet_va
       else:
              UpperCamelCase       			=		model

       UpperCamelCase       			=		"MobilenetV1/Conv2d_0/"
       UpperCamelCase       			=		backbone.conv_stem.convolution.weight
       UpperCamelCase       			=		backbone.conv_stem.normalization.bias
       UpperCamelCase       			=		backbone.conv_stem.normalization.weight
       UpperCamelCase       			=		backbone.conv_stem.normalization.running_mean
       UpperCamelCase       			=		backbone.conv_stem.normalization.running_var

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

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

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

       if isinstance(UpperCAmelCase_		,   UpperCAmelCase_					):
              UpperCamelCase       			=		"MobilenetV1/Logits/Conv2d_1c_1x1/"
              UpperCamelCase       			=		model.classifier.weight
              UpperCamelCase       			=		model.classifier.bias

       return tf_to_pt_map



def 						lowerCamelCase__      (				UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_					)->				List[str]:







       """simple docstring"""
       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
       UpperCamelCase       			=		tf.train.list_variables(UpperCAmelCase_					)
       UpperCamelCase       			=		{}
       for name, shape in init_vars:
              logger.info(F"Loading TF weight {name} with shape {shape}"					)
              UpperCamelCase       			=		tf.train.load_variable(UpperCAmelCase_		,   UpperCAmelCase_					)
              UpperCamelCase       			=		array

       # Build TF to PyTorch weights loading map
       UpperCamelCase       			=		_build_tf_to_pytorch_map(UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_					)

       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

              UpperCamelCase       			=		tf_weights[name]

              if "depthwise_weights" in name:
                     logger.info("Transposing depthwise"					)
                     UpperCamelCase       			=		np.transpose(UpperCAmelCase_		,   (2, 3, 0, 1)					)
              elif "weights" in name:
                     logger.info("Transposing"					)
                     if len(pointer.shape					) == 2:  # copying into linear layer
                            UpperCamelCase       			=		array.squeeze().transpose()
                     else:
                            UpperCamelCase       			=		np.transpose(UpperCAmelCase_		,   (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}"					)
              UpperCamelCase       			=		torch.from_numpy(UpperCAmelCase_					)

              tf_weights.pop(UpperCAmelCase_		,   UpperCAmelCase_					)
              tf_weights.pop(name + "/RMSProp"		,   UpperCAmelCase_					)
              tf_weights.pop(name + "/RMSProp_1"		,   UpperCAmelCase_					)
              tf_weights.pop(name + "/ExponentialMovingAverage"		,   UpperCAmelCase_					)

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



def 						lowerCamelCase__      (				UpperCAmelCase_		,   UpperCAmelCase_					)->				torch.Tensor:







       """simple docstring"""
       UpperCamelCase			,							UpperCamelCase       			=		features.shape[-2:]
       UpperCamelCase			,							UpperCamelCase       			=		conv_layer.stride
       UpperCamelCase			,							UpperCamelCase       			=		conv_layer.kernel_size

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

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

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

       UpperCamelCase       			=		(pad_left, pad_right, pad_top, pad_bottom)
       return nn.functional.pad(UpperCAmelCase_		,   UpperCAmelCase_		,   "constant"		,   0.0					)



class 	__a	(   nn.Module						):






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







              """simple docstring"""
              super().__init__()
              UpperCamelCase       			=		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."      )

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

              UpperCamelCase       			=		nn.Convad(
                  in_channels=UpperCAmelCase_							,   out_channels=UpperCAmelCase_							,   kernel_size=UpperCAmelCase_							,   stride=UpperCAmelCase_							,   padding=UpperCAmelCase_							,   groups=UpperCAmelCase_							,   bias=UpperCAmelCase_							,   padding_mode="zeros"							,   )

              if use_normalization:
                     UpperCamelCase       			=		nn.BatchNormad(
                         num_features=UpperCAmelCase_							,   eps=config.layer_norm_eps							,   momentum=0.9997							,   affine=UpperCAmelCase_							,   track_running_stats=UpperCAmelCase_							,   )
              else:
                     UpperCamelCase       			=		None

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






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Dict							,   UpperCAmelCase_			:				torch.Tensor      )->	torch.Tensor:







              """simple docstring"""
              if self.config.tf_padding:
                     UpperCamelCase       			=		apply_tf_padding(UpperCAmelCase_							,   self.convolution      )
              UpperCamelCase       			=		self.convolution(UpperCAmelCase_      )
              if self.normalization is not None:
                     UpperCamelCase       			=		self.normalization(UpperCAmelCase_      )
              if self.activation is not None:
                     UpperCamelCase       			=		self.activation(UpperCAmelCase_      )
              return features





class 	__a	(   _lowerCAmelCase						):
       UpperCamelCase_		:       List[str]				=      MobileNetVaConfig
       UpperCamelCase_		:       Dict				=      load_tf_weights_in_mobilenet_va
       UpperCamelCase_		:       List[str]				=      '''mobilenet_v1'''
       UpperCamelCase_		:       Optional[int]				=      '''pixel_values'''
       UpperCamelCase_		:       List[Any]				=      False






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Any							,   UpperCAmelCase_			:				Union[nn.Linear, nn.Convad]      )->	None:







              """simple docstring"""
              if isinstance(UpperCAmelCase_							,   (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(UpperCAmelCase_							,   nn.BatchNormad      ):
                     module.bias.data.zero_()
                     module.weight.data.fill_(1.0      )





SCREAMING_SNAKE_CASE			  =  R"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`MobileNetV1Config`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

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



@add_start_docstrings(
    '''The bare MobileNetV1 model outputting raw hidden-states without any specific head on top.'''     ,  _lowerCAmelCase     ,  )
class 	__a	(   _lowerCAmelCase						):






       def __init__(	self			:				str							,   UpperCAmelCase_			:				MobileNetVaConfig							,   UpperCAmelCase_			:				bool = True      )->	Optional[int]:







              """simple docstring"""
              super().__init__(UpperCAmelCase_      )
              UpperCamelCase       			=		config

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

              UpperCamelCase       			=		MobileNetVaConvLayer(
                  UpperCAmelCase_							,   in_channels=config.num_channels							,   out_channels=UpperCAmelCase_							,   kernel_size=3							,   stride=2							,   )

              UpperCamelCase       			=		[1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

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

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

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

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

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

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






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[int]							,   UpperCAmelCase_			:				Dict      )->	List[str]:







              """simple docstring"""
              raise NotImplementedError






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







              """simple docstring"""
              UpperCamelCase       			=		(
                  output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
              )
              UpperCamelCase       			=		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"      )

              UpperCamelCase       			=		self.conv_stem(UpperCAmelCase_      )

              UpperCamelCase       			=		() if output_hidden_states else None

              for i, layer_module in enumerate(self.layer      ):
                     UpperCamelCase       			=		layer_module(UpperCAmelCase_      )

                     if output_hidden_states:
                            UpperCamelCase       			=		all_hidden_states + (hidden_states,)

              UpperCamelCase       			=		hidden_states

              if self.pooler is not None:
                     UpperCamelCase       			=		torch.flatten(self.pooler(UpperCAmelCase_      )							,   start_dim=1      )
              else:
                     UpperCamelCase       			=		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=UpperCAmelCase_							,   pooler_output=UpperCAmelCase_							,   hidden_states=UpperCAmelCase_							,   )





@add_start_docstrings(
    '''
    MobileNetV1 model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    '''     ,  _lowerCAmelCase     ,  )
class 	__a	(   _lowerCAmelCase						):






       def __init__(	self			:				str							,   UpperCAmelCase_			:				MobileNetVaConfig      )->	None:







              """simple docstring"""
              super().__init__(UpperCAmelCase_      )

              UpperCamelCase       			=		config.num_labels
              UpperCamelCase       			=		MobileNetVaModel(UpperCAmelCase_      )

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

              # Classifier head
              UpperCamelCase       			=		nn.Dropout(config.classifier_dropout_prob							,   inplace=UpperCAmelCase_      )
              UpperCamelCase       			=		nn.Linear(UpperCAmelCase_							,   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(UpperCAmelCase_      )
       @add_code_sample_docstrings(
           checkpoint=_IMAGE_CLASS_CHECKPOINT							,   output_type=UpperCAmelCase_							,   config_class=_CONFIG_FOR_DOC							,   expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT							,   )
       def 						_SCREAMING_SNAKE_CASE							(	self			:				List[str]							,   UpperCAmelCase_			:				Optional[torch.Tensor] = None							,   UpperCAmelCase_			:				Optional[bool] = None							,   UpperCAmelCase_			:				Optional[torch.Tensor] = None							,   UpperCAmelCase_			:				Optional[bool] = None							,   )->	Union[tuple, ImageClassifierOutputWithNoAttention]:







              """simple docstring"""
              UpperCamelCase       			=		return_dict if return_dict is not None else self.config.use_return_dict

              UpperCamelCase       			=		self.mobilenet_va(UpperCAmelCase_							,   output_hidden_states=UpperCAmelCase_							,   return_dict=UpperCAmelCase_      )

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

              UpperCamelCase       			=		self.classifier(self.dropout(UpperCAmelCase_      )      )

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

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

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

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





