





from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFSequenceClassifierOutput,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import logging
from .configuration_regnet import RegNetConfig


A_												=     logging.get_logger(__name__)

# General docstring
A_												=     "RegNetConfig"

# Base docstring
A_												=     "facebook/regnet-y-040"
A_												=     [1, 1_0_8_8, 7, 7]

# Image classification docstring
A_												=     "facebook/regnet-y-040"
A_												=     "tabby, tabby cat"

A_												=     [
    "facebook/regnet-y-040",
    # See all regnet models at https://huggingface.co/models?filter=regnet
]

class SCREAMING_SNAKE_CASE_ (					tf.keras.layers.Layer   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase = 3	,  _lowerCAmelCase = 1	,  _lowerCAmelCase = 1	,  _lowerCAmelCase = "relu"	,  **_lowerCAmelCase	,  ):
										super().__init__(**A_	)
										# The padding and conv has been verified in
										# https://colab.research.google.com/gist/sayakpaul/854bc10eeaf21c9ee2119e0b9f3841a7/scratchpad.ipynb
										lowerCamelCase__       =       tf.keras.layers.ZeroPaddingaD(padding=kernel_size // 2	)
										lowerCamelCase__       =       tf.keras.layers.ConvaD(
										    filters=A_	,  kernel_size=A_	,  strides=A_	,  padding="VALID"	,  groups=A_	,  use_bias=A_	,  name="convolution"	,  )
										lowerCamelCase__       =       tf.keras.layers.BatchNormalization(epsilon=1E-5	,  momentum=0.9	,  name="normalization"	)
										lowerCamelCase__       =       ACTaFN[activation] if activation is not None else tf.identity






					def 					__magic_name__ ( self	,  _lowerCAmelCase	):
										lowerCamelCase__       =       self.convolution(self.padding(A_	)	)
										lowerCamelCase__       =       self.normalization(A_	)
										lowerCamelCase__       =       self.activation(A_	)
										return hidden_state

class SCREAMING_SNAKE_CASE_ (					tf.keras.layers.Layer   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  **_lowerCAmelCase	):
										super().__init__(**A_	)
										lowerCamelCase__       =       config.num_channels
										lowerCamelCase__       =       TFRegNetConvLayer(
										    out_channels=config.embedding_size	,  kernel_size=3	,  stride=2	,  activation=config.hidden_act	,  name="embedder"	,  )






					def 					__magic_name__ ( self	,  _lowerCAmelCase	):
										lowerCamelCase__       =       shape_list(A_	)[1]
										if tf.executing_eagerly() and num_channels != self.num_channels:
															raise ValueError(
															    "Make sure that the channel dimension of the pixel values match with the one set in the configuration."	)

										# When running on CPU, `tf.keras.layers.Conv2D` doesn't support `NCHW` format.
										# So change the input format from `NCHW` to `NHWC`.
										# shape = (batch_size, in_height, in_width, in_channels=num_channels)
										lowerCamelCase__       =       tf.transpose(A_	,  perm=(0, 2, 3, 1)	)
										lowerCamelCase__       =       self.embedder(A_	)
										return hidden_state

class SCREAMING_SNAKE_CASE_ (					tf.keras.layers.Layer   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase = 2	,  **_lowerCAmelCase	):
										super().__init__(**A_	)
										lowerCamelCase__       =       tf.keras.layers.ConvaD(
										    filters=A_	,  kernel_size=1	,  strides=A_	,  use_bias=A_	,  name="convolution"	)
										lowerCamelCase__       =       tf.keras.layers.BatchNormalization(epsilon=1E-5	,  momentum=0.9	,  name="normalization"	)






					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase = False	):
										return self.normalization(self.convolution(A_	)	,  training=A_	)

class SCREAMING_SNAKE_CASE_ (					tf.keras.layers.Layer   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  **_lowerCAmelCase	):
										super().__init__(**A_	)
										lowerCamelCase__       =       tf.keras.layers.GlobalAveragePoolingaD(keepdims=A_	,  name="pooler"	)
										lowerCamelCase__       =       [
										    tf.keras.layers.ConvaD(filters=A_	,  kernel_size=1	,  activation="relu"	,  name="attention.0"	),
										    tf.keras.layers.ConvaD(filters=A_	,  kernel_size=1	,  activation="sigmoid"	,  name="attention.2"	),
										]






					def 					__magic_name__ ( self	,  _lowerCAmelCase	):
										lowerCamelCase__       =       self.pooler(A_	)
										for layer_module in self.attention:
															lowerCamelCase__       =       layer_module(A_	)
										lowerCamelCase__       =       hidden_state * pooled
										return hidden_state

class SCREAMING_SNAKE_CASE_ (					tf.keras.layers.Layer   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase = 1	,  **_lowerCAmelCase	):
										super().__init__(**A_	)
										lowerCamelCase__       =       in_channels != out_channels or stride != 1
										lowerCamelCase__       =       max(1	,  out_channels // config.groups_width	)
										lowerCamelCase__       =       (
										    TFRegNetShortCut(A_	,  stride=A_	,  name="shortcut"	)
										    if should_apply_shortcut
										    else tf.keras.layers.Activation("linear"	,  name="shortcut"	)
										)
										# `self.layers` instead of `self.layer` because that is a reserved argument.
										lowerCamelCase__       =       [
										    TFRegNetConvLayer(A_	,  kernel_size=1	,  activation=config.hidden_act	,  name="layer.0"	),
										    TFRegNetConvLayer(
										        A_	,  stride=A_	,  groups=A_	,  activation=config.hidden_act	,  name="layer.1"	),
										    TFRegNetConvLayer(A_	,  kernel_size=1	,  activation=A_	,  name="layer.2"	),
										]
										lowerCamelCase__       =       ACTaFN[config.hidden_act]






					def 					__magic_name__ ( self	,  _lowerCAmelCase	):
										lowerCamelCase__       =       hidden_state
										for layer_module in self.layers:
															lowerCamelCase__       =       layer_module(A_	)
										lowerCamelCase__       =       self.shortcut(A_	)
										hidden_state += residual
										lowerCamelCase__       =       self.activation(A_	)
										return hidden_state

class SCREAMING_SNAKE_CASE_ (					tf.keras.layers.Layer   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase = 1	,  **_lowerCAmelCase	):
										super().__init__(**A_	)
										lowerCamelCase__       =       in_channels != out_channels or stride != 1
										lowerCamelCase__       =       max(1	,  out_channels // config.groups_width	)
										lowerCamelCase__       =       (
										    TFRegNetShortCut(A_	,  stride=A_	,  name="shortcut"	)
										    if should_apply_shortcut
										    else tf.keras.layers.Activation("linear"	,  name="shortcut"	)
										)
										lowerCamelCase__       =       [
										    TFRegNetConvLayer(A_	,  kernel_size=1	,  activation=config.hidden_act	,  name="layer.0"	),
										    TFRegNetConvLayer(
										        A_	,  stride=A_	,  groups=A_	,  activation=config.hidden_act	,  name="layer.1"	),
										    TFRegNetSELayer(A_	,  reduced_channels=int(round(in_channels / 4	)	)	,  name="layer.2"	),
										    TFRegNetConvLayer(A_	,  kernel_size=1	,  activation=A_	,  name="layer.3"	),
										]
										lowerCamelCase__       =       ACTaFN[config.hidden_act]






					def 					__magic_name__ ( self	,  _lowerCAmelCase	):
										lowerCamelCase__       =       hidden_state
										for layer_module in self.layers:
															lowerCamelCase__       =       layer_module(A_	)
										lowerCamelCase__       =       self.shortcut(A_	)
										hidden_state += residual
										lowerCamelCase__       =       self.activation(A_	)
										return hidden_state

class SCREAMING_SNAKE_CASE_ (					tf.keras.layers.Layer   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase = 2	,  _lowerCAmelCase = 2	,  **_lowerCAmelCase	):
										super().__init__(**A_	)

										lowerCamelCase__       =       TFRegNetXLayer if config.layer_type == "x" else TFRegNetYLayer
										lowerCamelCase__       =       [
										    # downsampling is done in the first layer with stride of 2
										    layer(A_	,  A_	,  A_	,  stride=A_	,  name="layers.0"	),
										    *[layer(A_	,  A_	,  A_	,  name=F"layers.{i+1}"	) for i in range(depth - 1	)],
										]






					def 					__magic_name__ ( self	,  _lowerCAmelCase	):
										for layer_module in self.layers:
															lowerCamelCase__       =       layer_module(A_	)
										return hidden_state

class SCREAMING_SNAKE_CASE_ (					tf.keras.layers.Layer   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  **_lowerCAmelCase	):
										super().__init__(**A_	)
										lowerCamelCase__       =       []
										# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
										self.stages.append(
										    TFRegNetStage(
										        A_	,  config.embedding_size	,  config.hidden_sizes[0]	,  stride=2 if config.downsample_in_first_stage else 1	,  depth=config.depths[0]	,  name="stages.0"	,  )	)
										lowerCamelCase__       =       zip(config.hidden_sizes	,  config.hidden_sizes[1:]	)
										for i, ((in_channels, out_channels), depth) in enumerate(zip(A_	,  config.depths[1:]	)	):
															self.stages.append(TFRegNetStage(A_	,  A_	,  A_	,  depth=A_	,  name=F"stages.{i+1}"	)	)






					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase = False	,  _lowerCAmelCase = True	):
										lowerCamelCase__       =       () if output_hidden_states else None

										for stage_module in self.stages:
															if output_hidden_states:
																				lowerCamelCase__       =       hidden_states + (hidden_state,)

															lowerCamelCase__       =       stage_module(A_	)

										if output_hidden_states:
															lowerCamelCase__       =       hidden_states + (hidden_state,)

										if not return_dict:
															return tuple(v for v in [hidden_state, hidden_states] if v is not None	)

										return TFBaseModelOutputWithNoAttention(last_hidden_state=A_	,  hidden_states=A_	)

@keras_serializable
class SCREAMING_SNAKE_CASE_ (					tf.keras.layers.Layer   ):





					"""simple docstring"""


					A__				 =    RegNetConfig





					def __init__( self	,  _lowerCAmelCase	,  **_lowerCAmelCase	):
										super().__init__(**A_	)
										lowerCamelCase__       =       config
										lowerCamelCase__       =       TFRegNetEmbeddings(A_	,  name="embedder"	)
										lowerCamelCase__       =       TFRegNetEncoder(A_	,  name="encoder"	)
										lowerCamelCase__       =       tf.keras.layers.GlobalAveragePoolingaD(keepdims=A_	,  name="pooler"	)






					@unpack_inputs
					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase = None	,  _lowerCAmelCase = None	,  _lowerCAmelCase = False	,  ):
										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

										lowerCamelCase__       =       self.embedder(A_	,  training=A_	)

										lowerCamelCase__       =       self.encoder(
										    A_	,  output_hidden_states=A_	,  return_dict=A_	,  training=A_	)

										lowerCamelCase__       =       encoder_outputs[0]
										lowerCamelCase__       =       self.pooler(A_	)

										# Change to NCHW output format have uniformity in the modules
										lowerCamelCase__       =       tf.transpose(A_	,  perm=(0, 3, 1, 2)	)
										lowerCamelCase__       =       tf.transpose(A_	,  perm=(0, 3, 1, 2)	)

										# Change the other hidden state outputs to NCHW as well
										if output_hidden_states:
															lowerCamelCase__       =       tuple([tf.transpose(A_	,  perm=(0, 3, 1, 2)	) for h in encoder_outputs[1]]	)

										if not return_dict:
															return (last_hidden_state, pooled_output) + encoder_outputs[1:]

										return TFBaseModelOutputWithPoolingAndNoAttention(
										    last_hidden_state=A_	,  pooler_output=A_	,  hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states	,  )

class SCREAMING_SNAKE_CASE_ (					_SCREAMING_SNAKE_CASE   ):





					"""simple docstring"""


					A__				 =    RegNetConfig
					A__				 =    "regnet"
					A__				 =    "pixel_values"





					@property
					def 					__magic_name__ ( self	):
										return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)	,  dtype=tf.floataa	)}

A_												=     R"\n    Parameters:\n    This model is a Tensorflow\n    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a\n    regular Tensorflow Module and refer to the Tensorflow documentation for all matter related to general usage and\n    behavior.\n        config ([`RegNetConfig`]): 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 [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.\n"

A_												=     R"\n    Args:\n        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConveNextImageProcessor.__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 RegNet model outputting raw features without any specific head on top."							,       _SCREAMING_SNAKE_CASE							,       )
class SCREAMING_SNAKE_CASE_ (					_SCREAMING_SNAKE_CASE   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  *_lowerCAmelCase	,  **_lowerCAmelCase	):
										super().__init__(A_	,  *A_	,  **A_	)
										lowerCamelCase__       =       TFRegNetMainLayer(A_	,  name="regnet"	)






					@unpack_inputs
					@add_start_docstrings_to_model_forward(A_	)
					@add_code_sample_docstrings(
					    checkpoint=_CHECKPOINT_FOR_DOC	,  output_type=A_	,  config_class=_CONFIG_FOR_DOC	,  modality="vision"	,  expected_output=_EXPECTED_OUTPUT_SHAPE	,  )
					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase = None	,  _lowerCAmelCase = None	,  _lowerCAmelCase=False	,  ):
										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

										lowerCamelCase__       =       self.regnet(
										    pixel_values=A_	,  output_hidden_states=A_	,  return_dict=A_	,  training=A_	,  )
										if not return_dict:
															return (outputs[0],) + outputs[1:]

										return TFBaseModelOutputWithPoolingAndNoAttention(
										    last_hidden_state=outputs.last_hidden_state	,  pooler_output=outputs.pooler_output	,  hidden_states=outputs.hidden_states	,  )

@add_start_docstrings(
    "\n    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    "							,       _SCREAMING_SNAKE_CASE							,       )
class SCREAMING_SNAKE_CASE_ (					_SCREAMING_SNAKE_CASE							,       _SCREAMING_SNAKE_CASE   ):





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  *_lowerCAmelCase	,  **_lowerCAmelCase	):
										super().__init__(A_	,  *A_	,  **A_	)
										lowerCamelCase__       =       config.num_labels
										lowerCamelCase__       =       TFRegNetMainLayer(A_	,  name="regnet"	)
										# classification head
										lowerCamelCase__       =       [
										    tf.keras.layers.Flatten(),
										    tf.keras.layers.Dense(config.num_labels	,  name="classifier.1"	) if config.num_labels > 0 else tf.identity,
										]






					@unpack_inputs
					@add_start_docstrings_to_model_forward(A_	)
					@add_code_sample_docstrings(
					    checkpoint=_IMAGE_CLASS_CHECKPOINT	,  output_type=A_	,  config_class=_CONFIG_FOR_DOC	,  expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT	,  )
					def 					__magic_name__ ( self	,  _lowerCAmelCase = None	,  _lowerCAmelCase = None	,  _lowerCAmelCase = None	,  _lowerCAmelCase = None	,  _lowerCAmelCase=False	,  ):
										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

										lowerCamelCase__       =       self.regnet(
										    A_	,  output_hidden_states=A_	,  return_dict=A_	,  training=A_	)

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

										lowerCamelCase__       =       self.classifier[0](A_	)
										lowerCamelCase__       =       self.classifier[1](A_	)

										lowerCamelCase__       =       None if labels is None else self.hf_compute_loss(labels=A_	,  logits=A_	)

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

										return TFSequenceClassifierOutput(loss=A_	,  logits=A_	,  hidden_states=outputs.hidden_states	)








from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig


A_												=     {
    "albert-base-v1": "https://huggingface.co/albert-base-v1/resolve/main/config.json",
    "albert-large-v1": "https://huggingface.co/albert-large-v1/resolve/main/config.json",
    "albert-xlarge-v1": "https://huggingface.co/albert-xlarge-v1/resolve/main/config.json",
    "albert-xxlarge-v1": "https://huggingface.co/albert-xxlarge-v1/resolve/main/config.json",
    "albert-base-v2": "https://huggingface.co/albert-base-v2/resolve/main/config.json",
    "albert-large-v2": "https://huggingface.co/albert-large-v2/resolve/main/config.json",
    "albert-xlarge-v2": "https://huggingface.co/albert-xlarge-v2/resolve/main/config.json",
    "albert-xxlarge-v2": "https://huggingface.co/albert-xxlarge-v2/resolve/main/config.json",
}

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					A__				 =    "albert"





					def __init__( self	,  _lowerCAmelCase=3_0000	,  _lowerCAmelCase=128	,  _lowerCAmelCase=4096	,  _lowerCAmelCase=12	,  _lowerCAmelCase=1	,  _lowerCAmelCase=64	,  _lowerCAmelCase=1_6384	,  _lowerCAmelCase=1	,  _lowerCAmelCase="gelu_new"	,  _lowerCAmelCase=0	,  _lowerCAmelCase=0	,  _lowerCAmelCase=512	,  _lowerCAmelCase=2	,  _lowerCAmelCase=0.02	,  _lowerCAmelCase=1E-12	,  _lowerCAmelCase=0.1	,  _lowerCAmelCase="absolute"	,  _lowerCAmelCase=0	,  _lowerCAmelCase=2	,  _lowerCAmelCase=3	,  **_lowerCAmelCase	,  ):
										super().__init__(pad_token_id=_lowerCAmelCase	,  bos_token_id=_lowerCAmelCase	,  eos_token_id=_lowerCAmelCase	,  **_lowerCAmelCase	)

										lowerCamelCase__       =       vocab_size
										lowerCamelCase__       =       embedding_size
										lowerCamelCase__       =       hidden_size
										lowerCamelCase__       =       num_hidden_layers
										lowerCamelCase__       =       num_hidden_groups
										lowerCamelCase__       =       num_attention_heads
										lowerCamelCase__       =       inner_group_num
										lowerCamelCase__       =       hidden_act
										lowerCamelCase__       =       intermediate_size
										lowerCamelCase__       =       hidden_dropout_prob
										lowerCamelCase__       =       attention_probs_dropout_prob
										lowerCamelCase__       =       max_position_embeddings
										lowerCamelCase__       =       type_vocab_size
										lowerCamelCase__       =       initializer_range
										lowerCamelCase__       =       layer_norm_eps
										lowerCamelCase__       =       classifier_dropout_prob
										lowerCamelCase__       =       position_embedding_type

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					@property
					def 					__magic_name__ ( self	):
										if self.task == "multiple-choice":
															lowerCamelCase__       =       {0: "batch", 1: "choice", 2: "sequence"}
										else:
															lowerCamelCase__       =       {0: "batch", 1: "sequence"}
										return OrderedDict(
										    [
										        ("input_ids", dynamic_axis),
										        ("attention_mask", dynamic_axis),
										        ("token_type_ids", dynamic_axis),
										    ]	)

