
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  lowercase__(				tf.keras.layers.Layer							):



	"""simple docstring"""




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




	def 						_lowercase  (       self  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       )       ->     Any:
					lowercase_       =					self.convolution(self.padding(SCREAMING_SNAKE_CASE_       )       )
					lowercase_       =					self.normalization(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					self.activation(SCREAMING_SNAKE_CASE_       )
					return hidden_state

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



	"""simple docstring"""




	def __init__(       self  :			str       ,							SCREAMING_SNAKE_CASE_  :			RegNetConfig       ,							**SCREAMING_SNAKE_CASE_  :			str       )       ->     Any:
					super().__init__(**SCREAMING_SNAKE_CASE_       )
					lowercase_       =					config.num_channels
					lowercase_       =					TFRegNetConvLayer(
					    out_channels=config.embedding_size       ,							kernel_size=3       ,							stride=2       ,							activation=config.hidden_act       ,							name='''embedder'''       ,							)




	def 						_lowercase  (       self  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			str       )       ->     Optional[int]:
					lowercase_       =					shape_list(SCREAMING_SNAKE_CASE_       )[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)
					lowercase_       =					tf.transpose(SCREAMING_SNAKE_CASE_       ,							perm=(0, 2, 3, 1)       )
					lowercase_       =					self.embedder(SCREAMING_SNAKE_CASE_       )
					return hidden_state

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



	"""simple docstring"""




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




	def 						_lowercase  (       self  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			tf.Tensor       ,							SCREAMING_SNAKE_CASE_  :			bool = False       )       ->     tf.Tensor:
					return self.normalization(self.convolution(SCREAMING_SNAKE_CASE_       )       ,							training=SCREAMING_SNAKE_CASE_       )

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



	"""simple docstring"""




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




	def 						_lowercase  (       self  :			str       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       )       ->     Dict:
					# [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
					lowercase_       =					self.pooler(SCREAMING_SNAKE_CASE_       )
					for layer_module in self.attention:
									lowercase_       =					layer_module(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					hidden_state * pooled
					return hidden_state

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



	"""simple docstring"""




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




	def 						_lowercase  (       self  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       )       ->     Any:
					lowercase_       =					hidden_state
					for layer_module in self.layers:
									lowercase_       =					layer_module(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					self.shortcut(SCREAMING_SNAKE_CASE_       )
					hidden_state += residual
					lowercase_       =					self.activation(SCREAMING_SNAKE_CASE_       )
					return hidden_state

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



	"""simple docstring"""




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




	def 						_lowercase  (       self  :			int       ,							SCREAMING_SNAKE_CASE_  :			Dict       )       ->     Optional[Any]:
					lowercase_       =					hidden_state
					for layer_module in self.layers:
									lowercase_       =					layer_module(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					self.shortcut(SCREAMING_SNAKE_CASE_       )
					hidden_state += residual
					lowercase_       =					self.activation(SCREAMING_SNAKE_CASE_       )
					return hidden_state

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



	"""simple docstring"""




	def __init__(       self  :			str       ,							SCREAMING_SNAKE_CASE_  :			RegNetConfig       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			int = 2       ,							SCREAMING_SNAKE_CASE_  :			int = 2       ,							**SCREAMING_SNAKE_CASE_  :			List[Any]       )       ->     List[str]:
					super().__init__(**SCREAMING_SNAKE_CASE_       )

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




	def 						_lowercase  (       self  :			Tuple       ,							SCREAMING_SNAKE_CASE_  :			int       )       ->     int:
					for layer_module in self.layers:
									lowercase_       =					layer_module(SCREAMING_SNAKE_CASE_       )
					return hidden_state

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



	"""simple docstring"""




	def __init__(       self  :			Union[str, Any]       ,							SCREAMING_SNAKE_CASE_  :			RegNetConfig       ,							**SCREAMING_SNAKE_CASE_  :			Dict       )       ->     Tuple:
					super().__init__(**SCREAMING_SNAKE_CASE_       )
					lowercase_       =					[]
					# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
					self.stages.append(
					    TFRegNetStage(
					        SCREAMING_SNAKE_CASE_       ,							config.embedding_size       ,							config.hidden_sizes[0]       ,							stride=2 if config.downsample_in_first_stage else 1       ,							depth=config.depths[0]       ,							name='''stages.0'''       ,							)       )
					lowercase_       =					zip(config.hidden_sizes       ,							config.hidden_sizes[1:]       )
					for i, ((in_channels, out_channels), depth) in enumerate(zip(SCREAMING_SNAKE_CASE_       ,							config.depths[1:]       )       ):
									self.stages.append(TFRegNetStage(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							depth=SCREAMING_SNAKE_CASE_       ,							name=f'''stages.{i+1}'''       )       )




	def 						_lowercase  (       self  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			tf.Tensor       ,							SCREAMING_SNAKE_CASE_  :			bool = False       ,							SCREAMING_SNAKE_CASE_  :			bool = True       )       ->     TFBaseModelOutputWithNoAttention:
					lowercase_       =					() if output_hidden_states else None

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

									lowercase_       =					stage_module(SCREAMING_SNAKE_CASE_       )

					if output_hidden_states:
									lowercase_       =					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=SCREAMING_SNAKE_CASE_       ,							hidden_states=SCREAMING_SNAKE_CASE_       )

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



	"""simple docstring"""




	a			:str         =	RegNetConfig
	def __init__(       self  :			str       ,							SCREAMING_SNAKE_CASE_  :			Dict       ,							**SCREAMING_SNAKE_CASE_  :			List[str]       )       ->     Any:
					super().__init__(**SCREAMING_SNAKE_CASE_       )
					lowercase_       =					config
					lowercase_       =					TFRegNetEmbeddings(SCREAMING_SNAKE_CASE_       ,							name='''embedder'''       )
					lowercase_       =					TFRegNetEncoder(SCREAMING_SNAKE_CASE_       ,							name='''encoder'''       )
					lowercase_       =					tf.keras.layers.GlobalAveragePoolingaD(keepdims=SCREAMING_SNAKE_CASE_       ,							name='''pooler'''       )




	@unpack_inputs
	def 						_lowercase  (       self  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			tf.Tensor       ,							SCREAMING_SNAKE_CASE_  :			Optional[bool] = None       ,							SCREAMING_SNAKE_CASE_  :			Optional[bool] = None       ,							SCREAMING_SNAKE_CASE_  :			bool = False       ,							)       ->     TFBaseModelOutputWithPoolingAndNoAttention:
					lowercase_       =					(
					    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
					)
					lowercase_       =					return_dict if return_dict is not None else self.config.use_return_dict

					lowercase_       =					self.embedder(SCREAMING_SNAKE_CASE_       ,							training=SCREAMING_SNAKE_CASE_       )

					lowercase_       =					self.encoder(
					    SCREAMING_SNAKE_CASE_       ,							output_hidden_states=SCREAMING_SNAKE_CASE_       ,							return_dict=SCREAMING_SNAKE_CASE_       ,							training=SCREAMING_SNAKE_CASE_       )

					lowercase_       =					encoder_outputs[0]
					lowercase_       =					self.pooler(SCREAMING_SNAKE_CASE_       )

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

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

class  lowercase__(				UpperCAmelCase							):



	"""simple docstring"""




	a			:Tuple         =	RegNetConfig
	a			:Any         =	'regnet'
	a			:List[str]         =	'pixel_values'
	@property
	def 						_lowercase  (       self  :			List[str]       )       ->     str:
					return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_2_4, 2_2_4)       ,							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.' ,						UpperCAmelCase ,						)
class  lowercase__(				UpperCAmelCase							):



	"""simple docstring"""




	def __init__(       self  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			RegNetConfig       ,							*SCREAMING_SNAKE_CASE_  :			Any       ,							**SCREAMING_SNAKE_CASE_  :			str       )       ->     List[str]:
					super().__init__(SCREAMING_SNAKE_CASE_       ,							*SCREAMING_SNAKE_CASE_       ,							**SCREAMING_SNAKE_CASE_       )
					lowercase_       =					TFRegNetMainLayer(SCREAMING_SNAKE_CASE_       ,							name='''regnet'''       )




	@unpack_inputs
	@add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE_       )
	@add_code_sample_docstrings(
	    checkpoint=_CHECKPOINT_FOR_DOC       ,							output_type=SCREAMING_SNAKE_CASE_       ,							config_class=_CONFIG_FOR_DOC       ,							modality='''vision'''       ,							expected_output=_EXPECTED_OUTPUT_SHAPE       ,							)
	def 						_lowercase  (       self  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			tf.Tensor       ,							SCREAMING_SNAKE_CASE_  :			Optional[bool] = None       ,							SCREAMING_SNAKE_CASE_  :			Optional[bool] = None       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]=False       ,							)       ->     Union[TFBaseModelOutputWithPoolingAndNoAttention, Tuple[tf.Tensor]]:
					lowercase_       =					(
					    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
					)
					lowercase_       =					return_dict if return_dict is not None else self.config.use_return_dict

					lowercase_       =					self.regnet(
					    pixel_values=SCREAMING_SNAKE_CASE_       ,							output_hidden_states=SCREAMING_SNAKE_CASE_       ,							return_dict=SCREAMING_SNAKE_CASE_       ,							training=SCREAMING_SNAKE_CASE_       ,							)
					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    ' ,						UpperCAmelCase ,						)
class  lowercase__(				UpperCAmelCase ,						UpperCAmelCase							):



	"""simple docstring"""




	def __init__(       self  :			Union[str, Any]       ,							SCREAMING_SNAKE_CASE_  :			RegNetConfig       ,							*SCREAMING_SNAKE_CASE_  :			str       ,							**SCREAMING_SNAKE_CASE_  :			int       )       ->     Union[str, Any]:
					super().__init__(SCREAMING_SNAKE_CASE_       ,							*SCREAMING_SNAKE_CASE_       ,							**SCREAMING_SNAKE_CASE_       )
					lowercase_       =					config.num_labels
					lowercase_       =					TFRegNetMainLayer(SCREAMING_SNAKE_CASE_       ,							name='''regnet'''       )
					# classification head
					lowercase_       =					[
					    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(SCREAMING_SNAKE_CASE_       )
	@add_code_sample_docstrings(
	    checkpoint=_IMAGE_CLASS_CHECKPOINT       ,							output_type=SCREAMING_SNAKE_CASE_       ,							config_class=_CONFIG_FOR_DOC       ,							expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT       ,							)
	def 						_lowercase  (       self  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			tf.Tensor = None       ,							SCREAMING_SNAKE_CASE_  :			tf.Tensor = None       ,							SCREAMING_SNAKE_CASE_  :			bool = None       ,							SCREAMING_SNAKE_CASE_  :			bool = None       ,							SCREAMING_SNAKE_CASE_  :			Dict=False       ,							)       ->     Union[TFSequenceClassifierOutput, Tuple[tf.Tensor]]:
					lowercase_       =					(
					    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
					)
					lowercase_       =					return_dict if return_dict is not None else self.config.use_return_dict

					lowercase_       =					self.regnet(
					    SCREAMING_SNAKE_CASE_       ,							output_hidden_states=SCREAMING_SNAKE_CASE_       ,							return_dict=SCREAMING_SNAKE_CASE_       ,							training=SCREAMING_SNAKE_CASE_       )

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

					lowercase_       =					self.classifier[0](SCREAMING_SNAKE_CASE_       )
					lowercase_       =					self.classifier[1](SCREAMING_SNAKE_CASE_       )

					lowercase_       =					None if labels is None else self.hf_compute_loss(labels=SCREAMING_SNAKE_CASE_       ,							logits=SCREAMING_SNAKE_CASE_       )

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

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









import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


SCREAMING_SNAKE_CASE_:List[Any]             =						"""▁"""

SCREAMING_SNAKE_CASE_:int             =						{"""vocab_file""": """spiece.model"""}

SCREAMING_SNAKE_CASE_:Any             =						{
    """vocab_file""": {"""google/pegasus-xsum""": """https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model"""}
}

SCREAMING_SNAKE_CASE_:Optional[Any]             =						{
    """google/pegasus-xsum""": 512,
}


SCREAMING_SNAKE_CASE_:Tuple             =						logging.get_logger(__name__)






class SCREAMING_SNAKE_CASE__     (		SCREAMING_SNAKE_CASE__    ):




    '''simple docstring'''






    __lowerCamelCase			: Union[str, Any]					=      VOCAB_FILES_NAMES

    __lowerCamelCase			: List[Any]					=      VOCAB_FILES_NAMES
    __lowerCamelCase			: Optional[Any]					=      PRETRAINED_VOCAB_FILES_MAP
    __lowerCamelCase			: str					=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    __lowerCamelCase			: str					=      ["input_ids", "attention_mask"]






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__="<pad>",		lowerCamelCase__="</s>",		lowerCamelCase__="<unk>",		lowerCamelCase__="<mask_2>",		lowerCamelCase__="<mask_1>",		lowerCamelCase__=None,		lowerCamelCase__=103,		lowerCamelCase__ = None,		**lowerCamelCase__,		):
        A :       int          =   offset
        if additional_special_tokens is not None:
            if not isinstance(lowerCamelCase__,		lowerCamelCase__		):
                raise TypeError(
                    f'''additional_special_tokens should be of type {type(lowerCamelCase__		)}, but is'''
                    f''' {type(lowerCamelCase__		)}'''		)

            A :       Dict          =   (
                ([mask_token_sent] + additional_special_tokens)
                if mask_token_sent not in additional_special_tokens and mask_token_sent is not None
                else additional_special_tokens
            )
            # fill additional tokens with ..., <unk_token_102> in case not all additional tokens are already taken
            additional_special_tokens_extended += [
                f'''<unk_{i}>''' for i in range(len(lowerCamelCase__		),		self.offset - 1		)
            ]

            if len(set(lowerCamelCase__		)		) != len(lowerCamelCase__		):
                raise ValueError(
                    """Please make sure that the provided additional_special_tokens do not contain an incorrectly"""
                    f''' shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}.'''		)
            A :       int          =   additional_special_tokens_extended
        else:
            A :       Optional[int]          =   [mask_token_sent] if mask_token_sent is not None else []
            additional_special_tokens += [f'''<unk_{i}>''' for i in range(2,		self.offset		)]

        A :       Any          =   {} if sp_model_kwargs is None else sp_model_kwargs

        super().__init__(
            eos_token=lowerCamelCase__,		unk_token=lowerCamelCase__,		mask_token=lowerCamelCase__,		pad_token=lowerCamelCase__,		mask_token_sent=lowerCamelCase__,		offset=lowerCamelCase__,		additional_special_tokens=lowerCamelCase__,		sp_model_kwargs=self.sp_model_kwargs,		**lowerCamelCase__,		)
        A :       Union[str, Any]          =   mask_token_sent
        A :       Optional[Any]          =   vocab_file
        A :       Optional[int]          =   spm.SentencePieceProcessor(**self.sp_model_kwargs		)
        self.sp_model.Load(lowerCamelCase__		)

        # add special tokens to encoder dict
        A :       Dict[int, str]          =   {
            0: self.pad_token,
            1: self.eos_token,
        }

        if self.mask_token_sent is not None:
            self.encoder.update(
                {
                    2: self.mask_token_sent,
                    3: self.mask_token,
                }		)

        if self.offset > 0:
            # entries 2-104 are only used for pretraining and called <mask_1>, <mask_2>, unk_2, ...unk_102
            # mask_token_sent is already added to list -> so start at 1
            self.encoder.update({i + 3: additional_special_tokens[i] for i in range(1,		self.offset - 1		)}		)

        A :       Dict[str, int]          =   {v: k for k, v in self.encoder.items()}






    @property
    def 			_lowerCAmelCase					(       self		):
        return len(self.sp_model		) + self.offset






    def 			_lowerCAmelCase					(       self		):
        A :       Optional[int]          =   {self.convert_ids_to_tokens(lowerCamelCase__		): i for i in range(self.vocab_size		)}
        vocab.update(self.added_tokens_encoder		)
        return vocab






    def __getstate__(       self		):
        A :       List[Any]          =   self.__dict__.copy()
        A :       Union[str, Any]          =   None
        return state






    def __setstate__(       self,		lowerCamelCase__		):
        A :       List[Any]          =   d

        # for backward compatibility
        if not hasattr(self,		"""sp_model_kwargs"""		):
            A :       int          =   {}

        A :       Tuple          =   spm.SentencePieceProcessor(**self.sp_model_kwargs		)
        self.sp_model.Load(self.vocab_file		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        return self.sp_model.encode(lowerCamelCase__,		out_type=lowerCamelCase__		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        if token in self.decoder:
            return self.decoder[token]
        elif token in self.added_tokens_decoder:
            return self.added_tokens_decoder[token]
        A :       List[str]          =   self.sp_model.piece_to_id(lowerCamelCase__		)
        return sp_id + self.offset






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        if index in self.encoder:
            return self.encoder[index]
        elif index in self.added_tokens_encoder:
            return self.added_tokens_encoder[index]
        else:
            A :       Dict          =   self.sp_model.IdToPiece(index - self.offset		)
        return token






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        A :       Optional[int]          =   []
        A :       Optional[Any]          =   """"""
        for token in tokens:
            # make sure that special tokens are not decoded using sentencepiece model
            if token in self.all_special_tokens:
                out_string += self.sp_model.decode(lowerCamelCase__		) + token
                A :       int          =   []
            else:
                current_sub_tokens.append(lowerCamelCase__		)
        out_string += self.sp_model.decode(lowerCamelCase__		)
        return out_string.strip()






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__=False		):
        return 1






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        A :       Optional[int]          =   set(self.all_special_ids		)  # call it once instead of inside list comp
        all_special_ids.remove(self.unk_token_id		)  # <unk> is only sometimes special

        return [1 if x in all_special_ids else 0 for x in seq]






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__ = None,		lowerCamelCase__ = False		):
        if already_has_special_tokens:
            return self._special_token_mask(lowerCamelCase__		)
        elif token_ids_a is None:
            return self._special_token_mask(lowerCamelCase__		) + [1]
        else:
            return self._special_token_mask(token_ids_a + token_ids_a		) + [1]






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__=None		):
        if token_ids_a is None:
            return token_ids_a + [self.eos_token_id]
        # We don't expect to process pairs, but leave the pair logic for API consistency
        return token_ids_a + token_ids_a + [self.eos_token_id]






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__ = None		):
        if not os.path.isdir(lowerCamelCase__		):
            logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''		)
            return
        A :       Any          =   os.path.join(
            lowerCamelCase__,		(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]		)

        if os.path.abspath(self.vocab_file		) != os.path.abspath(lowerCamelCase__		) and os.path.isfile(self.vocab_file		):
            copyfile(self.vocab_file,		lowerCamelCase__		)
        elif not os.path.isfile(self.vocab_file		):
            with open(lowerCamelCase__,		"""wb"""		) as fi:
                A :       Dict          =   self.sp_model.serialized_model_proto()
                fi.write(lowerCamelCase__		)

        return (out_vocab_file,)


