

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:    Optional[int] 							=			logging.get_logger(__name__)


# General docstring
_a:    List[str] 							=			"""MobileNetV1Config"""

# Base docstring
_a:    int 							=			"""google/mobilenet_v1_1.0_224"""
_a:    Dict 							=			[1, 1024, 7, 7]

# Image classification docstring
_a:    Tuple 							=			"""google/mobilenet_v1_1.0_224"""
_a:    Optional[int] 							=			"""tabby, tabby cat"""


_a:    Dict 							=			[
    """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      (						A       ,     A       ,     A=None    ):
	UpperCAmelCase_       				=			{}

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

	return tf_to_pt_map



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

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

		UpperCAmelCase_       				=			tf_weights[name]

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

class 						__UpperCamelCase    (					nn.Module						):




	def __init__(		self	:						List[str]		,	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		,	):


		'''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=lowerCAmelCase		,	out_channels=lowerCAmelCase		,	kernel_size=lowerCAmelCase		,	stride=lowerCAmelCase		,	padding=lowerCAmelCase		,	groups=lowerCAmelCase		,	bias=lowerCAmelCase		,	padding_mode="zeros"		,	)

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

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




	def       __A					(		self	:						str		,	lowerCAmelCase	:						torch.Tensor							):


		'''simple docstring'''







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


class 						__UpperCamelCase    (					lowercase						):
	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       __A					(		self	:						Tuple		,	lowerCAmelCase	:						Union[nn.Linear, nn.Convad]							):


		'''simple docstring'''







		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:    List[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.
"""

_a:    str 							=			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.'     ,     lowercase     ,     )
class 						__UpperCamelCase    (					lowercase						):




	def __init__(		self	:						Optional[int]		,	lowerCAmelCase	:						MobileNetVaConfig		,	lowerCAmelCase	:						bool = True							):


		'''simple docstring'''







		super().__init__(lowerCAmelCase							)
		UpperCAmelCase_       				=			config

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

		UpperCAmelCase_       				=			MobileNetVaConvLayer(
		    lowerCAmelCase		,	in_channels=config.num_channels		,	out_channels=lowerCAmelCase		,	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(
			        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		,	)							)

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

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




	def       __A					(		self	:						List[Any]		,	lowerCAmelCase	:						int							):


		'''simple docstring'''







		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       __A					(		self	:						List[Any]		,	lowerCAmelCase	:						Optional[torch.Tensor] = None		,	lowerCAmelCase	:						Optional[bool] = None		,	lowerCAmelCase	:						Optional[bool] = None		,	):


		'''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(lowerCAmelCase							)

		UpperCAmelCase_       				=			() if output_hidden_states else None

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

			if output_hidden_states:
				UpperCAmelCase_       				=			all_hidden_states + (hidden_states,)

		UpperCAmelCase_       				=			hidden_states

		if self.pooler is not None:
			UpperCAmelCase_       				=			torch.flatten(self.pooler(lowerCAmelCase							)		,	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=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    '     ,     lowercase     ,     )
class 						__UpperCamelCase    (					lowercase						):




	def __init__(		self	:						Tuple		,	lowerCAmelCase	:						MobileNetVaConfig							):


		'''simple docstring'''







		super().__init__(lowerCAmelCase							)

		UpperCAmelCase_       				=			config.num_labels
		UpperCAmelCase_       				=			MobileNetVaModel(lowerCAmelCase							)

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

		# Classifier head
		UpperCAmelCase_       				=			nn.Dropout(config.classifier_dropout_prob		,	inplace=lowerCAmelCase							)
		UpperCAmelCase_       				=			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       __A					(		self	:						Tuple		,	lowerCAmelCase	:						Optional[torch.Tensor] = None		,	lowerCAmelCase	:						Optional[bool] = None		,	lowerCAmelCase	:						Optional[torch.Tensor] = None		,	lowerCAmelCase	:						Optional[bool] = None		,	):


		'''simple docstring'''







		UpperCAmelCase_       				=			return_dict if return_dict is not None else self.config.use_return_dict

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

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

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

		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(lowerCAmelCase		,	lowerCAmelCase							)
			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(lowerCAmelCase		,	lowerCAmelCase							)

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

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


from __future__ import annotations

from collections import Counter
from random import random

class 						__UpperCamelCase    :




	def __init__(		self	:						Any							):


		'''simple docstring'''







		UpperCAmelCase_       				=			{}




	def       __A					(		self	:						List[str]		,	lowerCAmelCase	:						str							):


		'''simple docstring'''







		UpperCAmelCase_       				=			{}




	def       __A					(		self	:						Union[str, Any]		,	lowerCAmelCase	:						str		,	lowerCAmelCase	:						str		,	lowerCAmelCase	:						float							):


		'''simple docstring'''







		if nodea not in self.connections:
			self.add_node(lowerCAmelCase							)
		if nodea not in self.connections:
			self.add_node(lowerCAmelCase							)
		UpperCAmelCase_       				=			probability




	def       __A					(		self	:						Tuple							):


		'''simple docstring'''







		return list(self.connections							)




	def       __A					(		self	:						str		,	lowerCAmelCase	:						str							):


		'''simple docstring'''







		UpperCAmelCase_       				=			0
		UpperCAmelCase_       				=			random()

		for dest in self.connections[node]:
			current_probability += self.connections[node][dest]
			if current_probability > random_value:
				return dest
		return ""



def       __lowerCAmelCase      (						A       ,     A       ,     A    ):
	UpperCAmelCase_       				=			MarkovChainGraphUndirectedUnweighted()

	for nodea, nodea, probability in transitions:
		graph.add_transition_probability(A       ,     A       ,     A    )

	UpperCAmelCase_       				=			Counter(graph.get_nodes()    )
	UpperCAmelCase_       				=			start

	for _ in range(A    ):
		UpperCAmelCase_       				=			graph.transition(A    )
		visited[node] += 1

	return visited


if __name__ == "__main__":
				import doctest

				doctest.testmod()