




"""simple docstring"""






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


lowercase__			=		logging.get_logger(__name__)

# General docstring
lowercase__			=		"""RegNetConfig"""

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

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

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






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






			'''simple docstring'''



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






			def    A_       (    self    , lowercase    ):
							_lowerCamelCase	:     str       =       self.convolution(self.padding(lowercase    )    )
							_lowerCamelCase	:     Optional[Any]       =       self.normalization(lowercase    )
							_lowerCamelCase	:     Tuple       =       self.activation(lowercase    )
							return hidden_state






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






			'''simple docstring'''



			def __init__(    self    , lowercase    , **lowercase    ):
							super().__init__(**lowercase    )
							_lowerCamelCase	:     Optional[Any]       =       config.num_channels
							_lowerCamelCase	:     Dict       =       TFRegNetConvLayer(
							    out_channels=config.embedding_size    , kernel_size=3    , stride=2    , activation=config.hidden_act    , name='embedder'    , )






			def    A_       (    self    , lowercase    ):
							_lowerCamelCase	:     List[Any]       =       shape_list(lowercase    )[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	:     Union[str, Any]       =       tf.transpose(lowercase    , perm=(0, 2, 3, 1)    )
							_lowerCamelCase	:     Union[str, Any]       =       self.embedder(lowercase    )
							return hidden_state






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






			'''simple docstring'''



			def __init__(    self    , lowercase    , lowercase = 2    , **lowercase    ):
							super().__init__(**lowercase    )
							_lowerCamelCase	:     List[str]       =       tf.keras.layers.ConvaD(
							    filters=lowercase    , kernel_size=1    , strides=lowercase    , use_bias=lowercase    , name='convolution'    )
							_lowerCamelCase	:     Union[str, Any]       =       tf.keras.layers.BatchNormalization(epsilon=1E-5    , momentum=0.9    , name='normalization'    )






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






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






			'''simple docstring'''



			def __init__(    self    , lowercase    , lowercase    , **lowercase    ):
							super().__init__(**lowercase    )
							_lowerCamelCase	:     str       =       tf.keras.layers.GlobalAveragePoolingaD(keepdims=lowercase    , name='pooler'    )
							_lowerCamelCase	:     Optional[int]       =       [
							    tf.keras.layers.ConvaD(filters=lowercase    , kernel_size=1    , activation='relu'    , name='attention.0'    ),
							    tf.keras.layers.ConvaD(filters=lowercase    , kernel_size=1    , activation='sigmoid'    , name='attention.2'    ),
							]






			def    A_       (    self    , lowercase    ):
							# [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
							_lowerCamelCase	:     Optional[Any]       =       self.pooler(lowercase    )
							for layer_module in self.attention:
											_lowerCamelCase	:     List[Any]       =       layer_module(lowercase    )
							_lowerCamelCase	:     int       =       hidden_state * pooled
							return hidden_state






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






			'''simple docstring'''



			def __init__(    self    , lowercase    , lowercase    , lowercase    , lowercase = 1    , **lowercase    ):
							super().__init__(**lowercase    )
							_lowerCamelCase	:     Optional[Any]       =       in_channels != out_channels or stride != 1
							_lowerCamelCase	:     str       =       max(1    , out_channels // config.groups_width    )
							_lowerCamelCase	:     Tuple       =       (
							    TFRegNetShortCut(lowercase    , stride=lowercase    , 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	:     Union[str, Any]       =       [
							    TFRegNetConvLayer(lowercase    , kernel_size=1    , activation=config.hidden_act    , name='layer.0'    ),
							    TFRegNetConvLayer(
							        lowercase    , stride=lowercase    , groups=lowercase    , activation=config.hidden_act    , name='layer.1'    ),
							    TFRegNetConvLayer(lowercase    , kernel_size=1    , activation=lowercase    , name='layer.2'    ),
							]
							_lowerCamelCase	:     Any       =       ACTaFN[config.hidden_act]






			def    A_       (    self    , lowercase    ):
							_lowerCamelCase	:     List[str]       =       hidden_state
							for layer_module in self.layers:
											_lowerCamelCase	:     Dict       =       layer_module(lowercase    )
							_lowerCamelCase	:     List[Any]       =       self.shortcut(lowercase    )
							hidden_state += residual
							_lowerCamelCase	:     int       =       self.activation(lowercase    )
							return hidden_state






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






			'''simple docstring'''



			def __init__(    self    , lowercase    , lowercase    , lowercase    , lowercase = 1    , **lowercase    ):
							super().__init__(**lowercase    )
							_lowerCamelCase	:     Optional[int]       =       in_channels != out_channels or stride != 1
							_lowerCamelCase	:     List[str]       =       max(1    , out_channels // config.groups_width    )
							_lowerCamelCase	:     Dict       =       (
							    TFRegNetShortCut(lowercase    , stride=lowercase    , name='shortcut'    )
							    if should_apply_shortcut
							    else tf.keras.layers.Activation('linear'    , name='shortcut'    )
							)
							_lowerCamelCase	:     Union[str, Any]       =       [
							    TFRegNetConvLayer(lowercase    , kernel_size=1    , activation=config.hidden_act    , name='layer.0'    ),
							    TFRegNetConvLayer(
							        lowercase    , stride=lowercase    , groups=lowercase    , activation=config.hidden_act    , name='layer.1'    ),
							    TFRegNetSELayer(lowercase    , reduced_channels=int(round(in_channels / 4    )    )    , name='layer.2'    ),
							    TFRegNetConvLayer(lowercase    , kernel_size=1    , activation=lowercase    , name='layer.3'    ),
							]
							_lowerCamelCase	:     Dict       =       ACTaFN[config.hidden_act]






			def    A_       (    self    , lowercase    ):
							_lowerCamelCase	:     Any       =       hidden_state
							for layer_module in self.layers:
											_lowerCamelCase	:     Optional[Any]       =       layer_module(lowercase    )
							_lowerCamelCase	:     Union[str, Any]       =       self.shortcut(lowercase    )
							hidden_state += residual
							_lowerCamelCase	:     List[str]       =       self.activation(lowercase    )
							return hidden_state






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






			'''simple docstring'''



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

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






			def    A_       (    self    , lowercase    ):
							for layer_module in self.layers:
											_lowerCamelCase	:     Any       =       layer_module(lowercase    )
							return hidden_state






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






			'''simple docstring'''



			def __init__(    self    , lowercase    , **lowercase    ):
							super().__init__(**lowercase    )
							_lowerCamelCase	:     int       =       []
							# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
							self.stages.append(
							    TFRegNetStage(
							        lowercase    , 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	:     List[Any]       =       zip(config.hidden_sizes    , config.hidden_sizes[1:]    )
							for i, ((in_channels, out_channels), depth) in enumerate(zip(lowercase    , config.depths[1:]    )    ):
											self.stages.append(TFRegNetStage(lowercase    , lowercase    , lowercase    , depth=lowercase    , name=F'''stages.{i+1}'''    )    )






			def    A_       (    self    , lowercase    , lowercase = False    , lowercase = True    ):
							_lowerCamelCase	:     int       =       () if output_hidden_states else None

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

											_lowerCamelCase	:     List[str]       =       stage_module(lowercase    )

							if output_hidden_states:
											_lowerCamelCase	:     Dict       =       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=lowercase    , hidden_states=lowercase    )






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






			'''simple docstring'''



			lowerCamelCase__ =     RegNetConfig
			def __init__(    self    , lowercase    , **lowercase    ):
							super().__init__(**lowercase    )
							_lowerCamelCase	:     List[str]       =       config
							_lowerCamelCase	:     Union[str, Any]       =       TFRegNetEmbeddings(lowercase    , name='embedder'    )
							_lowerCamelCase	:     Union[str, Any]       =       TFRegNetEncoder(lowercase    , name='encoder'    )
							_lowerCamelCase	:     Union[str, Any]       =       tf.keras.layers.GlobalAveragePoolingaD(keepdims=lowercase    , name='pooler'    )






			@unpack_inputs
			def    A_       (    self    , lowercase    , lowercase = None    , lowercase = None    , lowercase = False    , ):
							_lowerCamelCase	:     Tuple       =       (
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							_lowerCamelCase	:     Tuple       =       return_dict if return_dict is not None else self.config.use_return_dict

							_lowerCamelCase	:     str       =       self.embedder(lowercase    , training=lowercase    )

							_lowerCamelCase	:     Union[str, Any]       =       self.encoder(
							    lowercase    , output_hidden_states=lowercase    , return_dict=lowercase    , training=lowercase    )

							_lowerCamelCase	:     Any       =       encoder_outputs[0]
							_lowerCamelCase	:     List[str]       =       self.pooler(lowercase    )

							# Change to NCHW output format have uniformity in the modules
							_lowerCamelCase	:     List[str]       =       tf.transpose(lowercase    , perm=(0, 3, 1, 2)    )
							_lowerCamelCase	:     str       =       tf.transpose(lowercase    , perm=(0, 3, 1, 2)    )

							# Change the other hidden state outputs to NCHW as well
							if output_hidden_states:
											_lowerCamelCase	:     List[Any]       =       tuple([tf.transpose(lowercase    , 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=lowercase    , pooler_output=lowercase    , hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states    , )






class lowerCAmelCase__      (	lowercase		):






			'''simple docstring'''



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

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

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






			'''simple docstring'''



			def __init__(    self    , lowercase    , *lowercase    , **lowercase    ):
							super().__init__(lowercase    , *lowercase    , **lowercase    )
							_lowerCamelCase	:     str       =       TFRegNetMainLayer(lowercase    , name='regnet'    )






			@unpack_inputs
			@add_start_docstrings_to_model_forward(lowercase    )
			@add_code_sample_docstrings(
			    checkpoint=_CHECKPOINT_FOR_DOC    , output_type=lowercase    , config_class=_CONFIG_FOR_DOC    , modality='vision'    , expected_output=_EXPECTED_OUTPUT_SHAPE    , )
			def    A_       (    self    , lowercase    , lowercase = None    , lowercase = None    , lowercase=False    , ):
							_lowerCamelCase	:     List[str]       =       (
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							_lowerCamelCase	:     List[Any]       =       return_dict if return_dict is not None else self.config.use_return_dict

							_lowerCamelCase	:     List[Any]       =       self.regnet(
							    pixel_values=lowercase    , output_hidden_states=lowercase    , return_dict=lowercase    , training=lowercase    , )
							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(
    """
    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    """,      lowercase,      )
class lowerCAmelCase__      (	lowercase,      lowercase		):






			'''simple docstring'''



			def __init__(    self    , lowercase    , *lowercase    , **lowercase    ):
							super().__init__(lowercase    , *lowercase    , **lowercase    )
							_lowerCamelCase	:     str       =       config.num_labels
							_lowerCamelCase	:     str       =       TFRegNetMainLayer(lowercase    , name='regnet'    )
							# classification head
							_lowerCamelCase	:     str       =       [
							    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(lowercase    )
			@add_code_sample_docstrings(
			    checkpoint=_IMAGE_CLASS_CHECKPOINT    , output_type=lowercase    , config_class=_CONFIG_FOR_DOC    , expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT    , )
			def    A_       (    self    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase = None    , lowercase=False    , ):
							_lowerCamelCase	:     Dict       =       (
							    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
							)
							_lowerCamelCase	:     Optional[Any]       =       return_dict if return_dict is not None else self.config.use_return_dict

							_lowerCamelCase	:     Optional[Any]       =       self.regnet(
							    lowercase    , output_hidden_states=lowercase    , return_dict=lowercase    , training=lowercase    )

							_lowerCamelCase	:     Dict       =       outputs.pooler_output if return_dict else outputs[1]

							_lowerCamelCase	:     int       =       self.classifier[0](lowercase    )
							_lowerCamelCase	:     Optional[Any]       =       self.classifier[1](lowercase    )

							_lowerCamelCase	:     Tuple       =       None if labels is None else self.hf_compute_loss(labels=lowercase    , logits=lowercase    )

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

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





"""simple docstring"""






from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase__			=		logging.get_logger(__name__)

lowercase__			=		{
    """facebook/dpr-ctx_encoder-single-nq-base""": (
        """https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/config.json"""
    ),
    """facebook/dpr-question_encoder-single-nq-base""": (
        """https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/config.json"""
    ),
    """facebook/dpr-reader-single-nq-base""": (
        """https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/config.json"""
    ),
    """facebook/dpr-ctx_encoder-multiset-base""": (
        """https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/config.json"""
    ),
    """facebook/dpr-question_encoder-multiset-base""": (
        """https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/config.json"""
    ),
    """facebook/dpr-reader-multiset-base""": (
        """https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/config.json"""
    ),
}






class lowerCAmelCase__      (	lowercase		):






			'''simple docstring'''



			lowerCamelCase__ =     """dpr"""
			def __init__(    self    , lowercase=30522    , lowercase=768    , lowercase=12    , lowercase=12    , lowercase=3072    , lowercase="gelu"    , lowercase=0.1    , lowercase=0.1    , lowercase=512    , lowercase=2    , lowercase=0.02    , lowercase=1E-12    , lowercase=0    , lowercase="absolute"    , lowercase = 0    , **lowercase    , ):
							super().__init__(pad_token_id=lowercase    , **lowercase    )

							_lowerCamelCase	:     int       =       vocab_size
							_lowerCamelCase	:     Optional[int]       =       hidden_size
							_lowerCamelCase	:     str       =       num_hidden_layers
							_lowerCamelCase	:     Optional[int]       =       num_attention_heads
							_lowerCamelCase	:     Tuple       =       hidden_act
							_lowerCamelCase	:     int       =       intermediate_size
							_lowerCamelCase	:     str       =       hidden_dropout_prob
							_lowerCamelCase	:     Tuple       =       attention_probs_dropout_prob
							_lowerCamelCase	:     Union[str, Any]       =       max_position_embeddings
							_lowerCamelCase	:     str       =       type_vocab_size
							_lowerCamelCase	:     int       =       initializer_range
							_lowerCamelCase	:     Optional[int]       =       layer_norm_eps
							_lowerCamelCase	:     Dict       =       projection_dim
							_lowerCamelCase	:     int       =       position_embedding_type