

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			:					Dict						=			logging.get_logger(__name__)


# General docstring
__A			:					List[Any]						=			'MobileNetV1Config'

# Base docstring
__A			:					List[str]						=			'google/mobilenet_v1_1.0_224'
__A			:					Any						=			[1, 1_0_2_4, 7, 7]

# Image classification docstring
__A			:					List[str]						=			'google/mobilenet_v1_1.0_224'
__A			:					str						=			'tabby, tabby cat'


__A			:					int						=			[
    '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 	__a     (      A__			:  Dict			,  A__			:  Union[str, Any]			,  A__			:  Optional[int]=None       ):
    SCREAMING_SNAKE_CASE										=   {}

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

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

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

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

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

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

    return tf_to_pt_map





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

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

        SCREAMING_SNAKE_CASE										=   tf_weights[name]

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

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

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

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

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





class 			_SCREAMING_SNAKE_CASE (   nn.Module		):







   '''simple docstring'''




   def __init__(       self					:			Any     , __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     , ):
       super().__init__()
       SCREAMING_SNAKE_CASE										=   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."			)

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

       SCREAMING_SNAKE_CASE										=   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:
           SCREAMING_SNAKE_CASE										=   nn.BatchNormad(
               num_features=__lowerCamelCase     , eps=config.layer_norm_eps     , momentum=0.9_997     , affine=__lowerCamelCase     , track_running_stats=__lowerCamelCase     , )
       else:
           SCREAMING_SNAKE_CASE										=   None

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





   def      _snake_case		(       self					:			Dict     , __lowerCamelCase					:			torch.Tensor			):
       if self.config.tf_padding:
           SCREAMING_SNAKE_CASE										=   apply_tf_padding(__lowerCamelCase     , self.convolution			)
       SCREAMING_SNAKE_CASE										=   self.convolution(__lowerCamelCase			)
       if self.normalization is not None:
           SCREAMING_SNAKE_CASE										=   self.normalization(__lowerCamelCase			)
       if self.activation is not None:
           SCREAMING_SNAKE_CASE										=   self.activation(__lowerCamelCase			)
       return features





class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   lowerCamelCase__				=      MobileNetVaConfig
   lowerCamelCase__				=      load_tf_weights_in_mobilenet_va
   lowerCamelCase__				=      "mobilenet_v1"
   lowerCamelCase__				=      "pixel_values"
   lowerCamelCase__				=      False





   def      _snake_case		(       self					:			List[Any]     , __lowerCamelCase					:			Union[nn.Linear, nn.Convad]			):
       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			:					int						=			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."					,       __snake_case					,       )
class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   def __init__(       self					:			str     , __lowerCamelCase					:			MobileNetVaConfig     , __lowerCamelCase					:			bool = True			):
       super().__init__(__lowerCamelCase			)
       SCREAMING_SNAKE_CASE										=   config

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

       SCREAMING_SNAKE_CASE										=   MobileNetVaConvLayer(
           __lowerCamelCase     , in_channels=config.num_channels     , out_channels=__lowerCamelCase     , kernel_size=3     , stride=2     , )

       SCREAMING_SNAKE_CASE										=   [1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1]

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

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

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

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





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

       SCREAMING_SNAKE_CASE										=   self.conv_stem(__lowerCamelCase			)

       SCREAMING_SNAKE_CASE										=   () if output_hidden_states else None

       for i, layer_module in enumerate(self.layer			):
           SCREAMING_SNAKE_CASE										=   layer_module(__lowerCamelCase			)

           if output_hidden_states:
               SCREAMING_SNAKE_CASE										=   all_hidden_states + (hidden_states,)

       SCREAMING_SNAKE_CASE										=   hidden_states

       if self.pooler is not None:
           SCREAMING_SNAKE_CASE										=   torch.flatten(self.pooler(__lowerCamelCase			)     , start_dim=1			)
       else:
           SCREAMING_SNAKE_CASE										=   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    "					,       __snake_case					,       )
class 			_SCREAMING_SNAKE_CASE (   __snake_case		):







   '''simple docstring'''




   def __init__(       self					:			Optional[Any]     , __lowerCamelCase					:			MobileNetVaConfig			):
       super().__init__(__lowerCamelCase			)

       SCREAMING_SNAKE_CASE										=   config.num_labels
       SCREAMING_SNAKE_CASE										=   MobileNetVaModel(__lowerCamelCase			)

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

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

       SCREAMING_SNAKE_CASE										=   self.mobilenet_va(__lowerCamelCase     , output_hidden_states=__lowerCamelCase     , return_dict=__lowerCamelCase			)

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

       SCREAMING_SNAKE_CASE										=   self.classifier(self.dropout(__lowerCamelCase			)			)

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

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

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

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


"""simple docstring"""





def     __SCREAMING_SNAKE_CASE    ( A_		,	A_					):
	if number < 0 or shift_amount < 0:
		raise ValueError('''both inputs must be positive integers'''					)

	lowerCAmelCase__      :		int    							=      str(bin(A_					)					)
	binary_number += "0" * shift_amount
	return binary_number

def     __SCREAMING_SNAKE_CASE    ( A_		,	A_					):
	if number < 0 or shift_amount < 0:
		raise ValueError('''both inputs must be positive integers'''					)

	lowerCAmelCase__      :		Any    							=      str(bin(A_					)					)[2:]
	if shift_amount >= len(A_					):
		return "0b0"
	lowerCAmelCase__      :		Tuple    							=      binary_number[: len(A_					) - shift_amount]
	return "0b" + shifted_binary_number

def     __SCREAMING_SNAKE_CASE    ( A_		,	A_					):
	if number >= 0:  # Get binary representation of positive number
		lowerCAmelCase__      :		Union[str, Any]    							=      '''0''' + str(bin(A_					)					).strip('''-'''					)[2:]
	else:  # Get binary (2's complement) representation of negative number
		lowerCAmelCase__      :		Optional[Any]    							=      len(bin(A_					)[3:]					)  # Find 2's complement of number
		lowerCAmelCase__      :		Optional[int]    							=      bin(abs(A_					) - (1 << binary_number_length)					)[3:]
		lowerCAmelCase__      :		int    							=      (
		    '''1''' + '''0''' * (binary_number_length - len(A_					)) + binary_number
		)

	if shift_amount >= len(A_					):
		return "0b" + binary_number[0] * len(A_					)
	return (
	    "0b"
	    + binary_number[0] * shift_amount
	    + binary_number[: len(A_					) - shift_amount]
	)


if __name__ == "__main__":
							import doctest

							doctest.testmod()




