






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


__UpperCAmelCase							:   Optional[Any]							=  logging.get_logger(__name__)

# General docstring
__UpperCAmelCase							:   int							=  "RegNetConfig"

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

# Image classification docstring
__UpperCAmelCase							:   Union[str, Any]							=  "facebook/regnet-y-040"
__UpperCAmelCase							:   Any							=  "tabby, tabby cat"

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


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






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






 def    lowerCAmelCase_							(		self					,UpperCamelCase       )						->    str:
     snake_case__      :Tuple         =   self.convolution(self.padding(_A       )       )
     snake_case__      :str         =   self.normalization(_A       )
     snake_case__      :Dict         =   self.activation(_A       )
     return hidden_state




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






 def __init__(		self					,UpperCamelCase					,**UpperCamelCase       )						->    Optional[int]:
     super().__init__(**_A       )
     snake_case__      :Dict         =   config.num_channels
     snake_case__      :str         =   TFRegNetConvLayer(
         out_channels=config.embedding_size					,kernel_size=3					,stride=2					,activation=config.hidden_act					,name="embedder"					,)






 def    lowerCAmelCase_							(		self					,UpperCamelCase       )						->    int:
     snake_case__      :int         =   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)
     snake_case__      :Tuple         =   tf.transpose(_A					,perm=(0, 2, 3, 1)       )
     snake_case__      :str         =   self.embedder(_A       )
     return hidden_state




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






 def __init__(		self					,UpperCamelCase					,UpperCamelCase = 2					,**UpperCamelCase       )						->    Dict:
     super().__init__(**_A       )
     snake_case__      :List[Any]         =   tf.keras.layers.ConvaD(
         filters=_A					,kernel_size=1					,strides=_A					,use_bias=_A					,name="convolution"       )
     snake_case__      :int         =   tf.keras.layers.BatchNormalization(epsilon=1E-5					,momentum=0.9					,name="normalization"       )






 def    lowerCAmelCase_							(		self					,UpperCamelCase					,UpperCamelCase = False       )						->    str:
     return self.normalization(self.convolution(_A       )					,training=_A       )




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






 def __init__(		self					,UpperCamelCase					,UpperCamelCase					,**UpperCamelCase       )						->    Tuple:
     super().__init__(**_A       )
     snake_case__      :Optional[int]         =   tf.keras.layers.GlobalAveragePoolingaD(keepdims=_A					,name="pooler"       )
     snake_case__      :str         =   [
         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    lowerCAmelCase_							(		self					,UpperCamelCase       )						->    Any:
     # [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
     snake_case__      :int         =   self.pooler(_A       )
     for layer_module in self.attention:
         snake_case__      :Tuple         =   layer_module(_A       )
     snake_case__      :int         =   hidden_state * pooled
     return hidden_state




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






 def __init__(		self					,UpperCamelCase					,UpperCamelCase					,UpperCamelCase					,UpperCamelCase = 1					,**UpperCamelCase       )						->    List[str]:
     super().__init__(**_A       )
     snake_case__      :int         =   in_channels != out_channels or stride != 1
     snake_case__      :Union[str, Any]         =   max(1					,out_channels // config.groups_width       )
     snake_case__      :List[Any]         =   (
         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.
     snake_case__      :Tuple         =   [
         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"       ),
     ]
     snake_case__      :Tuple         =   ACTaFN[config.hidden_act]






 def    lowerCAmelCase_							(		self					,UpperCamelCase       )						->    Optional[Any]:
     snake_case__      :List[str]         =   hidden_state
     for layer_module in self.layers:
         snake_case__      :Tuple         =   layer_module(_A       )
     snake_case__      :str         =   self.shortcut(_A       )
     hidden_state += residual
     snake_case__      :List[str]         =   self.activation(_A       )
     return hidden_state




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






 def __init__(		self					,UpperCamelCase					,UpperCamelCase					,UpperCamelCase					,UpperCamelCase = 1					,**UpperCamelCase       )						->    Optional[Any]:
     super().__init__(**_A       )
     snake_case__      :int         =   in_channels != out_channels or stride != 1
     snake_case__      :str         =   max(1					,out_channels // config.groups_width       )
     snake_case__      :Union[str, Any]         =   (
         TFRegNetShortCut(_A					,stride=_A					,name="shortcut"       )
         if should_apply_shortcut
         else tf.keras.layers.Activation("linear"					,name="shortcut"       )
     )
     snake_case__      :int         =   [
         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"       ),
     ]
     snake_case__      :Tuple         =   ACTaFN[config.hidden_act]






 def    lowerCAmelCase_							(		self					,UpperCamelCase       )						->    Union[str, Any]:
     snake_case__      :Optional[int]         =   hidden_state
     for layer_module in self.layers:
         snake_case__      :int         =   layer_module(_A       )
     snake_case__      :List[Any]         =   self.shortcut(_A       )
     hidden_state += residual
     snake_case__      :List[Any]         =   self.activation(_A       )
     return hidden_state




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






 def __init__(		self					,UpperCamelCase					,UpperCamelCase					,UpperCamelCase					,UpperCamelCase = 2					,UpperCamelCase = 2					,**UpperCamelCase       )						->    Union[str, Any]:
     super().__init__(**_A       )

     snake_case__      :Optional[Any]         =   TFRegNetXLayer if config.layer_type == "x" else TFRegNetYLayer
     snake_case__      :Optional[int]         =   [
         # 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    lowerCAmelCase_							(		self					,UpperCamelCase       )						->    Optional[Any]:
     for layer_module in self.layers:
         snake_case__      :Dict         =   layer_module(_A       )
     return hidden_state




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






 def __init__(		self					,UpperCamelCase					,**UpperCamelCase       )						->    List[str]:
     super().__init__(**_A       )
     snake_case__      :Optional[Any]         =   []
     # 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"					,)       )
     snake_case__      :Optional[Any]         =   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    lowerCAmelCase_							(		self					,UpperCamelCase					,UpperCamelCase = False					,UpperCamelCase = True       )						->    Optional[int]:
     snake_case__      :str         =   () if output_hidden_states else None

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

         snake_case__      :Optional[int]         =   stage_module(_A       )

     if output_hidden_states:
         snake_case__      :int         =   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 					_snake_case						(			tf.keras.layers.Layer ):
 _A          =   RegNetConfig






 def __init__(		self					,UpperCamelCase					,**UpperCamelCase       )						->    List[str]:
     super().__init__(**_A       )
     snake_case__      :Optional[Any]         =   config
     snake_case__      :List[str]         =   TFRegNetEmbeddings(_A					,name="embedder"       )
     snake_case__      :List[Any]         =   TFRegNetEncoder(_A					,name="encoder"       )
     snake_case__      :Optional[Any]         =   tf.keras.layers.GlobalAveragePoolingaD(keepdims=_A					,name="pooler"       )






 @unpack_inputs
 def    lowerCAmelCase_							(		self					,UpperCamelCase					,UpperCamelCase = None					,UpperCamelCase = None					,UpperCamelCase = False					,)						->    Tuple:
     snake_case__      :str         =   (
         output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
     )
     snake_case__      :Tuple         =   return_dict if return_dict is not None else self.config.use_return_dict

     snake_case__      :str         =   self.embedder(_A					,training=_A       )

     snake_case__      :Dict         =   self.encoder(
         _A					,output_hidden_states=_A					,return_dict=_A					,training=_A       )

     snake_case__      :str         =   encoder_outputs[0]
     snake_case__      :Union[str, Any]         =   self.pooler(_A       )

     # Change to NCHW output format have uniformity in the modules
     snake_case__      :List[Any]         =   tf.transpose(_A					,perm=(0, 3, 1, 2)       )
     snake_case__      :Tuple         =   tf.transpose(_A					,perm=(0, 3, 1, 2)       )

     # Change the other hidden state outputs to NCHW as well
     if output_hidden_states:
         snake_case__      :Union[str, Any]         =   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 					_snake_case						(			__lowercase ):
 _A          =   RegNetConfig
 _A          =   'regnet'
 _A          =   'pixel_values'






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




__UpperCAmelCase							:   int							=  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"

__UpperCAmelCase							:   List[Any]							=  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.'       ,    __lowercase       ,    )
class 					_snake_case						(			__lowercase ):






 def __init__(		self					,UpperCamelCase					,*UpperCamelCase					,**UpperCamelCase       )						->    List[str]:
     super().__init__(_A					,*_A					,**_A       )
     snake_case__      :Tuple         =   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    lowerCAmelCase_							(		self					,UpperCamelCase					,UpperCamelCase = None					,UpperCamelCase = None					,UpperCamelCase=False					,)						->    Dict:
     snake_case__      :Dict         =   (
         output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
     )
     snake_case__      :List[str]         =   return_dict if return_dict is not None else self.config.use_return_dict

     snake_case__      :str         =   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    '       ,    __lowercase       ,    )
class 					_snake_case						(			__lowercase       ,    __lowercase ):






 def __init__(		self					,UpperCamelCase					,*UpperCamelCase					,**UpperCamelCase       )						->    Any:
     super().__init__(_A					,*_A					,**_A       )
     snake_case__      :Union[str, Any]         =   config.num_labels
     snake_case__      :List[Any]         =   TFRegNetMainLayer(_A					,name="regnet"       )
     # classification head
     snake_case__      :int         =   [
         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    lowerCAmelCase_							(		self					,UpperCamelCase = None					,UpperCamelCase = None					,UpperCamelCase = None					,UpperCamelCase = None					,UpperCamelCase=False					,)						->    Dict:
     snake_case__      :Dict         =   (
         output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
     )
     snake_case__      :Any         =   return_dict if return_dict is not None else self.config.use_return_dict

     snake_case__      :Optional[Any]         =   self.regnet(
         _A					,output_hidden_states=_A					,return_dict=_A					,training=_A       )

     snake_case__      :Optional[int]         =   outputs.pooler_output if return_dict else outputs[1]

     snake_case__      :Dict         =   self.classifier[0](_A       )
     snake_case__      :List[str]         =   self.classifier[1](_A       )

     snake_case__      :Optional[Any]         =   None if labels is None else self.hf_compute_loss(labels=_A					,logits=_A       )

     if not return_dict:
         snake_case__      :List[str]         =   (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 typing import List, Optional, Tuple

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_herbert import HerbertTokenizer


__UpperCAmelCase							:   str							=  logging.get_logger(__name__)

__UpperCAmelCase							:   List[Any]							=  {"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}

__UpperCAmelCase							:   List[Any]							=  {
    "vocab_file": {
        "allegro/herbert-base-cased": "https://huggingface.co/allegro/herbert-base-cased/resolve/main/vocab.json"
    },
    "merges_file": {
        "allegro/herbert-base-cased": "https://huggingface.co/allegro/herbert-base-cased/resolve/main/merges.txt"
    },
}

__UpperCAmelCase							:   str							=  {"allegro/herbert-base-cased": 5_1_4}
__UpperCAmelCase							:   List[str]							=  {}


class 					_snake_case						(			_A ):
	_A          =   VOCAB_FILES_NAMES
	_A          =   PRETRAINED_VOCAB_FILES_MAP
	_A          =   PRETRAINED_INIT_CONFIGURATION
	_A          =   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	_A          =   HerbertTokenizer






	def __init__(		self					,UpperCamelCase=None					,UpperCamelCase=None					,UpperCamelCase=None					,UpperCamelCase="<s>"					,UpperCamelCase="<unk>"					,UpperCamelCase="<pad>"					,UpperCamelCase="<mask>"					,UpperCamelCase="</s>"					,**UpperCamelCase					,)						->    Dict:
					super().__init__(
					    UpperCamelCase					,UpperCamelCase					,tokenizer_file=UpperCamelCase					,cls_token=UpperCamelCase					,unk_token=UpperCamelCase					,pad_token=UpperCamelCase					,mask_token=UpperCamelCase					,sep_token=UpperCamelCase					,**UpperCamelCase					,)






	def    lowerCAmelCase_							(		self					,UpperCamelCase					,UpperCamelCase = None       )						->    List[int]:
					snake_case__      :Optional[int]         =   [self.cls_token_id]
					snake_case__      :Any         =   [self.sep_token_id]
					if token_ids_a is None:
									return cls + token_ids_a + sep

					return cls + token_ids_a + sep + token_ids_a + sep






	def    lowerCAmelCase_							(		self					,UpperCamelCase					,UpperCamelCase = None					,UpperCamelCase = False       )						->    List[int]:
					if already_has_special_tokens:
									return super().get_special_tokens_mask(
									    token_ids_a=UpperCamelCase					,token_ids_a=UpperCamelCase					,already_has_special_tokens=UpperCamelCase       )

					if token_ids_a is None:
									return [1] + ([0] * len(UpperCamelCase       )) + [1]
					return [1] + ([0] * len(UpperCamelCase       )) + [1] + ([0] * len(UpperCamelCase       )) + [1]






	def    lowerCAmelCase_							(		self					,UpperCamelCase					,UpperCamelCase = None       )						->    List[int]:
					snake_case__      :Any         =   [self.sep_token_id]
					snake_case__      :Union[str, Any]         =   [self.cls_token_id]

					if token_ids_a is None:
									return len(cls + token_ids_a + sep       ) * [0]
					return len(cls + token_ids_a + sep       ) * [0] + len(token_ids_a + sep       ) * [1]






	def    lowerCAmelCase_							(		self					,UpperCamelCase					,UpperCamelCase = None       )						->    Tuple[str]:
					snake_case__      :List[str]         =   self._tokenizer.model.save(UpperCamelCase					,name=UpperCamelCase       )
					return tuple(UpperCamelCase       )