






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__			:    List[Any]         =				logging.get_logger(__name__)

# General docstring
A__			:    Any         =				'ResNetConfig'

# Base docstring
A__			:    Union[str, Any]         =				'microsoft/resnet-50'
A__			:    int         =				[1, 20_48, 7, 7]

# Image classification docstring
A__			:    int         =				'microsoft/resnet-50'
A__			:    Any         =				'tiger cat'

A__			:    Tuple         =				[
    'microsoft/resnet-50',
    # See all resnet models at https://huggingface.co/models?filter=resnet
]
class    _UpperCAmelCase      ( nn.Module			):

				"""simple docstring"""




				def __init__(  self		:					Tuple,						lowerCamelCase		:					Any,						lowerCamelCase		:					Dict,						lowerCamelCase		:					List[Any] = 3,						lowerCamelCase		:					Optional[Any] = 1,						lowerCamelCase		:					Tuple = "relu"						):

								'''simple docstring'''
								super().__init__()
								lowercase__     				=    nn.Convad(
								    lowerCamelCase,						lowerCamelCase,						kernel_size=lowerCamelCase,						stride=lowerCamelCase,						padding=kernel_size // 2,						bias=lowerCamelCase						)
								lowercase__     				=    nn.BatchNormad(lowerCamelCase						)
								lowercase__     				=    ACTaFN[activation] if activation is not None else nn.Identity()




				def        lowercase__     (  self		:					Optional[int],						lowerCamelCase		:					Optional[Any]						):

								'''simple docstring'''
								lowercase__     				=    self.convolution(lowerCamelCase						)
								lowercase__     				=    self.normalization(lowerCamelCase						)
								lowercase__     				=    self.activation(lowerCamelCase						)
								return hidden_state



class    _UpperCAmelCase      ( nn.Module			):

				"""simple docstring"""




				def __init__(  self		:					List[Any],						lowerCamelCase		:					int						):

								'''simple docstring'''
								super().__init__()
								lowercase__     				=    ResNetConvLayer(
								    config.num_channels,						config.embedding_size,						kernel_size=7,						stride=2,						activation=config.hidden_act						)
								lowercase__     				=    nn.MaxPoolad(kernel_size=3,						stride=2,						padding=1						)
								lowercase__     				=    config.num_channels




				def        lowercase__     (  self		:					List[str],						lowerCamelCase		:					Optional[int]						):

								'''simple docstring'''
								lowercase__     				=    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.'''						)
								lowercase__     				=    self.embedder(lowerCamelCase						)
								lowercase__     				=    self.pooler(lowerCamelCase						)
								return embedding



class    _UpperCAmelCase      ( nn.Module			):

				"""simple docstring"""




				def __init__(  self		:					int,						lowerCamelCase		:					List[str],						lowerCamelCase		:					Tuple,						lowerCamelCase		:					str = 2						):

								'''simple docstring'''
								super().__init__()
								lowercase__     				=    nn.Convad(lowerCamelCase,						lowerCamelCase,						kernel_size=1,						stride=lowerCamelCase,						bias=lowerCamelCase						)
								lowercase__     				=    nn.BatchNormad(lowerCamelCase						)




				def        lowercase__     (  self		:					Union[str, Any],						lowerCamelCase		:					List[str]						):

								'''simple docstring'''
								lowercase__     				=    self.convolution(lowerCamelCase						)
								lowercase__     				=    self.normalization(lowerCamelCase						)
								return hidden_state



class    _UpperCAmelCase      ( nn.Module			):

				"""simple docstring"""




				def __init__(  self		:					List[str],						lowerCamelCase		:					int,						lowerCamelCase		:					Optional[Any],						lowerCamelCase		:					Union[str, Any] = 1,						lowerCamelCase		:					Dict = "relu"						):

								'''simple docstring'''
								super().__init__()
								lowercase__     				=    in_channels != out_channels or stride != 1
								lowercase__     				=    (
								    ResNetShortCut(lowerCamelCase,						lowerCamelCase,						stride=lowerCamelCase						) if should_apply_shortcut else nn.Identity()
								)
								lowercase__     				=    nn.Sequential(
								    ResNetConvLayer(lowerCamelCase,						lowerCamelCase,						stride=lowerCamelCase						),						ResNetConvLayer(lowerCamelCase,						lowerCamelCase,						activation=lowerCamelCase						),						)
								lowercase__     				=    ACTaFN[activation]




				def        lowercase__     (  self		:					Union[str, Any],						lowerCamelCase		:					Dict						):

								'''simple docstring'''
								lowercase__     				=    hidden_state
								lowercase__     				=    self.layer(lowerCamelCase						)
								lowercase__     				=    self.shortcut(lowerCamelCase						)
								hidden_state += residual
								lowercase__     				=    self.activation(lowerCamelCase						)
								return hidden_state



class    _UpperCAmelCase      ( nn.Module			):

				"""simple docstring"""




				def __init__(  self		:					Optional[int],						lowerCamelCase		:					Tuple,						lowerCamelCase		:					List[str],						lowerCamelCase		:					str = 1,						lowerCamelCase		:					Tuple = "relu",						lowerCamelCase		:					Optional[int] = 4						):

								'''simple docstring'''
								super().__init__()
								lowercase__     				=    in_channels != out_channels or stride != 1
								lowercase__     				=    out_channels // reduction
								lowercase__     				=    (
								    ResNetShortCut(lowerCamelCase,						lowerCamelCase,						stride=lowerCamelCase						) if should_apply_shortcut else nn.Identity()
								)
								lowercase__     				=    nn.Sequential(
								    ResNetConvLayer(lowerCamelCase,						lowerCamelCase,						kernel_size=1						),						ResNetConvLayer(lowerCamelCase,						lowerCamelCase,						stride=lowerCamelCase						),						ResNetConvLayer(lowerCamelCase,						lowerCamelCase,						kernel_size=1,						activation=lowerCamelCase						),						)
								lowercase__     				=    ACTaFN[activation]




				def        lowercase__     (  self		:					str,						lowerCamelCase		:					Tuple						):

								'''simple docstring'''
								lowercase__     				=    hidden_state
								lowercase__     				=    self.layer(lowerCamelCase						)
								lowercase__     				=    self.shortcut(lowerCamelCase						)
								hidden_state += residual
								lowercase__     				=    self.activation(lowerCamelCase						)
								return hidden_state



class    _UpperCAmelCase      ( nn.Module			):

				"""simple docstring"""




				def __init__(  self		:					Tuple,						lowerCamelCase		:					List[Any],						lowerCamelCase		:					Dict,						lowerCamelCase		:					List[Any],						lowerCamelCase		:					Optional[int] = 2,						lowerCamelCase		:					Union[str, Any] = 2,						):

								'''simple docstring'''
								super().__init__()

								lowercase__     				=    ResNetBottleNeckLayer if config.layer_type == '''bottleneck''' else ResNetBasicLayer

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




				def        lowercase__     (  self		:					Dict,						lowerCamelCase		:					str						):

								'''simple docstring'''
								lowercase__     				=    input
								for layer in self.layers:
												lowercase__     				=    layer(lowerCamelCase						)
								return hidden_state



class    _UpperCAmelCase      ( nn.Module			):

				"""simple docstring"""




				def __init__(  self		:					Optional[Any],						lowerCamelCase		:					Optional[Any]						):

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




				def        lowercase__     (  self		:					Dict,						lowerCamelCase		:					Optional[int],						lowerCamelCase		:					str = False,						lowerCamelCase		:					List[str] = True						):

								'''simple docstring'''
								lowercase__     				=    () if output_hidden_states else None

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

												lowercase__     				=    stage_module(lowerCamelCase						)

								if output_hidden_states:
												lowercase__     				=    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=lowerCamelCase,						hidden_states=lowerCamelCase,						)



class    _UpperCAmelCase      ( A__			):

				"""simple docstring"""




				lowercase__       							= ResNetConfig
				lowercase__       							= """resnet"""
				lowercase__       							= """pixel_values"""
				lowercase__       							= True




				def        lowercase__     (  self		:					Optional[Any],						lowerCamelCase		:					Optional[Any]						):

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




				def        lowercase__     (  self		:					List[Any],						lowerCamelCase		:					Tuple,						lowerCamelCase		:					Union[str, Any]=False						):

								'''simple docstring'''
								if isinstance(lowerCamelCase,						lowerCamelCase						):
												lowercase__     				=    value



A__			:    Optional[int]         =				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'

A__			:    Union[str, Any]         =				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."""	,A__	,)
class    _UpperCAmelCase      ( A__			):

				"""simple docstring"""




				def __init__(  self		:					Dict,						lowerCamelCase		:					Optional[Any]						):

								'''simple docstring'''
								super().__init__(lowerCamelCase						)
								lowercase__     				=    config
								lowercase__     				=    ResNetEmbeddings(lowerCamelCase						)
								lowercase__     				=    ResNetEncoder(lowerCamelCase						)
								lowercase__     				=    nn.AdaptiveAvgPoolad((1, 1)						)
								# Initialize weights and apply final processing
								self.post_init()




				@add_start_docstrings_to_model_forward(lowerCamelCase						)
				@add_code_sample_docstrings(
				    checkpoint=_CHECKPOINT_FOR_DOC,						output_type=lowerCamelCase,						config_class=_CONFIG_FOR_DOC,						modality='''vision''',						expected_output=_EXPECTED_OUTPUT_SHAPE,						)
				def        lowercase__     (  self		:					int,						lowerCamelCase		:					Optional[int],						lowerCamelCase		:					List[str] = None,						lowerCamelCase		:					Union[str, Any] = None						):

								'''simple docstring'''
								lowercase__     				=    (
								    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
								)
								lowercase__     				=    return_dict if return_dict is not None else self.config.use_return_dict

								lowercase__     				=    self.embedder(lowerCamelCase						)

								lowercase__     				=    self.encoder(
								    lowerCamelCase,						output_hidden_states=lowerCamelCase,						return_dict=lowerCamelCase						)

								lowercase__     				=    encoder_outputs[0]

								lowercase__     				=    self.pooler(lowerCamelCase						)

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

								return BaseModelOutputWithPoolingAndNoAttention(
								    last_hidden_state=lowerCamelCase,						pooler_output=lowerCamelCase,						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.
    """	,A__	,)
class    _UpperCAmelCase      ( A__			):

				"""simple docstring"""




				def __init__(  self		:					List[str],						lowerCamelCase		:					Optional[int]						):

								'''simple docstring'''
								super().__init__(lowerCamelCase						)
								lowercase__     				=    config.num_labels
								lowercase__     				=    ResNetModel(lowerCamelCase						)
								# classification head
								lowercase__     				=    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(lowerCamelCase						)
				@add_code_sample_docstrings(
				    checkpoint=_IMAGE_CLASS_CHECKPOINT,						output_type=lowerCamelCase,						config_class=_CONFIG_FOR_DOC,						expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,						)
				def        lowercase__     (  self		:					Any,						lowerCamelCase		:					Tuple = None,						lowerCamelCase		:					str = None,						lowerCamelCase		:					Optional[int] = None,						lowerCamelCase		:					str = None,						):

								'''simple docstring'''
								lowercase__     				=    return_dict if return_dict is not None else self.config.use_return_dict

								lowercase__     				=    self.resnet(lowerCamelCase,						output_hidden_states=lowerCamelCase,						return_dict=lowerCamelCase						)

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

								lowercase__     				=    self.classifier(lowerCamelCase						)

								lowercase__     				=    None

								if labels is not None:
												if self.config.problem_type is None:
																if self.num_labels == 1:
																				lowercase__     				=    '''regression'''
																elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																				lowercase__     				=    '''single_label_classification'''
																else:
																				lowercase__     				=    '''multi_label_classification'''
												if self.config.problem_type == "regression":
																lowercase__     				=    MSELoss()
																if self.num_labels == 1:
																				lowercase__     				=    loss_fct(logits.squeeze(),						labels.squeeze()						)
																else:
																				lowercase__     				=    loss_fct(lowerCamelCase,						lowerCamelCase						)
												elif self.config.problem_type == "single_label_classification":
																lowercase__     				=    CrossEntropyLoss()
																lowercase__     				=    loss_fct(logits.view(-1,						self.num_labels						),						labels.view(-1						)						)
												elif self.config.problem_type == "multi_label_classification":
																lowercase__     				=    BCEWithLogitsLoss()
																lowercase__     				=    loss_fct(lowerCamelCase,						lowerCamelCase						)

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

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



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

				"""simple docstring"""




				def __init__(  self		:					Union[str, Any],						lowerCamelCase		:					List[Any]						):

								'''simple docstring'''
								super().__init__(lowerCamelCase						)
								super()._init_backbone(lowerCamelCase						)

								lowercase__     				=    [config.embedding_size] + config.hidden_sizes
								lowercase__     				=    ResNetEmbeddings(lowerCamelCase						)
								lowercase__     				=    ResNetEncoder(lowerCamelCase						)

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




				@add_start_docstrings_to_model_forward(lowerCamelCase						)
				@replace_return_docstrings(output_type=lowerCamelCase,						config_class=_CONFIG_FOR_DOC						)
				def        lowercase__     (  self		:					Optional[Any],						lowerCamelCase		:					Any,						lowerCamelCase		:					Any = None,						lowerCamelCase		:					str = None						):

								'''simple docstring'''
								lowercase__     				=    return_dict if return_dict is not None else self.config.use_return_dict
								lowercase__     				=    (
								    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
								)

								lowercase__     				=    self.embedder(lowerCamelCase						)

								lowercase__     				=    self.encoder(lowerCamelCase,						output_hidden_states=lowerCamelCase,						return_dict=lowerCamelCase						)

								lowercase__     				=    outputs.hidden_states

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

								if not return_dict:
												lowercase__     				=    (feature_maps,)
												if output_hidden_states:
																output += (outputs.hidden_states,)
												return output

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












import json
import os
import tempfile

import transformers

import datasets
from utils import generate_example_dataset, get_duration


A__			:    Dict         =				50_00_00

A__     ,     A__			:    str         =				os.path.split(__file__)
A__			:    Optional[Any]         =				os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))






@get_duration
def 						a    (		lowerCamelCase_	,     **lowerCamelCase_       ):






				'''simple docstring'''
				lowercase__     				=    dataset.map(**lowerCamelCase_       )






@get_duration
def 						a    (		lowerCamelCase_	,     **lowerCamelCase_       ):






				'''simple docstring'''
				lowercase__     				=    dataset.filter(**lowerCamelCase_       )






def 						a    (		):






				'''simple docstring'''
				lowercase__     				=    {'''num examples''': SPEED_TEST_N_EXAMPLES}
				with tempfile.TemporaryDirectory() as tmp_dir:
								lowercase__     				=    datasets.Features({'''text''': datasets.Value('''string'''       ), '''numbers''': datasets.Value('''float32'''       )}       )
								lowercase__     				=    generate_example_dataset(
								    os.path.join(lowerCamelCase_	,     '''dataset.arrow'''       )	,     lowerCamelCase_	,     num_examples=lowerCamelCase_       )

								lowercase__     				=    transformers.AutoTokenizer.from_pretrained('''bert-base-cased'''	,     use_fast=lowerCamelCase_       )

								def tokenize(lowerCamelCase_       ):
												return tokenizer(examples['''text''']       )

								lowercase__     				=    map(lowerCamelCase_       )

								lowercase__     				=    map(lowerCamelCase_	,     batched=lowerCamelCase_       )

								lowercase__     				=    map(lowerCamelCase_	,     function=lambda lowerCamelCase_       : None	,     batched=lowerCamelCase_       )

								with dataset.formatted_as(type='''numpy'''       ):
												lowercase__     				=    map(lowerCamelCase_	,     function=lambda lowerCamelCase_       : None	,     batched=lowerCamelCase_       )

								with dataset.formatted_as(type='''pandas'''       ):
												lowercase__     				=    map(lowerCamelCase_	,     function=lambda lowerCamelCase_       : None	,     batched=lowerCamelCase_       )

								with dataset.formatted_as(type='''torch'''	,     columns='''numbers'''       ):
												lowercase__     				=    map(lowerCamelCase_	,     function=lambda lowerCamelCase_       : None	,     batched=lowerCamelCase_       )

								with dataset.formatted_as(type='''tensorflow'''	,     columns='''numbers'''       ):
												lowercase__     				=    map(lowerCamelCase_	,     function=lambda lowerCamelCase_       : None	,     batched=lowerCamelCase_       )

								lowercase__     				=    map(lowerCamelCase_	,     function=lowerCamelCase_	,     batched=lowerCamelCase_       )

								lowercase__     				=    filter(lowerCamelCase_       )

								# Activate later when tokenizer support batched inputs
								# with dataset.formatted_as(type='numpy'):
								#     times[func.__name__ + " fast-tokenizer batched numpy"] = func(dataset, function=tokenize, batched=True)

				with open(lowerCamelCase_	,     '''wb'''       ) as f:
								f.write(json.dumps(lowerCamelCase_       ).encode('''utf-8'''       )       )


if __name__ == "__main__":  # useful to run the profiler
							benchmark_map_filter()




