

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


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


# General docstring
_lowerCAmelCase					:			Union[str, Any]  				=			"MobileNetV1Config"

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

# Image classification docstring
_lowerCAmelCase					:			List[Any]  				=			"google/mobilenet_v1_1.0_224"
_lowerCAmelCase					:			Dict  				=			"tabby, tabby cat"


_lowerCAmelCase					:			List[Any]  				=			[
    "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        UpperCamelCase_(							_snake_case   :  int       ,     _snake_case   :  Optional[Any]       ,     _snake_case   :  Optional[int]=None	):




			"""simple docstring"""






			__a        ={}

			if isinstance(_snake_case       ,     _snake_case	):
						__a        =model.mobilenet_va
			else:
						__a        =model

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

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

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

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

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

			return tf_to_pt_map






def        UpperCamelCase_(							_snake_case   :  Optional[int]       ,     _snake_case   :  List[str]       ,     _snake_case   :  int	):




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

			# Build TF to PyTorch weights loading map
			__a        =_build_tf_to_pytorch_map(_snake_case       ,     _snake_case       ,     _snake_case	)

			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

						__a        =tf_weights[name]

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

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

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






def        UpperCamelCase_(							_snake_case   :  torch.Tensor       ,     _snake_case   :  nn.Convad	):




			"""simple docstring"""






			__a      ,    __a        =features.shape[-2:]
			__a      ,    __a        =conv_layer.stride
			__a      ,    __a        =conv_layer.kernel_size

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

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

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

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






class 	__magic_name__    (			nn.Module      ):





			def __init__(		self					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case = 1					,  __snake_case = 1					,  __snake_case = False					,  __snake_case = True					,  __snake_case = True					,  )   -> None:




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

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

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

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

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





			def 		__magic_name__    (		self					,  __snake_case					)   -> torch.Tensor:




						'''simple docstring'''
						if self.config.tf_padding:
									__a        =apply_tf_padding(__snake_case					,  self.convolution					)
						__a        =self.convolution(__snake_case					)
						if self.normalization is not None:
									__a        =self.normalization(__snake_case					)
						if self.activation is not None:
									__a        =self.activation(__snake_case					)
						return features






class 	__magic_name__    (			lowerCAmelCase_      ):
			SCREAMING_SNAKE_CASE							   =						MobileNetVaConfig
			SCREAMING_SNAKE_CASE							   =						load_tf_weights_in_mobilenet_va
			SCREAMING_SNAKE_CASE							   =						'mobilenet_v1'
			SCREAMING_SNAKE_CASE							   =						'pixel_values'
			SCREAMING_SNAKE_CASE							   =						False





			def 		__magic_name__    (		self					,  __snake_case					)   -> None:




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


_lowerCAmelCase					:			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"

_lowerCAmelCase					:			Optional[int]  				=			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 	__magic_name__    (			lowerCAmelCase_      ):





			def __init__(		self					,  __snake_case					,  __snake_case = True					)   -> Dict:




						'''simple docstring'''
						super().__init__(__snake_case					)
						__a        =config

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

						__a        =MobileNetVaConvLayer(
						    __snake_case					,  in_channels=config.num_channels					,  out_channels=__snake_case					,  kernel_size=3					,  stride=2					,  )

						__a        =[1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

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

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

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

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

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

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





			def 		__magic_name__    (		self					,  __snake_case					)   -> Dict:




						'''simple docstring'''
						raise NotImplementedError





			@add_start_docstrings_to_model_forward(__snake_case					)
			@add_code_sample_docstrings(
			    checkpoint=_CHECKPOINT_FOR_DOC					,  output_type=__snake_case					,  config_class=_CONFIG_FOR_DOC					,  modality='vision'					,  expected_output=_EXPECTED_OUTPUT_SHAPE					,  )
			def 		__magic_name__    (		self					,  __snake_case = None					,  __snake_case = None					,  __snake_case = None					,  )   -> Union[tuple, BaseModelOutputWithPoolingAndNoAttention]:




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

						__a        =self.conv_stem(__snake_case					)

						__a        =() if output_hidden_states else None

						for i, layer_module in enumerate(self.layer					):
									__a        =layer_module(__snake_case					)

									if output_hidden_states:
												__a        =all_hidden_states + (hidden_states,)

						__a        =hidden_states

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






@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 	__magic_name__    (			lowerCAmelCase_      ):





			def __init__(		self					,  __snake_case					)   -> None:




						'''simple docstring'''
						super().__init__(__snake_case					)

						__a        =config.num_labels
						__a        =MobileNetVaModel(__snake_case					)

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

						# Classifier head
						__a        =nn.Dropout(config.classifier_dropout_prob					,  inplace=__snake_case					)
						__a        =nn.Linear(__snake_case					,  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(__snake_case					)
			@add_code_sample_docstrings(
			    checkpoint=_IMAGE_CLASS_CHECKPOINT					,  output_type=__snake_case					,  config_class=_CONFIG_FOR_DOC					,  expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT					,  )
			def 		__magic_name__    (		self					,  __snake_case = None					,  __snake_case = None					,  __snake_case = None					,  __snake_case = None					,  )   -> Union[tuple, ImageClassifierOutputWithNoAttention]:




						'''simple docstring'''
						__a        =return_dict if return_dict is not None else self.config.use_return_dict

						__a        =self.mobilenet_va(__snake_case					,  output_hidden_states=__snake_case					,  return_dict=__snake_case					)

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

						__a        =self.classifier(self.dropout(__snake_case					)					)

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

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

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

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









import argparse

import torch

from transformers import OpenAIGPTConfig, OpenAIGPTModel, load_tf_weights_in_openai_gpt
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()






def        UpperCamelCase_(							_snake_case   :  Dict       ,     _snake_case   :  Optional[int]       ,     _snake_case   :  str	):




			"""simple docstring"""






			if openai_config_file == "":
						__a        =OpenAIGPTConfig()
			else:
						__a        =OpenAIGPTConfig.from_json_file(_snake_case	)
			__a        =OpenAIGPTModel(_snake_case	)

			# Load weights from numpy
			load_tf_weights_in_openai_gpt(_snake_case       ,     _snake_case       ,     _snake_case	)

			# Save pytorch-model
			__a        =pytorch_dump_folder_path + '/' + WEIGHTS_NAME
			__a        =pytorch_dump_folder_path + '/' + CONFIG_NAME
			print(F'Save PyTorch model to {pytorch_weights_dump_path}'	)
			torch.save(model.state_dict()       ,     _snake_case	)
			print(F'Save configuration file to {pytorch_config_dump_path}'	)
			with open(_snake_case       ,     'w'       ,     encoding='utf-8'	) as f:
						f.write(config.to_json_string()	)


if __name__ == "__main__":
						_lowerCAmelCase					:			Tuple  				=			argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    "--openai_checkpoint_folder_path",
						    default=None,
						    type=str,
						    required=True,
						    help="Path to the TensorFlow checkpoint path.",
						)
						parser.add_argument(
						    "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
						)
						parser.add_argument(
						    "--openai_config_file",
						    default="",
						    type=str,
						    help=(
						        "An optional config json file corresponding to the pre-trained OpenAI model. \n"
						        "This specifies the model architecture."
						    ),
						)
						_lowerCAmelCase					:			int  				=			parser.parse_args()
						convert_openai_checkpoint_to_pytorch(
						    args.openai_checkpoint_folder_path, args.openai_config_file, args.pytorch_dump_folder_path
						)






