"""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


__lowerCamelCase											=  logging.get_logger(__name__)

# General docstring
__lowerCamelCase											=  "RegNetConfig"

# Base docstring
__lowerCamelCase											=  "facebook/regnet-y-040"
__lowerCamelCase											=  [1, 10_88, 7, 7]

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

__lowerCamelCase											=  [
    "facebook/regnet-y-040",
    # See all regnet models at https://huggingface.co/models?filter=regnet
]
class     UpperCamelCase__(			tf.keras.layers.Layer ):





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





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				Optional[Any]:
													A__	     =							self.convolution(self.padding(__UpperCAmelCase     )     )
													A__	     =							self.normalization(__UpperCAmelCase     )
													A__	     =							self.activation(__UpperCAmelCase     )
													return hidden_state






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





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





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				Any:
													A__	     =							shape_list(__UpperCAmelCase     )[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)
													A__	     =							tf.transpose(__UpperCAmelCase		,perm=(0, 2, 3, 1)     )
													A__	     =							self.embedder(__UpperCAmelCase     )
													return hidden_state






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





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





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase = False     )					->				tf.Tensor:
													return self.normalization(self.convolution(__UpperCAmelCase     )		,training=__UpperCAmelCase     )






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





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





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				int:
													# [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
													A__	     =							self.pooler(__UpperCAmelCase     )
													for layer_module in self.attention:
																				A__	     =							layer_module(__UpperCAmelCase     )
													A__	     =							hidden_state * pooled
													return hidden_state






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





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





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				List[Any]:
													A__	     =							hidden_state
													for layer_module in self.layers:
																				A__	     =							layer_module(__UpperCAmelCase     )
													A__	     =							self.shortcut(__UpperCAmelCase     )
													hidden_state += residual
													A__	     =							self.activation(__UpperCAmelCase     )
													return hidden_state






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





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





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				Tuple:
													A__	     =							hidden_state
													for layer_module in self.layers:
																				A__	     =							layer_module(__UpperCAmelCase     )
													A__	     =							self.shortcut(__UpperCAmelCase     )
													hidden_state += residual
													A__	     =							self.activation(__UpperCAmelCase     )
													return hidden_state






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





						def __init__(       self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase = 2		,__UpperCAmelCase = 2		,**__UpperCAmelCase     )					->				Dict:
													super().__init__(**__UpperCAmelCase     )

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





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				Dict:
													for layer_module in self.layers:
																				A__	     =							layer_module(__UpperCAmelCase     )
													return hidden_state






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





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





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase = False		,__UpperCAmelCase = True     )					->				TFBaseModelOutputWithNoAttention:
													A__	     =							() if output_hidden_states else None

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

																				A__	     =							stage_module(__UpperCAmelCase     )

													if output_hidden_states:
																				A__	     =							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=__UpperCAmelCase		,hidden_states=__UpperCAmelCase     )






@keras_serializable
class     UpperCamelCase__(			tf.keras.layers.Layer ):
						lowerCAmelCase__  :						List[str]													=		RegNetConfig





						def __init__(       self		,__UpperCAmelCase		,**__UpperCAmelCase     )					->				Optional[int]:
													super().__init__(**__UpperCAmelCase     )
													A__	     =							config
													A__	     =							TFRegNetEmbeddings(__UpperCAmelCase		,name='embedder'     )
													A__	     =							TFRegNetEncoder(__UpperCAmelCase		,name='encoder'     )
													A__	     =							tf.keras.layers.GlobalAveragePoolingaD(keepdims=__UpperCAmelCase		,name='pooler'     )





						@unpack_inputs
						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase = None		,__UpperCAmelCase = None		,__UpperCAmelCase = False		,)					->				TFBaseModelOutputWithPoolingAndNoAttention:
													A__	     =							(
													    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
													)
													A__	     =							return_dict if return_dict is not None else self.config.use_return_dict

													A__	     =							self.embedder(__UpperCAmelCase		,training=__UpperCAmelCase     )

													A__	     =							self.encoder(
													    __UpperCAmelCase		,output_hidden_states=__UpperCAmelCase		,return_dict=__UpperCAmelCase		,training=__UpperCAmelCase     )

													A__	     =							encoder_outputs[0]
													A__	     =							self.pooler(__UpperCAmelCase     )

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

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






class     UpperCamelCase__(			__A ):
						lowerCAmelCase__  :						Any													=		RegNetConfig
						lowerCAmelCase__  :						List[str]													=		'regnet'
						lowerCAmelCase__  :						Union[str, Any]													=		'pixel_values'





						@property
						def 				snake_case__	(       self     )					->				str:
													return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 2_24, 2_24)		,dtype=tf.floataa     )}






__lowerCamelCase											=  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"

__lowerCamelCase											=  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.'		,						__A		,						)
class     UpperCamelCase__(			__A ):





						def __init__(       self		,__UpperCAmelCase		,*__UpperCAmelCase		,**__UpperCAmelCase     )					->				int:
													super().__init__(__UpperCAmelCase		,*__UpperCAmelCase		,**__UpperCAmelCase     )
													A__	     =							TFRegNetMainLayer(__UpperCAmelCase		,name='regnet'     )





						@unpack_inputs
						@add_start_docstrings_to_model_forward(__UpperCAmelCase     )
						@add_code_sample_docstrings(
						    checkpoint=_CHECKPOINT_FOR_DOC		,output_type=__UpperCAmelCase		,config_class=_CONFIG_FOR_DOC		,modality='vision'		,expected_output=_EXPECTED_OUTPUT_SHAPE		,)
						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase = None		,__UpperCAmelCase = None		,__UpperCAmelCase=False		,)					->				Union[TFBaseModelOutputWithPoolingAndNoAttention, Tuple[tf.Tensor]]:
													A__	     =							(
													    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
													)
													A__	     =							return_dict if return_dict is not None else self.config.use_return_dict

													A__	     =							self.regnet(
													    pixel_values=__UpperCAmelCase		,output_hidden_states=__UpperCAmelCase		,return_dict=__UpperCAmelCase		,training=__UpperCAmelCase		,)
													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    '		,						__A		,						)
class     UpperCamelCase__(			__A		,						__A ):





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

													A__	     =							self.regnet(
													    __UpperCAmelCase		,output_hidden_states=__UpperCAmelCase		,return_dict=__UpperCAmelCase		,training=__UpperCAmelCase     )

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

													A__	     =							self.classifier[0](__UpperCAmelCase     )
													A__	     =							self.classifier[1](__UpperCAmelCase     )

													A__	     =							None if labels is None else self.hf_compute_loss(labels=__UpperCAmelCase		,logits=__UpperCAmelCase     )

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

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






"""simple docstring"""




import importlib
import shutil
import threading
import warnings
from typing import List

import fsspec
import fsspec.asyn

from . import compression
from .hffilesystem import HfFileSystem


__lowerCamelCase											=  importlib.util.find_spec("s3fs") is not None

if _has_safs:
				from .safilesystem import SaFileSystem  # noqa: F401

__lowerCamelCase											=  [
    compression.BzaFileSystem,
    compression.GzipFileSystem,
    compression.LzaFileSystem,
    compression.XzFileSystem,
    compression.ZstdFileSystem,
]

# Register custom filesystems
for fs_class in COMPRESSION_FILESYSTEMS + [HfFileSystem]:
				if fs_class.protocol in fsspec.registry and fsspec.registry[fs_class.protocol] is not fs_class:
								warnings.warn(F'''A filesystem protocol was already set for {fs_class.protocol} and will be overwritten.''')
				fsspec.register_implementation(fs_class.protocol, fs_class, clobber=True)






def   UpperCAmelCase					(	UpperCamelCase__				):






							"""simple docstring"""





							if "://" in dataset_path:
														A__	     =							dataset_path.split('://'				)[1]
							return dataset_path






def   UpperCAmelCase					(	UpperCamelCase__				):






							"""simple docstring"""





							if fs is not None and fs.protocol != "file":
														return True
							else:
														return False






def   UpperCAmelCase					(	UpperCamelCase__       ,    UpperCamelCase__       ,    UpperCamelCase__				):






							"""simple docstring"""





							A__	     =							not is_remote_filesystem(UpperCamelCase__				)
							if is_local:
														# LocalFileSystem.mv does copy + rm, it is more efficient to simply move a local directory
														shutil.move(fs._strip_protocol(UpperCamelCase__				)       ,    fs._strip_protocol(UpperCamelCase__				)				)
							else:
														fs.mv(UpperCamelCase__       ,    UpperCamelCase__       ,    recursive=UpperCamelCase__				)






def   UpperCAmelCase					(	):






							"""simple docstring"""





							if hasattr(fsspec.asyn       ,    'reset_lock'				):
														# for future fsspec>2022.05.0
														fsspec.asyn.reset_lock()
							else:
														A__	     =							None
														A__	     =							None
														A__	     =							threading.Lock()





