
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


lowercase   =	logging.get_logger(__name__)


# General docstring
lowercase   =	'''MobileNetV1Config'''

# Base docstring
lowercase   =	'''google/mobilenet_v1_1.0_224'''
lowercase   =	[1, 1_0_2_4, 7, 7]

# Image classification docstring
lowercase   =	'''google/mobilenet_v1_1.0_224'''
lowercase   =	'''tabby, tabby cat'''


lowercase   =	[
    '''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__					:   Optional[Any]      ,					UpperCAmelCase__					:   List[Any]      ,					UpperCAmelCase__					:   Optional[Any]=None							)	->     str:
							lowerCamelCase_					=       {}

							if isinstance(UpperCAmelCase__      ,					UpperCAmelCase__							):
														lowerCamelCase_					=       model.mobilenet_va
							else:
														lowerCamelCase_					=       model

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

							for i in range(1_3							):
														lowerCamelCase_					=       i + 1
														lowerCamelCase_					=       i * 2

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

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

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

							return tf_to_pt_map



def   __lowerCAmelCase    (				UpperCAmelCase__					:   List[str]      ,					UpperCAmelCase__					:   Optional[Any]      ,					UpperCAmelCase__					:   int							)	->     int:
							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
							lowerCamelCase_					=       tf.train.list_variables(UpperCAmelCase__							)
							lowerCamelCase_					=       {}
							for name, shape in init_vars:
														logger.info(F'''Loading TF weight {name} with shape {shape}'''							)
														lowerCamelCase_					=       tf.train.load_variable(UpperCAmelCase__      ,					UpperCAmelCase__							)
														lowerCamelCase_					=       array

							# Build TF to PyTorch weights loading map
							lowerCamelCase_					=       _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

														lowerCamelCase_					=       tf_weights[name]

														if "depthwise_weights" in name:
																					logger.info("""Transposing depthwise"""							)
																					lowerCamelCase_					=       np.transpose(UpperCAmelCase__      ,					(2, 3, 0, 1)							)
														elif "weights" in name:
																					logger.info("""Transposing"""							)
																					if len(pointer.shape							) == 2:  # copying into linear layer
																												lowerCamelCase_					=       array.squeeze().transpose()
																					else:
																												lowerCamelCase_					=       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}'''							)
														lowerCamelCase_					=       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__					:   torch.Tensor      ,					UpperCAmelCase__					:   nn.Convad							)	->     torch.Tensor:
							lowerCamelCase_	, lowerCamelCase_					=       features.shape[-2:]
							lowerCamelCase_	, lowerCamelCase_					=       conv_layer.stride
							lowerCamelCase_	, lowerCamelCase_					=       conv_layer.kernel_size

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

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

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

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





class    __A(					nn.Module							):

					def __init__(						self       :  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 ,       ):
												super().__init__()
												lowerCamelCase_					=       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.''' )

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

												lowerCamelCase_					=       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:
																			lowerCamelCase_					=       nn.BatchNormad(
																			    num_features=__UpperCamelCase ,       eps=config.layer_norm_eps ,       momentum=0.9997 ,       affine=__UpperCamelCase ,       track_running_stats=__UpperCamelCase ,       )
												else:
																			lowerCamelCase_					=       None

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



					def  lowercase__   (						self       :  List[Any] ,       __UpperCamelCase       :  torch.Tensor ):
												if self.config.tf_padding:
																			lowerCamelCase_					=       apply_tf_padding(__UpperCamelCase ,       self.convolution )
												lowerCamelCase_					=       self.convolution(__UpperCamelCase )
												if self.normalization is not None:
																			lowerCamelCase_					=       self.normalization(__UpperCamelCase )
												if self.activation is not None:
																			lowerCamelCase_					=       self.activation(__UpperCamelCase )
												return features





class    __A(					UpperCAmelCase							):
					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  lowercase__   (						self       :  Union[str, Any] ,       __UpperCamelCase       :  Union[nn.Linear, nn.Convad] ):
												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 )
lowercase   =	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.
'''

lowercase   =	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.'''  ,		UpperCAmelCase  ,		)
class    __A(					UpperCAmelCase							):

					def __init__(						self       :  Any ,       __UpperCamelCase       :  MobileNetVaConfig ,       __UpperCamelCase       :  bool = True ):
												super().__init__(__UpperCamelCase )
												lowerCamelCase_					=       config

												lowerCamelCase_					=       3_2
												lowerCamelCase_					=       max(int(depth * config.depth_multiplier ) ,       config.min_depth )

												lowerCamelCase_					=       MobileNetVaConvLayer(
												    __UpperCamelCase ,       in_channels=config.num_channels ,       out_channels=__UpperCamelCase ,       kernel_size=3 ,       stride=2 ,       )

												lowerCamelCase_					=       [1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

												lowerCamelCase_					=       nn.ModuleList()
												for i in range(1_3 ):
																			lowerCamelCase_					=       out_channels

																			if strides[i] == 2 or i == 0:
																										depth *= 2
																										lowerCamelCase_					=       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 ,       ) )

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

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

					def  lowercase__   (						self       :  int ,       __UpperCamelCase       :  Any ):
												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  lowercase__   (						self       :  List[str] ,       __UpperCamelCase       :  Optional[torch.Tensor] = None ,       __UpperCamelCase       :  Optional[bool] = None ,       __UpperCamelCase       :  Optional[bool] = None ,       ):
												lowerCamelCase_					=       (
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)
												lowerCamelCase_					=       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""" )

												lowerCamelCase_					=       self.conv_stem(__UpperCamelCase )

												lowerCamelCase_					=       () if output_hidden_states else None

												for i, layer_module in enumerate(self.layer ):
																			lowerCamelCase_					=       layer_module(__UpperCamelCase )

																			if output_hidden_states:
																										lowerCamelCase_					=       all_hidden_states + (hidden_states,)

												lowerCamelCase_					=       hidden_states

												if self.pooler is not None:
																			lowerCamelCase_					=       torch.flatten(self.pooler(__UpperCamelCase ) ,       start_dim=1 )
												else:
																			lowerCamelCase_					=       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.
    '''  ,		UpperCAmelCase  ,		)
class    __A(					UpperCAmelCase							):

					def __init__(						self       :  List[Any] ,       __UpperCamelCase       :  MobileNetVaConfig ):
												super().__init__(__UpperCamelCase )

												lowerCamelCase_					=       config.num_labels
												lowerCamelCase_					=       MobileNetVaModel(__UpperCamelCase )

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

												# Classifier head
												lowerCamelCase_					=       nn.Dropout(config.classifier_dropout_prob ,       inplace=__UpperCamelCase )
												lowerCamelCase_					=       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  lowercase__   (						self       :  List[Any] ,       __UpperCamelCase       :  Optional[torch.Tensor] = None ,       __UpperCamelCase       :  Optional[bool] = None ,       __UpperCamelCase       :  Optional[torch.Tensor] = None ,       __UpperCamelCase       :  Optional[bool] = None ,       ):
												lowerCamelCase_					=       return_dict if return_dict is not None else self.config.use_return_dict

												lowerCamelCase_					=       self.mobilenet_va(__UpperCamelCase ,       output_hidden_states=__UpperCamelCase ,       return_dict=__UpperCamelCase )

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

												lowerCamelCase_					=       self.classifier(self.dropout(__UpperCamelCase ) )

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

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

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

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







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available


lowercase   =	{}


try:
		if not is_sentencepiece_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		lowercase   =	['''MLukeTokenizer''']

if TYPE_CHECKING:
		try:
				if not is_sentencepiece_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .tokenization_mluke import MLukeTokenizer


else:
		import sys

		lowercase   =	_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





