






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      :       Union[str, Any]    =						logging.get_logger(__name__)


# General docstring
lowerCamelCase      :       str    =						'''MobileNetV1Config'''

# Base docstring
lowerCamelCase      :       Union[str, Any]    =						'''google/mobilenet_v1_1.0_224'''
lowerCamelCase      :       Union[str, Any]    =						[1, 1024, 7, 7]

# Image classification docstring
lowerCamelCase      :       str    =						'''google/mobilenet_v1_1.0_224'''
lowerCamelCase      :       Union[str, Any]    =						'''tabby, tabby cat'''


lowerCamelCase      :       Tuple    =						[
    '''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   (    __snake_case       ,			__snake_case       ,			__snake_case=None							):
     __lowerCAmelCase             =       {}

     if isinstance(__snake_case       ,			__snake_case							):
          __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(13							):
          __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(__snake_case       ,			__snake_case							):
          __lowerCAmelCase             =       "MobilenetV1/Logits/Conv2d_1c_1x1/"
          __lowerCAmelCase             =       model.classifier.weight
          __lowerCAmelCase             =       model.classifier.bias

     return tf_to_pt_map





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

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

          __lowerCAmelCase             =       tf_weights[name]

          if "depthwise_weights" in name:
               logger.info("Transposing depthwise"							)
               __lowerCAmelCase             =       np.transpose(__snake_case       ,			(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(__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}"""							)
          __lowerCAmelCase             =       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    __lowerCAmelCase   (    __snake_case       ,			__snake_case							):
     __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(__snake_case       ,			__snake_case       ,			"constant"       ,			0.0							)



class 				_UpperCamelCase      (nn.Module      ):


      def __init__(						self       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase = 1       ,    __UpperCamelCase = 1       ,    __UpperCamelCase = False       ,    __UpperCamelCase = True       ,    __UpperCamelCase = True       ,    )->						None:
           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.9_9_9_7       ,    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 			__UpperCAmelCase		(						self       ,    __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 				_UpperCamelCase      (a_      ):
      snake_case_	=				MobileNetVaConfig
      snake_case_	=				load_tf_weights_in_mobilenet_va
      snake_case_	=				"""mobilenet_v1"""
      snake_case_	=				"""pixel_values"""
      snake_case_	=				False


      def 			__UpperCAmelCase		(						self       ,    __UpperCamelCase						)->						None:
           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						)

lowerCamelCase      :       Union[str, Any]    =						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.
'''

lowerCamelCase      :       int    =						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."""				,		a_				,		)
class 				_UpperCamelCase      (a_      ):


      def __init__(						self       ,    __UpperCamelCase       ,    __UpperCamelCase = True						)->						Any:
           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 			__UpperCAmelCase		(						self       ,    __UpperCamelCase						)->						int:
           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 			__UpperCAmelCase		(						self       ,    __UpperCamelCase = None       ,    __UpperCamelCase = None       ,    __UpperCamelCase = None       ,    )->						Union[tuple, BaseModelOutputWithPoolingAndNoAttention]:
           __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.
    """				,		a_				,		)
class 				_UpperCamelCase      (a_      ):


      def __init__(						self       ,    __UpperCamelCase						)->						None:
           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 			__UpperCAmelCase		(						self       ,    __UpperCamelCase = None       ,    __UpperCamelCase = None       ,    __UpperCamelCase = None       ,    __UpperCamelCase = None       ,    )->						Union[tuple, ImageClassifierOutputWithNoAttention]:
           __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       ,    )














import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
 import torch

 from transformers.generation import DisjunctiveConstraint



@require_torch
class 				_UpperCamelCase      (unittest.TestCase      ):


      def 			__UpperCAmelCase		(						self						)->						List[Any]:
           # For consistency across different places the DisjunctiveConstraint is called,
           # dc.token_ids is a list of integers. It is also initialized only by integers.

           __lowerCAmelCase             =       [[1, 2, 4], [1, 2, 3, 4]]
           __lowerCAmelCase             =       DisjunctiveConstraint(__UpperCamelCase						)
           self.assertTrue(isinstance(dc.token_ids       ,    __UpperCamelCase						)						)

           with self.assertRaises(__UpperCamelCase						):
                DisjunctiveConstraint(torch.LongTensor([[1, 2, 4], [1, 2, 3]]						)						)

           with self.assertRaises(__UpperCamelCase						):
                DisjunctiveConstraint([torch.LongTensor([1, 2, 4]						), torch.LongTensor([1, 2, 3, 4, 5]						)]						)


      def 			__UpperCAmelCase		(						self						)->						Tuple:
           # We can't have constraints that are complete subsets of another. This leads to a preverse
           # interpretation of "constraint fulfillment": does generating [1,2,3] fulfill the constraint?
           # It would mean that it generated [1,2] which fulfills it, but it's in the middle of potentially
           # fulfilling [1,2,3,4]. If we believe that [1,2,3] does fulfill the constraint, then the algorithm
           # will necessarily never reach [1,2,3,4], giving users a false sense of control (better to just not allow it).
           __lowerCAmelCase             =       [[1, 2], [1, 2, 3, 4]]

           with self.assertRaises(__UpperCamelCase						):
                DisjunctiveConstraint(__UpperCamelCase						)  # fails here


      def 			__UpperCAmelCase		(						self						)->						List[Any]:
           __lowerCAmelCase             =       [[1, 2, 3], [1, 2, 4]]

           __lowerCAmelCase             =       DisjunctiveConstraint(__UpperCamelCase						)

           __lowerCAmelCase      ,			__lowerCAmelCase      ,			__lowerCAmelCase             =       dc.update(1						)
           __lowerCAmelCase             =       stepped is True and completed is False and reset is False
           self.assertTrue(__UpperCamelCase						)
           self.assertTrue(not dc.completed						)
           self.assertTrue(dc.current_seq == [1]						)

           __lowerCAmelCase      ,			__lowerCAmelCase      ,			__lowerCAmelCase             =       dc.update(2						)
           __lowerCAmelCase             =       stepped is True and completed is False and reset is False
           self.assertTrue(__UpperCamelCase						)
           self.assertTrue(not dc.completed						)
           self.assertTrue(dc.current_seq == [1, 2]						)

           __lowerCAmelCase      ,			__lowerCAmelCase      ,			__lowerCAmelCase             =       dc.update(3						)
           __lowerCAmelCase             =       stepped is True and completed is True and reset is False
           self.assertTrue(__UpperCamelCase						)
           self.assertTrue(dc.completed						)  # Completed!
           self.assertTrue(dc.current_seq == [1, 2, 3]						)


      def 			__UpperCAmelCase		(						self						)->						int:
           __lowerCAmelCase             =       [[1, 2, 3], [1, 2, 4, 5], [1, 2, 5]]

           __lowerCAmelCase             =       DisjunctiveConstraint(__UpperCamelCase						)

           __lowerCAmelCase      ,			__lowerCAmelCase      ,			__lowerCAmelCase             =       dc.update(1						)
           self.assertTrue(not dc.completed						)
           self.assertTrue(dc.current_seq == [1]						)

           __lowerCAmelCase      ,			__lowerCAmelCase      ,			__lowerCAmelCase             =       dc.update(2						)
           self.assertTrue(not dc.completed						)
           self.assertTrue(dc.current_seq == [1, 2]						)

           __lowerCAmelCase      ,			__lowerCAmelCase      ,			__lowerCAmelCase             =       dc.update(4						)
           self.assertTrue(not dc.completed						)
           self.assertTrue(dc.current_seq == [1, 2, 4]						)

           __lowerCAmelCase      ,			__lowerCAmelCase      ,			__lowerCAmelCase             =       dc.update(5						)
           self.assertTrue(dc.completed						)  # Completed!
           self.assertTrue(dc.current_seq == [1, 2, 4, 5]						)

           dc.reset()

           __lowerCAmelCase      ,			__lowerCAmelCase      ,			__lowerCAmelCase             =       dc.update(1						)
           self.assertTrue(not dc.completed						)
           self.assertTrue(dc.remaining() == 3						)
           self.assertTrue(dc.current_seq == [1]						)

           __lowerCAmelCase      ,			__lowerCAmelCase      ,			__lowerCAmelCase             =       dc.update(2						)
           self.assertTrue(not dc.completed						)
           self.assertTrue(dc.remaining() == 2						)
           self.assertTrue(dc.current_seq == [1, 2]						)

           __lowerCAmelCase      ,			__lowerCAmelCase      ,			__lowerCAmelCase             =       dc.update(5						)
           self.assertTrue(dc.completed						)  # Completed!
           self.assertTrue(dc.remaining() == 0						)
           self.assertTrue(dc.current_seq == [1, 2, 5]						)






