

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


_UpperCamelCase		     =    logging.get_logger(__name__)

# General docstring
_UpperCamelCase		     =    """ResNetConfig"""

# Base docstring
_UpperCamelCase		     =    """microsoft/resnet-50"""
_UpperCamelCase		     =    [1, 2_0_4_8, 7, 7]

# Image classification docstring
_UpperCamelCase		     =    """microsoft/resnet-50"""
_UpperCamelCase		     =    """tiger cat"""

_UpperCamelCase		     =    [
    """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__()
        lowerCAmelCase__   :				str 			=				nn.Convad(
            snake_case			,					snake_case			,					kernel_size=snake_case			,					stride=snake_case			,					padding=kernel_size // 2			,					bias=snake_case   )
        lowerCAmelCase__   :				str 			=				nn.BatchNormad(snake_case   )
        lowerCAmelCase__   :				Dict 			=				ACTaFN[activation] if activation is not None else nn.Identity()






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[Any] 			=				self.convolution(snake_case   )
        lowerCAmelCase__   :				str 			=				self.normalization(snake_case   )
        lowerCAmelCase__   :				Optional[Any] 			=				self.activation(snake_case   )
        return hidden_state






class 			__a					(						nn.Module       ):





       """simple docstring"""


       def __init__(   self			,					snake_case   ):





        """simple docstring"""

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






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[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."   )
        lowerCAmelCase__   :				Optional[Any] 			=				self.embedder(snake_case   )
        lowerCAmelCase__   :				int 			=				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__()
        lowerCAmelCase__   :				Any 			=				nn.Convad(snake_case			,					snake_case			,					kernel_size=1			,					stride=snake_case			,					bias=snake_case   )
        lowerCAmelCase__   :				List[str] 			=				nn.BatchNormad(snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[int] 			=				self.convolution(snake_case   )
        lowerCAmelCase__   :				Optional[Any] 			=				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__()
        lowerCAmelCase__   :				List[Any] 			=				in_channels != out_channels or stride != 1
        lowerCAmelCase__   :				Dict 			=				(
            ResNetShortCut(snake_case			,					snake_case			,					stride=snake_case   ) if should_apply_shortcut else nn.Identity()
        )
        lowerCAmelCase__   :				Union[str, Any] 			=				nn.Sequential(
            ResNetConvLayer(snake_case			,					snake_case			,					stride=snake_case   )			,					ResNetConvLayer(snake_case			,					snake_case			,					activation=snake_case   )			,					)
        lowerCAmelCase__   :				List[str] 			=				ACTaFN[activation]






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case   ):





        """simple docstring"""

        lowerCAmelCase__   :				List[Any] 			=				hidden_state
        lowerCAmelCase__   :				Optional[int] 			=				self.layer(snake_case   )
        lowerCAmelCase__   :				List[Any] 			=				self.shortcut(snake_case   )
        hidden_state += residual
        lowerCAmelCase__   :				Dict 			=				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__()
        lowerCAmelCase__   :				Union[str, Any] 			=				in_channels != out_channels or stride != 1
        lowerCAmelCase__   :				Any 			=				out_channels // reduction
        lowerCAmelCase__   :				Optional[Any] 			=				(
            ResNetShortCut(snake_case			,					snake_case			,					stride=snake_case   ) if should_apply_shortcut else nn.Identity()
        )
        lowerCAmelCase__   :				Optional[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   )			,					)
        lowerCAmelCase__   :				Dict 			=				ACTaFN[activation]






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case   ):





        """simple docstring"""

        lowerCAmelCase__   :				Tuple 			=				hidden_state
        lowerCAmelCase__   :				Tuple 			=				self.layer(snake_case   )
        lowerCAmelCase__   :				str 			=				self.shortcut(snake_case   )
        hidden_state += residual
        lowerCAmelCase__   :				List[str] 			=				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__()

        lowerCAmelCase__   :				str 			=				ResNetBottleNeckLayer if config.layer_type == "bottleneck" else ResNetBasicLayer

        lowerCAmelCase__   :				Dict 			=				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 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case   ):





        """simple docstring"""

        lowerCAmelCase__   :				Tuple 			=				input
        for layer in self.layers:
         lowerCAmelCase__   :				Optional[int] 			=				layer(snake_case   )
        return hidden_state






class 			__a					(						nn.Module       ):





       """simple docstring"""


       def __init__(   self			,					snake_case   ):





        """simple docstring"""

        super().__init__()
        lowerCAmelCase__   :				str 			=				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]			,					)   )
        lowerCAmelCase__   :				List[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 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case = False			,					snake_case = True   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[Any] 			=				() if output_hidden_states else None

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

         lowerCAmelCase__   :				Tuple 			=				stage_module(snake_case   )

        if output_hidden_states:
         lowerCAmelCase__   :				List[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					(						__magic_name__       ):





       """simple docstring"""


       __UpperCamelCase					:      Union[str, Any]    =						ResNetConfig
       __UpperCamelCase					:      List[str]    =						'resnet'
       __UpperCamelCase					:      Union[str, Any]    =						'pixel_values'
       __UpperCamelCase					:      int    =						True






       def 			SCREAMING_SNAKE_CASE_	(   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 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case=False   ):





        """simple docstring"""

        if isinstance(snake_case			,					snake_case   ):
         lowerCAmelCase__   :				int 			=				value




_UpperCamelCase		     =    r"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

_UpperCamelCase		     =    r"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConvNextImageProcessor.__call__`] for details.

        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
"""






@add_start_docstrings(
    'The bare ResNet model outputting raw features without any specific head on top.'       ,	__magic_name__       ,	)
class 			__a					(						__magic_name__       ):





       """simple docstring"""


       def __init__(   self			,					snake_case   ):





        """simple docstring"""

        super().__init__(snake_case   )
        lowerCAmelCase__   :				List[str] 			=				config
        lowerCAmelCase__   :				Optional[int] 			=				ResNetEmbeddings(snake_case   )
        lowerCAmelCase__   :				str 			=				ResNetEncoder(snake_case   )
        lowerCAmelCase__   :				Optional[int] 			=				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 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case = None			,					snake_case = None   ):





        """simple docstring"""

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

        lowerCAmelCase__   :				Optional[int] 			=				self.embedder(snake_case   )

        lowerCAmelCase__   :				Union[str, Any] 			=				self.encoder(
            snake_case			,					output_hidden_states=snake_case			,					return_dict=snake_case   )

        lowerCAmelCase__   :				List[str] 			=				encoder_outputs[0]

        lowerCAmelCase__   :				Tuple 			=				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(
    '\n    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    '       ,	__magic_name__       ,	)
class 			__a					(						__magic_name__       ):





       """simple docstring"""


       def __init__(   self			,					snake_case   ):





        """simple docstring"""

        super().__init__(snake_case   )
        lowerCAmelCase__   :				int 			=				config.num_labels
        lowerCAmelCase__   :				List[str] 			=				ResNetModel(snake_case   )
        # classification head
        lowerCAmelCase__   :				Optional[int] 			=				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 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case = None			,					snake_case = None			,					snake_case = None			,					snake_case = None			,					):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[Any] 			=				return_dict if return_dict is not None else self.config.use_return_dict

        lowerCAmelCase__   :				Union[str, Any] 			=				self.resnet(snake_case			,					output_hidden_states=snake_case			,					return_dict=snake_case   )

        lowerCAmelCase__   :				List[str] 			=				outputs.pooler_output if return_dict else outputs[1]

        lowerCAmelCase__   :				Optional[int] 			=				self.classifier(snake_case   )

        lowerCAmelCase__   :				Union[str, Any] 			=				None

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

        if not return_dict:
         lowerCAmelCase__   :				Tuple 			=				(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(
    '\n    ResNet backbone, to be used with frameworks like DETR and MaskFormer.\n    '       ,	__magic_name__       ,	)
class 			__a					(						__magic_name__       ,	__magic_name__       ):





       """simple docstring"""


       def __init__(   self			,					snake_case   ):





        """simple docstring"""

        super().__init__(snake_case   )
        super()._init_backbone(snake_case   )

        lowerCAmelCase__   :				int 			=				[config.embedding_size] + config.hidden_sizes
        lowerCAmelCase__   :				Optional[Any] 			=				ResNetEmbeddings(snake_case   )
        lowerCAmelCase__   :				List[Any] 			=				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 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case = None			,					snake_case = None   ):





        """simple docstring"""

        lowerCAmelCase__   :				List[str] 			=				return_dict if return_dict is not None else self.config.use_return_dict
        lowerCAmelCase__   :				Union[str, Any] 			=				(
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )

        lowerCAmelCase__   :				Union[str, Any] 			=				self.embedder(snake_case   )

        lowerCAmelCase__   :				Any 			=				self.encoder(snake_case			,					output_hidden_states=snake_case			,					return_dict=snake_case   )

        lowerCAmelCase__   :				str 			=				outputs.hidden_states

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

        if not return_dict:
         lowerCAmelCase__   :				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_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


if is_sentencepiece_available():
    from ..ta.tokenization_ta import TaTokenizer
else:
    from ...utils.dummy_sentencepiece_objects import TaTokenizer

_UpperCamelCase		     =    TaTokenizer

if is_tokenizers_available():
    from ..ta.tokenization_ta_fast import TaTokenizerFast
else:
    from ...utils.dummy_tokenizers_objects import TaTokenizerFast

_UpperCamelCase		     =    TaTokenizerFast

_UpperCamelCase		     =    {"""configuration_mt5""": ["""MT5Config""", """MT5OnnxConfig"""]}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    _UpperCamelCase		     =    [
        """MT5EncoderModel""",
        """MT5ForConditionalGeneration""",
        """MT5ForQuestionAnswering""",
        """MT5Model""",
        """MT5PreTrainedModel""",
        """MT5Stack""",
    ]

try:
    if not is_tf_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    _UpperCamelCase		     =    ["""TFMT5EncoderModel""", """TFMT5ForConditionalGeneration""", """TFMT5Model"""]

try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    _UpperCamelCase		     =    ["""FlaxMT5EncoderModel""", """FlaxMT5ForConditionalGeneration""", """FlaxMT5Model"""]


if TYPE_CHECKING:
    from .configuration_mta import MTaConfig, MTaOnnxConfig

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_mta import (
            MTaEncoderModel,
            MTaForConditionalGeneration,
            MTaForQuestionAnswering,
            MTaModel,
            MTaPreTrainedModel,
            MTaStack,
        )

    try:
        if not is_tf_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_tf_mta import TFMTaEncoderModel, TFMTaForConditionalGeneration, TFMTaModel

    try:
        if not is_flax_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_flax_mta import FlaxMTaEncoderModel, FlaxMTaForConditionalGeneration, FlaxMTaModel

else:
    import sys

    _UpperCamelCase		     =    _LazyModule(
        __name__,
        globals()["""__file__"""],
        _import_structure,
        extra_objects={"""MT5Tokenizer""": MTaTokenizer, """MT5TokenizerFast""": MTaTokenizerFast},
        module_spec=__spec__,
    )






