

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


a__				=		logging.get_logger(__name__)

# General docstring
a__				=		'''ResNetConfig'''

# Base docstring
a__				=		'''microsoft/resnet-50'''
a__				=		[1, 2048, 7, 7]

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

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

class 				UpperCAmelCase_		(			nn.Module   ):
							"""simple docstring"""







							def __init__(			self       ,							_a       ,							_a       ,							_a = 3       ,							_a = 1       ,							_a = "relu"       )    ->    str:
												super().__init__()
												_a	:    Dict         =		nn.Convad(
												    _a       ,							_a       ,							kernel_size=_a       ,							stride=_a       ,							padding=kernel_size // 2       ,							bias=_a       )
												_a	:    Tuple         =		nn.BatchNormad(_a       )
												_a	:    Dict         =		ACTaFN[activation] if activation is not None else nn.Identity()




							def 				__lowercase				(			self       ,							_a       )    ->    Tensor:
												_a	:    Optional[Any]         =		self.convolution(_a       )
												_a	:    Dict         =		self.normalization(_a       )
												_a	:    Dict         =		self.activation(_a       )
												return hidden_state



class 				UpperCAmelCase_		(			nn.Module   ):
							"""simple docstring"""







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




							def 				__lowercase				(			self       ,							_a       )    ->    Tensor:
												_a	:    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.'''       )
												_a	:    List[str]         =		self.embedder(_a       )
												_a	:    int         =		self.pooler(_a       )
												return embedding



class 				UpperCAmelCase_		(			nn.Module   ):
							"""simple docstring"""







							def __init__(			self       ,							_a       ,							_a       ,							_a = 2       )    ->    Dict:
												super().__init__()
												_a	:    Tuple         =		nn.Convad(_a       ,							_a       ,							kernel_size=1       ,							stride=_a       ,							bias=_a       )
												_a	:    List[Any]         =		nn.BatchNormad(_a       )




							def 				__lowercase				(			self       ,							_a       )    ->    Tensor:
												_a	:    Any         =		self.convolution(_a       )
												_a	:    Optional[int]         =		self.normalization(_a       )
												return hidden_state



class 				UpperCAmelCase_		(			nn.Module   ):
							"""simple docstring"""







							def __init__(			self       ,							_a       ,							_a       ,							_a = 1       ,							_a = "relu"       )    ->    Any:
												super().__init__()
												_a	:    List[Any]         =		in_channels != out_channels or stride != 1
												_a	:    Union[str, Any]         =		(
												    ResNetShortCut(_a       ,							_a       ,							stride=_a       ) if should_apply_shortcut else nn.Identity()
												)
												_a	:    List[str]         =		nn.Sequential(
												    ResNetConvLayer(_a       ,							_a       ,							stride=_a       )       ,							ResNetConvLayer(_a       ,							_a       ,							activation=_a       )       ,							)
												_a	:    int         =		ACTaFN[activation]




							def 				__lowercase				(			self       ,							_a       )    ->    Any:
												_a	:    List[str]         =		hidden_state
												_a	:    Optional[Any]         =		self.layer(_a       )
												_a	:    Union[str, Any]         =		self.shortcut(_a       )
												hidden_state += residual
												_a	:    Optional[Any]         =		self.activation(_a       )
												return hidden_state



class 				UpperCAmelCase_		(			nn.Module   ):
							"""simple docstring"""







							def __init__(			self       ,							_a       ,							_a       ,							_a = 1       ,							_a = "relu"       ,							_a = 4       )    ->    int:
												super().__init__()
												_a	:    List[str]         =		in_channels != out_channels or stride != 1
												_a	:    Dict         =		out_channels // reduction
												_a	:    List[str]         =		(
												    ResNetShortCut(_a       ,							_a       ,							stride=_a       ) if should_apply_shortcut else nn.Identity()
												)
												_a	:    Tuple         =		nn.Sequential(
												    ResNetConvLayer(_a       ,							_a       ,							kernel_size=1       )       ,							ResNetConvLayer(_a       ,							_a       ,							stride=_a       )       ,							ResNetConvLayer(_a       ,							_a       ,							kernel_size=1       ,							activation=_a       )       ,							)
												_a	:    int         =		ACTaFN[activation]




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



class 				UpperCAmelCase_		(			nn.Module   ):
							"""simple docstring"""







							def __init__(			self       ,							_a       ,							_a       ,							_a       ,							_a = 2       ,							_a = 2       ,							)    ->    int:
												super().__init__()

												_a	:    Tuple         =		ResNetBottleNeckLayer if config.layer_type == '''bottleneck''' else ResNetBasicLayer

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




							def 				__lowercase				(			self       ,							_a       )    ->    Tensor:
												_a	:    str         =		input
												for layer in self.layers:
																	_a	:    List[Any]         =		layer(_a       )
												return hidden_state



class 				UpperCAmelCase_		(			nn.Module   ):
							"""simple docstring"""







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




							def 				__lowercase				(			self       ,							_a       ,							_a = False       ,							_a = True       )    ->    BaseModelOutputWithNoAttention:
												_a	:    Optional[Any]         =		() if output_hidden_states else None

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

																	_a	:    Union[str, Any]         =		stage_module(_a       )

												if output_hidden_states:
																	_a	:    Any         =		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=_a       ,							hidden_states=_a       ,							)



class 				UpperCAmelCase_		(			__lowercase   ):
							"""simple docstring"""







							UpperCAmelCase__       :						List[Any]	    =    ResNetConfig
							UpperCAmelCase__       :						Optional[int]	    =    "resnet"
							UpperCAmelCase__       :						Dict	    =    "pixel_values"
							UpperCAmelCase__       :						Tuple	    =    True


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




							def 				__lowercase				(			self       ,							_a       ,							_a=False       )    ->    List[Any]:
												if isinstance(_a       ,							_a       ):
																	_a	:    List[Any]         =		value



a__				=		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.
'''

a__				=		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."				,       __lowercase				,       )
class 				UpperCAmelCase_		(			__lowercase   ):
							"""simple docstring"""







							def __init__(			self       ,							_a       )    ->    List[str]:
												super().__init__(_a       )
												_a	:    List[str]         =		config
												_a	:    int         =		ResNetEmbeddings(_a       )
												_a	:    str         =		ResNetEncoder(_a       )
												_a	:    List[Any]         =		nn.AdaptiveAvgPoolad((1, 1)       )
												# Initialize weights and apply final processing
												self.post_init()




							@add_start_docstrings_to_model_forward(_a       )
							@add_code_sample_docstrings(
							    checkpoint=_CHECKPOINT_FOR_DOC       ,							output_type=_a       ,							config_class=_CONFIG_FOR_DOC       ,							modality='''vision'''       ,							expected_output=_EXPECTED_OUTPUT_SHAPE       ,							)
							def 				__lowercase				(			self       ,							_a       ,							_a = None       ,							_a = None       )    ->    BaseModelOutputWithPoolingAndNoAttention:
												_a	:    str         =		(
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)
												_a	:    Optional[Any]         =		return_dict if return_dict is not None else self.config.use_return_dict

												_a	:    int         =		self.embedder(_a       )

												_a	:    Tuple         =		self.encoder(
												    _a       ,							output_hidden_states=_a       ,							return_dict=_a       )

												_a	:    List[Any]         =		encoder_outputs[0]

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

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

												return BaseModelOutputWithPoolingAndNoAttention(
												    last_hidden_state=_a       ,							pooler_output=_a       ,							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    "				,       __lowercase				,       )
class 				UpperCAmelCase_		(			__lowercase   ):
							"""simple docstring"""







							def __init__(			self       ,							_a       )    ->    Optional[int]:
												super().__init__(_a       )
												_a	:    List[str]         =		config.num_labels
												_a	:    Optional[Any]         =		ResNetModel(_a       )
												# classification head
												_a	:    List[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(_a       )
							@add_code_sample_docstrings(
							    checkpoint=_IMAGE_CLASS_CHECKPOINT       ,							output_type=_a       ,							config_class=_CONFIG_FOR_DOC       ,							expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT       ,							)
							def 				__lowercase				(			self       ,							_a = None       ,							_a = None       ,							_a = None       ,							_a = None       ,							)    ->    ImageClassifierOutputWithNoAttention:
												_a	:    Optional[int]         =		return_dict if return_dict is not None else self.config.use_return_dict

												_a	:    Tuple         =		self.resnet(_a       ,							output_hidden_states=_a       ,							return_dict=_a       )

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

												_a	:    int         =		self.classifier(_a       )

												_a	:    Tuple         =		None

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

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

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



@add_start_docstrings(
    "\n    ResNet backbone, to be used with frameworks like DETR and MaskFormer.\n    "				,       __lowercase				,       )
class 				UpperCAmelCase_		(			__lowercase				,       __lowercase   ):
							"""simple docstring"""







							def __init__(			self       ,							_a       )    ->    Union[str, Any]:
												super().__init__(_a       )
												super()._init_backbone(_a       )

												_a	:    int         =		[config.embedding_size] + config.hidden_sizes
												_a	:    str         =		ResNetEmbeddings(_a       )
												_a	:    Any         =		ResNetEncoder(_a       )

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




							@add_start_docstrings_to_model_forward(_a       )
							@replace_return_docstrings(output_type=_a       ,							config_class=_CONFIG_FOR_DOC       )
							def 				__lowercase				(			self       ,							_a       ,							_a = None       ,							_a = None       )    ->    BackboneOutput:
												_a	:    List[Any]         =		return_dict if return_dict is not None else self.config.use_return_dict
												_a	:    List[str]         =		(
												    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
												)

												_a	:    List[str]         =		self.embedder(_a       )

												_a	:    Optional[int]         =		self.encoder(_a       ,							output_hidden_states=_a       ,							return_dict=_a       )

												_a	:    List[str]         =		outputs.hidden_states

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

												if not return_dict:
																	_a	:    List[Any]         =		(feature_maps,)
																	if output_hidden_states:
																						output += (outputs.hidden_states,)
																	return output

												return BackboneOutput(
												    feature_maps=_a       ,							hidden_states=outputs.hidden_states if output_hidden_states else None       ,							attentions=_a       ,							)






from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


a__				=		logging.get_logger(__name__)

a__				=		{
    '''junnyu/roformer_chinese_small''': '''https://huggingface.co/junnyu/roformer_chinese_small/resolve/main/config.json''',
    '''junnyu/roformer_chinese_base''': '''https://huggingface.co/junnyu/roformer_chinese_base/resolve/main/config.json''',
    '''junnyu/roformer_chinese_char_small''': (
        '''https://huggingface.co/junnyu/roformer_chinese_char_small/resolve/main/config.json'''
    ),
    '''junnyu/roformer_chinese_char_base''': (
        '''https://huggingface.co/junnyu/roformer_chinese_char_base/resolve/main/config.json'''
    ),
    '''junnyu/roformer_small_discriminator''': (
        '''https://huggingface.co/junnyu/roformer_small_discriminator/resolve/main/config.json'''
    ),
    '''junnyu/roformer_small_generator''': (
        '''https://huggingface.co/junnyu/roformer_small_generator/resolve/main/config.json'''
    ),
    # See all RoFormer models at https://huggingface.co/models?filter=roformer
}

class 				UpperCAmelCase_		(			__lowercase   ):
							"""simple docstring"""







							UpperCAmelCase__       :						Optional[Any]	    =    "roformer"


							def __init__(			self       ,							_a=5_0_0_0_0       ,							_a=None       ,							_a=7_6_8       ,							_a=1_2       ,							_a=1_2       ,							_a=3_0_7_2       ,							_a="gelu"       ,							_a=0.1       ,							_a=0.1       ,							_a=1_5_3_6       ,							_a=2       ,							_a=0.02       ,							_a=1e-1_2       ,							_a=0       ,							_a=False       ,							_a=True       ,							**_a       ,							)    ->    List[str]:
												super().__init__(pad_token_id=_a       ,							**_a       )

												_a	:    Tuple         =		vocab_size
												_a	:    List[Any]         =		hidden_size if embedding_size is None else embedding_size
												_a	:    Any         =		hidden_size
												_a	:    Any         =		num_hidden_layers
												_a	:    List[Any]         =		num_attention_heads
												_a	:    str         =		hidden_act
												_a	:    Any         =		intermediate_size
												_a	:    Dict         =		hidden_dropout_prob
												_a	:    Optional[Any]         =		attention_probs_dropout_prob
												_a	:    str         =		max_position_embeddings
												_a	:    Dict         =		type_vocab_size
												_a	:    List[Any]         =		initializer_range
												_a	:    Dict         =		layer_norm_eps
												_a	:    Dict         =		rotary_value
												_a	:    Dict         =		use_cache



class 				UpperCAmelCase_		(			__lowercase   ):
							"""simple docstring"""







							@property
							def 				__lowercase				(			self       )    ->    Mapping[str, Mapping[int, str]]:
												if self.task == "multiple-choice":
																	_a	:    Dict         =		{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
												else:
																	_a	:    List[Any]         =		{0: '''batch''', 1: '''sequence'''}
												_a	:    Union[str, Any]         =		{0: '''batch''', 1: '''sequence'''}
												return OrderedDict(
												    [
												        ('''input_ids''', dynamic_axis),
												        ('''attention_mask''', dynamic_axis),
												        ('''token_type_ids''', dynamic_axis),
												    ]       )



