






'''simple docstring'''
from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_outputs import (
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import logging
from .configuration_regnet import RegNetConfig


a_ 		=       logging.get_logger(__name__)

# General docstring
a_ 		=       'RegNetConfig'

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

# Image classification docstring
a_ 		=       'facebook/regnet-y-040'
a_ 		=       'tabby, tabby cat'

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






class    __SCREAMING_SNAKE_CASE					(	nn.Module ):




 def __init__(					self		:							Optional[Any]    ,					__lowercase		:							int    ,					__lowercase		:							int    ,					__lowercase		:							int = 3    ,					__lowercase		:							int = 1    ,					__lowercase		:							int = 1    ,					__lowercase		:							Optional[str] = "relu"    ,					)		->  List[str]:
       super().__init__()
       SCREAMING_SNAKE_CASE__   : Optional[int]													=nn.Convad(
           _UpperCAmelCase    ,					_UpperCAmelCase    ,					kernel_size=_UpperCAmelCase    ,					stride=_UpperCAmelCase    ,					padding=kernel_size // 2    ,					groups=_UpperCAmelCase    ,					bias=_UpperCAmelCase    ,					)
       SCREAMING_SNAKE_CASE__   : Any													=nn.BatchNormad(_UpperCAmelCase					)
       SCREAMING_SNAKE_CASE__   : Optional[int]													=ACTaFN[activation] if activation is not None else nn.Identity()




 def    __magic_name__		(					self		:							int    ,					__lowercase		:							List[str]					)		->  List[str]:
       SCREAMING_SNAKE_CASE__   : Dict													=self.convolution(_UpperCAmelCase					)
       SCREAMING_SNAKE_CASE__   : List[str]													=self.normalization(_UpperCAmelCase					)
       SCREAMING_SNAKE_CASE__   : List[Any]													=self.activation(_UpperCAmelCase					)
       return hidden_state






class    __SCREAMING_SNAKE_CASE					(	nn.Module ):




 def __init__(					self		:							Any    ,					__lowercase		:							RegNetConfig					)		->  Optional[Any]:
       super().__init__()
       SCREAMING_SNAKE_CASE__   : List[Any]													=RegNetConvLayer(
           config.num_channels    ,					config.embedding_size    ,					kernel_size=3    ,					stride=2    ,					activation=config.hidden_act					)
       SCREAMING_SNAKE_CASE__   : List[str]													=config.num_channels




 def    __magic_name__		(					self		:							str    ,					__lowercase		:							Tuple					)		->  str:
       SCREAMING_SNAKE_CASE__   : Tuple													=pixel_values.shape[1]
       if 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.'''					)
       SCREAMING_SNAKE_CASE__   : Optional[int]													=self.embedder(_UpperCAmelCase					)
       return hidden_state






class    __SCREAMING_SNAKE_CASE					(	nn.Module ):




 def __init__(					self		:							List[Any]    ,					__lowercase		:							int    ,					__lowercase		:							int    ,					__lowercase		:							int = 2					)		->  List[str]:
       super().__init__()
       SCREAMING_SNAKE_CASE__   : str													=nn.Convad(_UpperCAmelCase    ,					_UpperCAmelCase    ,					kernel_size=1    ,					stride=_UpperCAmelCase    ,					bias=_UpperCAmelCase					)
       SCREAMING_SNAKE_CASE__   : Optional[Any]													=nn.BatchNormad(_UpperCAmelCase					)




 def    __magic_name__		(					self		:							Dict    ,					__lowercase		:							Tensor					)		->  Any:
       SCREAMING_SNAKE_CASE__   : int													=self.convolution(_UpperCAmelCase					)
       SCREAMING_SNAKE_CASE__   : Any													=self.normalization(_UpperCAmelCase					)
       return hidden_state






class    __SCREAMING_SNAKE_CASE					(	nn.Module ):




 def __init__(					self		:							Optional[int]    ,					__lowercase		:							int    ,					__lowercase		:							int					)		->  str:
       super().__init__()

       SCREAMING_SNAKE_CASE__   : Optional[int]													=nn.AdaptiveAvgPoolad((1, 1)					)
       SCREAMING_SNAKE_CASE__   : int													=nn.Sequential(
           nn.Convad(_UpperCAmelCase    ,					_UpperCAmelCase    ,					kernel_size=1					)    ,					nn.ReLU()    ,					nn.Convad(_UpperCAmelCase    ,					_UpperCAmelCase    ,					kernel_size=1					)    ,					nn.Sigmoid()    ,					)




 def    __magic_name__		(					self		:							Optional[Any]    ,					__lowercase		:							Optional[int]					)		->  Any:
       # b c h w -> b c 1 1
       SCREAMING_SNAKE_CASE__   : int													=self.pooler(_UpperCAmelCase					)
       SCREAMING_SNAKE_CASE__   : Dict													=self.attention(_UpperCAmelCase					)
       SCREAMING_SNAKE_CASE__   : List[str]													=hidden_state * attention
       return hidden_state






class    __SCREAMING_SNAKE_CASE					(	nn.Module ):




 def __init__(					self		:							List[Any]    ,					__lowercase		:							RegNetConfig    ,					__lowercase		:							int    ,					__lowercase		:							int    ,					__lowercase		:							int = 1					)		->  List[Any]:
       super().__init__()
       SCREAMING_SNAKE_CASE__   : str													=in_channels != out_channels or stride != 1
       SCREAMING_SNAKE_CASE__   : int													=max(1    ,					out_channels // config.groups_width					)
       SCREAMING_SNAKE_CASE__   : Tuple													=(
           RegNetShortCut(_UpperCAmelCase    ,					_UpperCAmelCase    ,					stride=_UpperCAmelCase					) if should_apply_shortcut else nn.Identity()
       )
       SCREAMING_SNAKE_CASE__   : Any													=nn.Sequential(
           RegNetConvLayer(_UpperCAmelCase    ,					_UpperCAmelCase    ,					kernel_size=1    ,					activation=config.hidden_act					)    ,					RegNetConvLayer(_UpperCAmelCase    ,					_UpperCAmelCase    ,					stride=_UpperCAmelCase    ,					groups=_UpperCAmelCase    ,					activation=config.hidden_act					)    ,					RegNetConvLayer(_UpperCAmelCase    ,					_UpperCAmelCase    ,					kernel_size=1    ,					activation=_UpperCAmelCase					)    ,					)
       SCREAMING_SNAKE_CASE__   : Any													=ACTaFN[config.hidden_act]




 def    __magic_name__		(					self		:							List[Any]    ,					__lowercase		:							str					)		->  int:
       SCREAMING_SNAKE_CASE__   : str													=hidden_state
       SCREAMING_SNAKE_CASE__   : str													=self.layer(_UpperCAmelCase					)
       SCREAMING_SNAKE_CASE__   : Tuple													=self.shortcut(_UpperCAmelCase					)
       hidden_state += residual
       SCREAMING_SNAKE_CASE__   : Optional[Any]													=self.activation(_UpperCAmelCase					)
       return hidden_state






class    __SCREAMING_SNAKE_CASE					(	nn.Module ):




 def __init__(					self		:							Any    ,					__lowercase		:							RegNetConfig    ,					__lowercase		:							int    ,					__lowercase		:							int    ,					__lowercase		:							int = 1					)		->  str:
       super().__init__()
       SCREAMING_SNAKE_CASE__   : Optional[int]													=in_channels != out_channels or stride != 1
       SCREAMING_SNAKE_CASE__   : str													=max(1    ,					out_channels // config.groups_width					)
       SCREAMING_SNAKE_CASE__   : Dict													=(
           RegNetShortCut(_UpperCAmelCase    ,					_UpperCAmelCase    ,					stride=_UpperCAmelCase					) if should_apply_shortcut else nn.Identity()
       )
       SCREAMING_SNAKE_CASE__   : int													=nn.Sequential(
           RegNetConvLayer(_UpperCAmelCase    ,					_UpperCAmelCase    ,					kernel_size=1    ,					activation=config.hidden_act					)    ,					RegNetConvLayer(_UpperCAmelCase    ,					_UpperCAmelCase    ,					stride=_UpperCAmelCase    ,					groups=_UpperCAmelCase    ,					activation=config.hidden_act					)    ,					RegNetSELayer(_UpperCAmelCase    ,					reduced_channels=int(round(in_channels / 4					)					)					)    ,					RegNetConvLayer(_UpperCAmelCase    ,					_UpperCAmelCase    ,					kernel_size=1    ,					activation=_UpperCAmelCase					)    ,					)
       SCREAMING_SNAKE_CASE__   : Optional[int]													=ACTaFN[config.hidden_act]




 def    __magic_name__		(					self		:							Dict    ,					__lowercase		:							Any					)		->  Union[str, Any]:
       SCREAMING_SNAKE_CASE__   : List[str]													=hidden_state
       SCREAMING_SNAKE_CASE__   : Tuple													=self.layer(_UpperCAmelCase					)
       SCREAMING_SNAKE_CASE__   : Dict													=self.shortcut(_UpperCAmelCase					)
       hidden_state += residual
       SCREAMING_SNAKE_CASE__   : List[str]													=self.activation(_UpperCAmelCase					)
       return hidden_state






class    __SCREAMING_SNAKE_CASE					(	nn.Module ):




 def __init__(					self		:							Any    ,					__lowercase		:							RegNetConfig    ,					__lowercase		:							int    ,					__lowercase		:							int    ,					__lowercase		:							int = 2    ,					__lowercase		:							int = 2    ,					)		->  Optional[int]:
       super().__init__()

       SCREAMING_SNAKE_CASE__   : Dict													=RegNetXLayer if config.layer_type == 'x' else RegNetYLayer

       SCREAMING_SNAKE_CASE__   : Any													=nn.Sequential(
           # downsampling is done in the first layer with stride of 2
           layer(
               _UpperCAmelCase    ,					_UpperCAmelCase    ,					_UpperCAmelCase    ,					stride=_UpperCAmelCase    ,					)    ,					*[layer(_UpperCAmelCase    ,					_UpperCAmelCase    ,					_UpperCAmelCase					) for _ in range(depth - 1					)]    ,					)




 def    __magic_name__		(					self		:							Union[str, Any]    ,					__lowercase		:							Optional[Any]					)		->  str:
       SCREAMING_SNAKE_CASE__   : Union[str, Any]													=self.layers(_UpperCAmelCase					)
       return hidden_state






class    __SCREAMING_SNAKE_CASE					(	nn.Module ):




 def __init__(					self		:							str    ,					__lowercase		:							RegNetConfig					)		->  int:
       super().__init__()
       SCREAMING_SNAKE_CASE__   : Dict													=nn.ModuleList([]					)
       # based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
       self.stages.append(
           RegNetStage(
               _UpperCAmelCase    ,					config.embedding_size    ,					config.hidden_sizes[0]    ,					stride=2 if config.downsample_in_first_stage else 1    ,					depth=config.depths[0]    ,					)					)
       SCREAMING_SNAKE_CASE__   : int													=zip(config.hidden_sizes    ,					config.hidden_sizes[1:]					)
       for (in_channels, out_channels), depth in zip(_UpperCAmelCase    ,					config.depths[1:]					):
             self.stages.append(RegNetStage(_UpperCAmelCase    ,					_UpperCAmelCase    ,					_UpperCAmelCase    ,					depth=_UpperCAmelCase					)					)




 def    __magic_name__		(					self		:							Union[str, Any]    ,					__lowercase		:							Tensor    ,					__lowercase		:							bool = False    ,					__lowercase		:							bool = True					)		->  List[str]:
       SCREAMING_SNAKE_CASE__   : Optional[int]													=() if output_hidden_states else None

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

             SCREAMING_SNAKE_CASE__   : Any													=stage_module(_UpperCAmelCase					)

       if output_hidden_states:
             SCREAMING_SNAKE_CASE__   : str													=hidden_states + (hidden_state,)

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

       return BaseModelOutputWithNoAttention(last_hidden_state=_UpperCAmelCase    ,					hidden_states=_UpperCAmelCase					)






class    __SCREAMING_SNAKE_CASE					(	_UpperCamelCase ):
 snake_case_     							=							RegNetConfig
 snake_case_     							=							'regnet'
 snake_case_     							=							'pixel_values'
 snake_case_     							=							True




 def    __magic_name__		(					self		:							Dict    ,					__lowercase		:							Dict					)		->  Any:
       if isinstance(_UpperCAmelCase    ,					nn.Convad					):
             nn.init.kaiming_normal_(module.weight    ,					mode='''fan_out'''    ,					nonlinearity='''relu'''					)
       elif isinstance(_UpperCAmelCase    ,					(nn.BatchNormad, nn.GroupNorm)					):
             nn.init.constant_(module.weight    ,					1					)
             nn.init.constant_(module.bias    ,					0					)




 def    __magic_name__		(					self		:							Union[str, Any]    ,					__lowercase		:							List[str]    ,					__lowercase		:							Dict=False					)		->  Union[str, Any]:
       if isinstance(_UpperCAmelCase    ,					_UpperCAmelCase					):
             SCREAMING_SNAKE_CASE__   : Union[str, Any]													=value


a_ 		=       R'\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\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 [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n'

a_ 		=       R'\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConvNextImageProcessor.__call__`] for details.\n\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 [`~file_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		, )
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class    __SCREAMING_SNAKE_CASE					(	_UpperCamelCase ):




 def __init__(					self		:							Optional[Any]    ,					__lowercase		:							Any					)		->  List[Any]:
       super().__init__(_UpperCAmelCase					)
       SCREAMING_SNAKE_CASE__   : List[Any]													=config
       SCREAMING_SNAKE_CASE__   : int													=RegNetEmbeddings(_UpperCAmelCase					)
       SCREAMING_SNAKE_CASE__   : str													=RegNetEncoder(_UpperCAmelCase					)
       SCREAMING_SNAKE_CASE__   : Optional[int]													=nn.AdaptiveAvgPoolad((1, 1)					)
       # Initialize weights and apply final processing
       self.post_init()




 @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    __magic_name__		(					self		:							List[str]    ,					__lowercase		:							Tensor    ,					__lowercase		:							Optional[bool] = None    ,					__lowercase		:							Optional[bool] = None					)		->  Dict:
       SCREAMING_SNAKE_CASE__   : Any													=(
           output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
       )
       SCREAMING_SNAKE_CASE__   : Optional[int]													=return_dict if return_dict is not None else self.config.use_return_dict

       SCREAMING_SNAKE_CASE__   : List[str]													=self.embedder(_UpperCAmelCase					)

       SCREAMING_SNAKE_CASE__   : Tuple													=self.encoder(
           _UpperCAmelCase    ,					output_hidden_states=_UpperCAmelCase    ,					return_dict=_UpperCAmelCase					)

       SCREAMING_SNAKE_CASE__   : List[Any]													=encoder_outputs[0]

       SCREAMING_SNAKE_CASE__   : int													=self.pooler(_UpperCAmelCase					)

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

       return BaseModelOutputWithPoolingAndNoAttention(
           last_hidden_state=_UpperCAmelCase    ,					pooler_output=_UpperCAmelCase    ,					hidden_states=encoder_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		, )
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class    __SCREAMING_SNAKE_CASE					(	_UpperCamelCase ):




 def __init__(					self		:							Union[str, Any]    ,					__lowercase		:							Dict					)		->  List[str]:
       super().__init__(_UpperCAmelCase					)
       SCREAMING_SNAKE_CASE__   : List[Any]													=config.num_labels
       SCREAMING_SNAKE_CASE__   : Any													=RegNetModel(_UpperCAmelCase					)
       # classification head
       SCREAMING_SNAKE_CASE__   : Union[str, Any]													=nn.Sequential(
           nn.Flatten()    ,					nn.Linear(config.hidden_sizes[-1]    ,					config.num_labels					) if config.num_labels > 0 else nn.Identity()    ,					)
       # initialize weights and apply final processing
       self.post_init()




 @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    __magic_name__		(					self		:							Dict    ,					__lowercase		:							Optional[torch.FloatTensor] = None    ,					__lowercase		:							Optional[torch.LongTensor] = None    ,					__lowercase		:							Optional[bool] = None    ,					__lowercase		:							Optional[bool] = None    ,					)		->  str:
       SCREAMING_SNAKE_CASE__   : Any													=return_dict if return_dict is not None else self.config.use_return_dict

       SCREAMING_SNAKE_CASE__   : Optional[int]													=self.regnet(_UpperCAmelCase    ,					output_hidden_states=_UpperCAmelCase    ,					return_dict=_UpperCAmelCase					)

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

       SCREAMING_SNAKE_CASE__   : List[Any]													=self.classifier(_UpperCAmelCase					)

       SCREAMING_SNAKE_CASE__   : Dict													=None

       if labels is not None:
             if self.config.problem_type is None:
                   if self.num_labels == 1:
                         SCREAMING_SNAKE_CASE__   : Tuple													='regression'
                   elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                         SCREAMING_SNAKE_CASE__   : Any													='single_label_classification'
                   else:
                         SCREAMING_SNAKE_CASE__   : Union[str, Any]													='multi_label_classification'
             if self.config.problem_type == "regression":
                   SCREAMING_SNAKE_CASE__   : List[str]													=MSELoss()
                   if self.num_labels == 1:
                         SCREAMING_SNAKE_CASE__   : int													=loss_fct(logits.squeeze()    ,					labels.squeeze()					)
                   else:
                         SCREAMING_SNAKE_CASE__   : Union[str, Any]													=loss_fct(_UpperCAmelCase    ,					_UpperCAmelCase					)
             elif self.config.problem_type == "single_label_classification":
                   SCREAMING_SNAKE_CASE__   : List[Any]													=CrossEntropyLoss()
                   SCREAMING_SNAKE_CASE__   : Tuple													=loss_fct(logits.view(-1    ,					self.num_labels					)    ,					labels.view(-1					)					)
             elif self.config.problem_type == "multi_label_classification":
                   SCREAMING_SNAKE_CASE__   : List[Any]													=BCEWithLogitsLoss()
                   SCREAMING_SNAKE_CASE__   : int													=loss_fct(_UpperCAmelCase    ,					_UpperCAmelCase					)

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

       return ImageClassifierOutputWithNoAttention(loss=_UpperCAmelCase    ,					logits=_UpperCAmelCase    ,					hidden_states=outputs.hidden_states					)




'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__lowerCAmelCase									=   {
    '''configuration_bigbird_pegasus''': [
        '''BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''BigBirdPegasusConfig''',
        '''BigBirdPegasusOnnxConfig''',
    ],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__lowerCAmelCase									=   [
							    '''BIGBIRD_PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST''',
							    '''BigBirdPegasusForCausalLM''',
							    '''BigBirdPegasusForConditionalGeneration''',
							    '''BigBirdPegasusForQuestionAnswering''',
							    '''BigBirdPegasusForSequenceClassification''',
							    '''BigBirdPegasusModel''',
							    '''BigBirdPegasusPreTrainedModel''',
							]


if TYPE_CHECKING:
							from .configuration_bigbird_pegasus import (
							    BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP,
							    BigBirdPegasusConfig,
							    BigBirdPegasusOnnxConfig,
							)

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_bigbird_pegasus import (
														    BIGBIRD_PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST,
														    BigBirdPegasusForCausalLM,
														    BigBirdPegasusForConditionalGeneration,
														    BigBirdPegasusForQuestionAnswering,
														    BigBirdPegasusForSequenceClassification,
														    BigBirdPegasusModel,
														    BigBirdPegasusPreTrainedModel,
														)


else:
							import sys

							__lowerCAmelCase									=   _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)




