'''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 ...modeling_outputs import (
    BackboneOutput,
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ...utils.backbone_utils import BackboneMixin
from .configuration_resnet import ResNetConfig


SCREAMING_SNAKE_CASE_	   =			logging.get_logger(__name__)

# General docstring
SCREAMING_SNAKE_CASE_	   =			'ResNetConfig'

# Base docstring
SCREAMING_SNAKE_CASE_	   =			'microsoft/resnet-50'
SCREAMING_SNAKE_CASE_	   =			[1, 20_48, 7, 7]

# Image classification docstring
SCREAMING_SNAKE_CASE_	   =			'microsoft/resnet-50'
SCREAMING_SNAKE_CASE_	   =			'tiger cat'

SCREAMING_SNAKE_CASE_	   =			[
    'microsoft/resnet-50',
    # See all resnet models at https://huggingface.co/models?filter=resnet
]




class    a (       nn.Module				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_ = 3   ,    snake_case_ = 1   ,    snake_case_ = "relu"					):

				'''simple docstring'''







				super().__init__()
				__UpperCAmelCase:					Union[str, Any]		= nn.Convad(
				    snake_case_   ,    snake_case_   ,    kernel_size=snake_case_   ,    stride=snake_case_   ,    padding=kernel_size // 2   ,    bias=snake_case_					)
				__UpperCAmelCase:					Any		= nn.BatchNormad(snake_case_					)
				__UpperCAmelCase:					Tuple		= ACTaFN[activation] if activation is not None else nn.Identity()







		def    lowercase_  (							self   ,    snake_case_					):

				'''simple docstring'''







				__UpperCAmelCase:					List[Any]		= self.convolution(snake_case_					)
				__UpperCAmelCase:					Dict		= self.normalization(snake_case_					)
				__UpperCAmelCase:					str		= self.activation(snake_case_					)
				return hidden_state




class    a (       nn.Module				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_					):

				'''simple docstring'''







				super().__init__()
				__UpperCAmelCase:					Any		= ResNetConvLayer(
				    config.num_channels   ,    config.embedding_size   ,    kernel_size=7   ,    stride=2   ,    activation=config.hidden_act					)
				__UpperCAmelCase:					List[Any]		= nn.MaxPoolad(kernel_size=3   ,    stride=2   ,    padding=1					)
				__UpperCAmelCase:					int		= config.num_channels







		def    lowercase_  (							self   ,    snake_case_					):

				'''simple docstring'''







				__UpperCAmelCase:					Any		= 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."""					)
				__UpperCAmelCase:					Optional[Any]		= self.embedder(snake_case_					)
				__UpperCAmelCase:					Dict		= self.pooler(snake_case_					)
				return embedding




class    a (       nn.Module				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_ = 2					):

				'''simple docstring'''







				super().__init__()
				__UpperCAmelCase:					Optional[int]		= nn.Convad(snake_case_   ,    snake_case_   ,    kernel_size=1   ,    stride=snake_case_   ,    bias=snake_case_					)
				__UpperCAmelCase:					Union[str, Any]		= nn.BatchNormad(snake_case_					)







		def    lowercase_  (							self   ,    snake_case_					):

				'''simple docstring'''







				__UpperCAmelCase:					str		= self.convolution(snake_case_					)
				__UpperCAmelCase:					Tuple		= self.normalization(snake_case_					)
				return hidden_state




class    a (       nn.Module				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_ = 1   ,    snake_case_ = "relu"					):

				'''simple docstring'''







				super().__init__()
				__UpperCAmelCase:					List[Any]		= in_channels != out_channels or stride != 1
				__UpperCAmelCase:					Optional[Any]		= (
				    ResNetShortCut(snake_case_   ,    snake_case_   ,    stride=snake_case_					) if should_apply_shortcut else nn.Identity()
				)
				__UpperCAmelCase:					Tuple		= nn.Sequential(
				    ResNetConvLayer(snake_case_   ,    snake_case_   ,    stride=snake_case_					)   ,    ResNetConvLayer(snake_case_   ,    snake_case_   ,    activation=snake_case_					)   ,    )
				__UpperCAmelCase:					List[Any]		= ACTaFN[activation]







		def    lowercase_  (							self   ,    snake_case_					):

				'''simple docstring'''







				__UpperCAmelCase:					List[str]		= hidden_state
				__UpperCAmelCase:					Tuple		= self.layer(snake_case_					)
				__UpperCAmelCase:					List[Any]		= self.shortcut(snake_case_					)
				hidden_state += residual
				__UpperCAmelCase:					List[str]		= self.activation(snake_case_					)
				return hidden_state




class    a (       nn.Module				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_ = 1   ,    snake_case_ = "relu"   ,    snake_case_ = 4					):

				'''simple docstring'''







				super().__init__()
				__UpperCAmelCase:					Optional[int]		= in_channels != out_channels or stride != 1
				__UpperCAmelCase:					Union[str, Any]		= out_channels // reduction
				__UpperCAmelCase:					Optional[Any]		= (
				    ResNetShortCut(snake_case_   ,    snake_case_   ,    stride=snake_case_					) if should_apply_shortcut else nn.Identity()
				)
				__UpperCAmelCase:					int		= nn.Sequential(
				    ResNetConvLayer(snake_case_   ,    snake_case_   ,    kernel_size=1					)   ,    ResNetConvLayer(snake_case_   ,    snake_case_   ,    stride=snake_case_					)   ,    ResNetConvLayer(snake_case_   ,    snake_case_   ,    kernel_size=1   ,    activation=snake_case_					)   ,    )
				__UpperCAmelCase:					Dict		= ACTaFN[activation]







		def    lowercase_  (							self   ,    snake_case_					):

				'''simple docstring'''







				__UpperCAmelCase:					Tuple		= hidden_state
				__UpperCAmelCase:					str		= self.layer(snake_case_					)
				__UpperCAmelCase:					str		= self.shortcut(snake_case_					)
				hidden_state += residual
				__UpperCAmelCase:					Union[str, Any]		= self.activation(snake_case_					)
				return hidden_state




class    a (       nn.Module				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_ = 2   ,    snake_case_ = 2   ,    ):

				'''simple docstring'''







				super().__init__()

				__UpperCAmelCase:					Any		= ResNetBottleNeckLayer if config.layer_type == """bottleneck""" else ResNetBasicLayer

				__UpperCAmelCase:					List[str]		= nn.Sequential(
				    # downsampling is done in the first layer with stride of 2
				    layer(snake_case_   ,    snake_case_   ,    stride=snake_case_   ,    activation=config.hidden_act					)   ,    *[layer(snake_case_   ,    snake_case_   ,    activation=config.hidden_act					) for _ in range(depth - 1					)]   ,    )







		def    lowercase_  (							self   ,    snake_case_					):

				'''simple docstring'''







				__UpperCAmelCase:					List[Any]		= input
				for layer in self.layers:
						__UpperCAmelCase:					Tuple		= layer(snake_case_					)
				return hidden_state




class    a (       nn.Module				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_					):

				'''simple docstring'''







				super().__init__()
				__UpperCAmelCase:					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(
				    ResNetStage(
				        snake_case_   ,    config.embedding_size   ,    config.hidden_sizes[0]   ,    stride=2 if config.downsample_in_first_stage else 1   ,    depth=config.depths[0]   ,    )					)
				__UpperCAmelCase:					Any		= zip(config.hidden_sizes   ,    config.hidden_sizes[1:]					)
				for (in_channels, out_channels), depth in zip(snake_case_   ,    config.depths[1:]					):
						self.stages.append(ResNetStage(snake_case_   ,    snake_case_   ,    snake_case_   ,    depth=snake_case_					)					)







		def    lowercase_  (							self   ,    snake_case_   ,    snake_case_ = False   ,    snake_case_ = True					):

				'''simple docstring'''







				__UpperCAmelCase:					Any		= () if output_hidden_states else None

				for stage_module in self.stages:
						if output_hidden_states:
								__UpperCAmelCase:					Optional[int]		= hidden_states + (hidden_state,)

						__UpperCAmelCase:					Union[str, Any]		= stage_module(snake_case_					)

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




class    a (       __lowerCAmelCase				):



		"""simple docstring"""
		__lowerCAmelCase				      =					ResNetConfig
		__lowerCAmelCase				      =					"""resnet"""
		__lowerCAmelCase				      =					"""pixel_values"""
		__lowerCAmelCase				      =					True







		def    lowercase_  (							self   ,    snake_case_					):

				'''simple docstring'''







				if isinstance(snake_case_   ,    nn.Convad					):
						nn.init.kaiming_normal_(module.weight   ,    mode="""fan_out"""   ,    nonlinearity="""relu"""					)
				elif isinstance(snake_case_   ,    (nn.BatchNormad, nn.GroupNorm)					):
						nn.init.constant_(module.weight   ,    1					)
						nn.init.constant_(module.bias   ,    0					)







		def    lowercase_  (							self   ,    snake_case_   ,    snake_case_=False					):

				'''simple docstring'''







				if isinstance(snake_case_   ,    snake_case_					):
						__UpperCAmelCase:					int		= value



SCREAMING_SNAKE_CASE_	   =			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 ([`ResNetConfig`]): 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'

SCREAMING_SNAKE_CASE_	   =			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 [`~utils.ModelOutput`] instead of a plain tuple.\n'




@add_start_docstrings(
    """The bare ResNet model outputting raw features without any specific head on top."""						,  __lowerCAmelCase						,  )
class    a (       __lowerCAmelCase				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_					):

				'''simple docstring'''







				super().__init__(snake_case_					)
				__UpperCAmelCase:					List[str]		= config
				__UpperCAmelCase:					List[Any]		= ResNetEmbeddings(snake_case_					)
				__UpperCAmelCase:					Any		= ResNetEncoder(snake_case_					)
				__UpperCAmelCase:					Optional[Any]		= nn.AdaptiveAvgPoolad((1, 1)					)
				# Initialize weights and apply final processing
				self.post_init()







		@add_start_docstrings_to_model_forward(snake_case_					)
		@add_code_sample_docstrings(
		    checkpoint=_CHECKPOINT_FOR_DOC   ,    output_type=snake_case_   ,    config_class=_CONFIG_FOR_DOC   ,    modality="""vision"""   ,    expected_output=_EXPECTED_OUTPUT_SHAPE   ,    )
		def    lowercase_  (							self   ,    snake_case_   ,    snake_case_ = None   ,    snake_case_ = None					):

				'''simple docstring'''







				__UpperCAmelCase:					Optional[Any]		= (
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				__UpperCAmelCase:					Tuple		= return_dict if return_dict is not None else self.config.use_return_dict

				__UpperCAmelCase:					Union[str, Any]		= self.embedder(snake_case_					)

				__UpperCAmelCase:					List[str]		= self.encoder(
				    snake_case_   ,    output_hidden_states=snake_case_   ,    return_dict=snake_case_					)

				__UpperCAmelCase:					str		= encoder_outputs[0]

				__UpperCAmelCase:					Union[str, Any]		= self.pooler(snake_case_					)

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

				return BaseModelOutputWithPoolingAndNoAttention(
				    last_hidden_state=snake_case_   ,    pooler_output=snake_case_   ,    hidden_states=encoder_outputs.hidden_states   ,    )




@add_start_docstrings(
    """
    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    """						,  __lowerCAmelCase						,  )
class    a (       __lowerCAmelCase				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_					):

				'''simple docstring'''







				super().__init__(snake_case_					)
				__UpperCAmelCase:					str		= config.num_labels
				__UpperCAmelCase:					List[str]		= ResNetModel(snake_case_					)
				# classification head
				__UpperCAmelCase:					str		= 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(snake_case_					)
		@add_code_sample_docstrings(
		    checkpoint=_IMAGE_CLASS_CHECKPOINT   ,    output_type=snake_case_   ,    config_class=_CONFIG_FOR_DOC   ,    expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT   ,    )
		def    lowercase_  (							self   ,    snake_case_ = None   ,    snake_case_ = None   ,    snake_case_ = None   ,    snake_case_ = None   ,    ):

				'''simple docstring'''







				__UpperCAmelCase:					str		= return_dict if return_dict is not None else self.config.use_return_dict

				__UpperCAmelCase:					Any		= self.resnet(snake_case_   ,    output_hidden_states=snake_case_   ,    return_dict=snake_case_					)

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

				__UpperCAmelCase:					Optional[int]		= self.classifier(snake_case_					)

				__UpperCAmelCase:					int		= None

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

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

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




@add_start_docstrings(
    """
    ResNet backbone, to be used with frameworks like DETR and MaskFormer.
    """						,  __lowerCAmelCase						,  )
class    a (       __lowerCAmelCase						,  __lowerCAmelCase				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_					):

				'''simple docstring'''







				super().__init__(snake_case_					)
				super()._init_backbone(snake_case_					)

				__UpperCAmelCase:					Any		= [config.embedding_size] + config.hidden_sizes
				__UpperCAmelCase:					List[Any]		= ResNetEmbeddings(snake_case_					)
				__UpperCAmelCase:					Dict		= ResNetEncoder(snake_case_					)

				# initialize weights and apply final processing
				self.post_init()







		@add_start_docstrings_to_model_forward(snake_case_					)
		@replace_return_docstrings(output_type=snake_case_   ,    config_class=_CONFIG_FOR_DOC					)
		def    lowercase_  (							self   ,    snake_case_   ,    snake_case_ = None   ,    snake_case_ = None					):

				'''simple docstring'''







				__UpperCAmelCase:					str		= return_dict if return_dict is not None else self.config.use_return_dict
				__UpperCAmelCase:					Any		= (
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)

				__UpperCAmelCase:					Union[str, Any]		= self.embedder(snake_case_					)

				__UpperCAmelCase:					List[Any]		= self.encoder(snake_case_   ,    output_hidden_states=snake_case_   ,    return_dict=snake_case_					)

				__UpperCAmelCase:					int		= outputs.hidden_states

				__UpperCAmelCase:					str		= ()
				for idx, stage in enumerate(self.stage_names					):
						if stage in self.out_features:
								feature_maps += (hidden_states[idx],)

				if not return_dict:
						__UpperCAmelCase:					Union[str, Any]		= (feature_maps,)
						if output_hidden_states:
								output += (outputs.hidden_states,)
						return output

				return BackboneOutput(
				    feature_maps=snake_case_   ,    hidden_states=outputs.hidden_states if output_hidden_states else None   ,    attentions=snake_case_   ,    )
'''simple docstring'''






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


SCREAMING_SNAKE_CASE_	   =			{
    'configuration_rembert': ['REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'RemBertConfig', 'RemBertOnnxConfig']
}

try:
    if not is_sentencepiece_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    SCREAMING_SNAKE_CASE_	   =			['RemBertTokenizer']

try:
    if not is_tokenizers_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    SCREAMING_SNAKE_CASE_	   =			['RemBertTokenizerFast']

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    SCREAMING_SNAKE_CASE_	   =			[
        'REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
        'RemBertForCausalLM',
        'RemBertForMaskedLM',
        'RemBertForMultipleChoice',
        'RemBertForQuestionAnswering',
        'RemBertForSequenceClassification',
        'RemBertForTokenClassification',
        'RemBertLayer',
        'RemBertModel',
        'RemBertPreTrainedModel',
        'load_tf_weights_in_rembert',
    ]


try:
    if not is_tf_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    SCREAMING_SNAKE_CASE_	   =			[
        'TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
        'TFRemBertForCausalLM',
        'TFRemBertForMaskedLM',
        'TFRemBertForMultipleChoice',
        'TFRemBertForQuestionAnswering',
        'TFRemBertForSequenceClassification',
        'TFRemBertForTokenClassification',
        'TFRemBertLayer',
        'TFRemBertModel',
        'TFRemBertPreTrainedModel',
    ]


if TYPE_CHECKING:
    from .configuration_rembert import REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RemBertConfig, RemBertOnnxConfig

    try:
        if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .tokenization_rembert import RemBertTokenizer

    try:
        if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .tokenization_rembert_fast import RemBertTokenizerFast

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_rembert import (
            REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
            RemBertForCausalLM,
            RemBertForMaskedLM,
            RemBertForMultipleChoice,
            RemBertForQuestionAnswering,
            RemBertForSequenceClassification,
            RemBertForTokenClassification,
            RemBertLayer,
            RemBertModel,
            RemBertPreTrainedModel,
            load_tf_weights_in_rembert,
        )

    try:
        if not is_tf_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_tf_rembert import (
            TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
            TFRemBertForCausalLM,
            TFRemBertForMaskedLM,
            TFRemBertForMultipleChoice,
            TFRemBertForQuestionAnswering,
            TFRemBertForSequenceClassification,
            TFRemBertForTokenClassification,
            TFRemBertLayer,
            TFRemBertModel,
            TFRemBertPreTrainedModel,
        )


else:
    import sys

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