

import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

# Register SEW's fairseq modules
from sew_asapp import tasks  # noqa: F401

from transformers import (
    SEWConfig,
    SEWForCTC,
    SEWModel,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
    logging,
)


logging.set_verbosity_info()
SCREAMING_SNAKE_CASE_            =			logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_            =			{
    'post_extract_proj': 'feature_projection',
    'encoder.pos_conv.0': 'encoder.pos_conv_embed.conv',
    'self_attn.k_proj': 'encoder.layers.*.attention.k_proj',
    'self_attn.v_proj': 'encoder.layers.*.attention.v_proj',
    'self_attn.q_proj': 'encoder.layers.*.attention.q_proj',
    'self_attn.out_proj': 'encoder.layers.*.attention.out_proj',
    'self_attn_layer_norm': 'encoder.layers.*.layer_norm',
    'fc1': 'encoder.layers.*.feed_forward.intermediate_dense',
    'fc2': 'encoder.layers.*.feed_forward.output_dense',
    'final_layer_norm': 'encoder.layers.*.final_layer_norm',
    'encoder.upsample.0': 'encoder.upsample.projection',
    'encoder.layer_norm': 'encoder.layer_norm',
    'w2v_model.layer_norm': 'layer_norm',
    'w2v_encoder.proj': 'lm_head',
    'mask_emb': 'masked_spec_embed',
}

def 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					List[str]					, lowerCAmelCase:					str					, lowerCAmelCase:					Optional[int]					, lowerCAmelCase:					Union[str, Any]					, lowerCAmelCase:					Union[str, Any]	)				->    Optional[int]:


				for attribute in key.split("."	):
								_UpperCAmelCase    :						Tuple							=      getattr(lowerCAmelCase__					, lowerCAmelCase__	)

				if weight_type is not None:
								_UpperCAmelCase    :						Optional[int]							=      getattr(lowerCAmelCase__					, lowerCAmelCase__	).shape
				else:
								_UpperCAmelCase    :						str							=      hf_pointer.shape

				assert hf_shape == value.shape, (
				    F'Shape of hf {key + "." + weight_type if weight_type is not None else ""} is {hf_shape}, but should be'
				    F' {value.shape} for {full_name}'
				)

				if weight_type == "weight":
								_UpperCAmelCase    :						Tuple							=      value
				elif weight_type == "weight_g":
								_UpperCAmelCase    :						str							=      value
				elif weight_type == "weight_v":
								_UpperCAmelCase    :						List[str]							=      value
				elif weight_type == "bias":
								_UpperCAmelCase    :						Tuple							=      value
				else:
								_UpperCAmelCase    :						Union[str, Any]							=      value

				logger.info(F'{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.'	)

def 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					str					, lowerCAmelCase:					Any					, lowerCAmelCase:					Optional[Any]	)				->    Any:
				_UpperCAmelCase    :						Any							=      []
				_UpperCAmelCase    :						Dict							=      fairseq_model.state_dict()

				_UpperCAmelCase    :						Union[str, Any]							=      hf_model.sew.feature_extractor if is_finetuned else hf_model.feature_extractor

				for name, value in fairseq_dict.items():
								_UpperCAmelCase    :						Optional[Any]							=      False
								if "conv_layers" in name:
												load_conv_layer(
												    lowerCAmelCase__					, lowerCAmelCase__					, lowerCAmelCase__					, lowerCAmelCase__					, hf_model.config.feat_extract_norm == "group"					, )
												_UpperCAmelCase    :						Union[str, Any]							=      True
								else:
												for key, mapped_key in MAPPING.items():
																_UpperCAmelCase    :						List[Any]							=      """sew.""" + mapped_key if (is_finetuned and mapped_key != """lm_head""") else mapped_key

																if key in name or key.split("w2v_model."	)[-1] == name.split("."	)[0]:
																				_UpperCAmelCase    :						int							=      True
																				if "*" in mapped_key:
																								_UpperCAmelCase    :						Dict							=      name.split(lowerCAmelCase__	)[0].split("."	)[-2]
																								_UpperCAmelCase    :						List[str]							=      mapped_key.replace("*"					, lowerCAmelCase__	)
																				if "weight_g" in name:
																								_UpperCAmelCase    :						List[str]							=      """weight_g"""
																				elif "weight_v" in name:
																								_UpperCAmelCase    :						List[str]							=      """weight_v"""
																				elif "weight" in name:
																								_UpperCAmelCase    :						int							=      """weight"""
																				elif "bias" in name:
																								_UpperCAmelCase    :						Optional[int]							=      """bias"""
																				else:
																								_UpperCAmelCase    :						List[str]							=      None
																				set_recursively(lowerCAmelCase__					, lowerCAmelCase__					, lowerCAmelCase__					, lowerCAmelCase__					, lowerCAmelCase__	)
																continue
								if not is_used:
												unused_weights.append(lowerCAmelCase__	)

				logger.warning(F'Unused weights: {unused_weights}'	)

def 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					int					, lowerCAmelCase:					Tuple					, lowerCAmelCase:					str					, lowerCAmelCase:					Tuple					, lowerCAmelCase:					Any	)				->    Any:
				_UpperCAmelCase    :						Tuple							=      full_name.split("conv_layers."	)[-1]
				_UpperCAmelCase    :						Tuple							=      name.split("."	)
				_UpperCAmelCase    :						List[Any]							=      int(items[0]	)
				_UpperCAmelCase    :						Optional[int]							=      int(items[1]	)

				if type_id == 0:
								if "bias" in name:
												assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
												    F'{full_name} has size {value.shape}, but'
												    F' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'
												)
												_UpperCAmelCase    :						Dict							=      value
												logger.info(F'Feat extract conv layer {layer_id} was initialized from {full_name}.'	)
								elif "weight" in name:
												assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
												    F'{full_name} has size {value.shape}, but'
												    F' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'
												)
												_UpperCAmelCase    :						Union[str, Any]							=      value
												logger.info(F'Feat extract conv layer {layer_id} was initialized from {full_name}.'	)
				elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
								if "bias" in name:
												assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
												    F'{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was'
												    " found."
												)
												_UpperCAmelCase    :						Union[str, Any]							=      value
												logger.info(F'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'	)
								elif "weight" in name:
												assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
												    F'{full_name} has size {value.shape}, but'
												    F' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'
												)
												_UpperCAmelCase    :						List[str]							=      value
												logger.info(F'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'	)
				else:
								unused_weights.append(lowerCAmelCase__	)

def 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					Tuple					, lowerCAmelCase:					Dict	)				->    int:
				_UpperCAmelCase    :						Dict							=      SEWConfig()
				if is_finetuned:
								_UpperCAmelCase    :						Optional[int]							=      model.wav_encoder.wav_model.cfg
				else:
								_UpperCAmelCase    :						int							=      model.cfg

				_UpperCAmelCase    :						Union[str, Any]							=      fs_config.conv_bias
				_UpperCAmelCase    :						Optional[Any]							=      eval(fs_config.conv_feature_layers	)
				_UpperCAmelCase    :						Any							=      [x[0] for x in conv_layers]
				_UpperCAmelCase    :						str							=      [x[1] for x in conv_layers]
				_UpperCAmelCase    :						Union[str, Any]							=      [x[2] for x in conv_layers]
				_UpperCAmelCase    :						Optional[int]							=      """gelu"""
				_UpperCAmelCase    :						Optional[Any]							=      """layer""" if fs_config.extractor_mode == """layer_norm""" else """group"""
				_UpperCAmelCase    :						List[Any]							=      0.0
				_UpperCAmelCase    :						int							=      fs_config.activation_fn.name
				_UpperCAmelCase    :						Optional[Any]							=      fs_config.encoder_embed_dim
				_UpperCAmelCase    :						List[Any]							=      0.02
				_UpperCAmelCase    :						Optional[int]							=      fs_config.encoder_ffn_embed_dim
				_UpperCAmelCase    :						List[str]							=      1E-5
				_UpperCAmelCase    :						str							=      fs_config.encoder_layerdrop
				_UpperCAmelCase    :						str							=      fs_config.encoder_attention_heads
				_UpperCAmelCase    :						List[str]							=      fs_config.conv_pos_groups
				_UpperCAmelCase    :						int							=      fs_config.conv_pos
				_UpperCAmelCase    :						int							=      len(lowerCAmelCase__	)
				_UpperCAmelCase    :						Any							=      fs_config.encoder_layers
				_UpperCAmelCase    :						List[Any]							=      fs_config.squeeze_factor

				# take care of any params that are overridden by the Wav2VecCtc model
				if is_finetuned:
								_UpperCAmelCase    :						int							=      model.cfg
								_UpperCAmelCase    :						Any							=      fs_config.final_dropout
								_UpperCAmelCase    :						Union[str, Any]							=      fs_config.layerdrop
				_UpperCAmelCase    :						Dict							=      fs_config.activation_dropout
				_UpperCAmelCase    :						List[Any]							=      fs_config.mask_prob > 0 or fs_config.mask_channel_prob > 0
				_UpperCAmelCase    :						Optional[Any]							=      fs_config.attention_dropout
				_UpperCAmelCase    :						Dict							=      fs_config.dropout_input
				_UpperCAmelCase    :						Optional[int]							=      fs_config.dropout
				_UpperCAmelCase    :						Dict							=      fs_config.mask_channel_length
				_UpperCAmelCase    :						Optional[int]							=      fs_config.mask_channel_prob
				_UpperCAmelCase    :						Union[str, Any]							=      fs_config.mask_length
				_UpperCAmelCase    :						List[Any]							=      fs_config.mask_prob

				_UpperCAmelCase    :						Dict							=      """Wav2Vec2FeatureExtractor"""
				_UpperCAmelCase    :						Dict							=      """Wav2Vec2CTCTokenizer"""

				return config

@torch.no_grad()
def 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					Optional[int]					, lowerCAmelCase:					Dict					, lowerCAmelCase:					Optional[int]=None					, lowerCAmelCase:					List[str]=None					, lowerCAmelCase:					List[Any]=True	)				->    Union[str, Any]:


				if is_finetuned:
								_UpperCAmelCase    :						List[Any]							=      fairseq.checkpoint_utils.load_model_ensemble_and_task(
								    [checkpoint_path]					, arg_overrides={"data": "/".join(dict_path.split("/"	)[:-1]	)}	)
				else:
								_UpperCAmelCase    :						int							=      fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]	)

				if config_path is not None:
								_UpperCAmelCase    :						Optional[Any]							=      SEWConfig.from_pretrained(lowerCAmelCase__	)
				else:
								_UpperCAmelCase    :						int							=      convert_config(model[0]					, lowerCAmelCase__	)
				_UpperCAmelCase    :						str							=      model[0].eval()

				_UpperCAmelCase    :						Tuple							=      True if config.feat_extract_norm == """layer""" else False
				_UpperCAmelCase    :						Union[str, Any]							=      WavaVecaFeatureExtractor(
				    feature_size=1					, sampling_rate=1_6000					, padding_value=0					, do_normalize=lowerCAmelCase__					, return_attention_mask=lowerCAmelCase__					, )

				if is_finetuned:
								if dict_path:
												_UpperCAmelCase    :						Union[str, Any]							=      Dictionary.load(lowerCAmelCase__	)

												# important change bos & pad token id since CTC symbol is <pad> and
												# not <s> as in fairseq
												_UpperCAmelCase    :						Dict							=      target_dict.pad_index
												_UpperCAmelCase    :						Tuple							=      target_dict.bos_index
												_UpperCAmelCase    :						Optional[int]							=      target_dict.pad_index
												_UpperCAmelCase    :						str							=      target_dict.bos_index
												_UpperCAmelCase    :						List[Any]							=      target_dict.eos_index
												_UpperCAmelCase    :						List[str]							=      len(target_dict.symbols	)
												_UpperCAmelCase    :						Dict							=      os.path.join(lowerCAmelCase__					, "vocab.json"	)
												if not os.path.isdir(lowerCAmelCase__	):
																logger.error("--pytorch_dump_folder_path ({}) should be a directory".format(lowerCAmelCase__	)	)
																return
												os.makedirs(lowerCAmelCase__					, exist_ok=lowerCAmelCase__	)
												with open(lowerCAmelCase__					, "w"					, encoding="utf-8"	) as vocab_handle:
																json.dump(target_dict.indices					, lowerCAmelCase__	)
												_UpperCAmelCase    :						Union[str, Any]							=      WavaVecaCTCTokenizer(
												    lowerCAmelCase__					, unk_token=target_dict.unk_word					, pad_token=target_dict.pad_word					, bos_token=target_dict.bos_word					, eos_token=target_dict.eos_word					, word_delimiter_token="|"					, do_lower_case=lowerCAmelCase__					, )
												_UpperCAmelCase    :						Optional[Any]							=      WavaVecaProcessor(feature_extractor=lowerCAmelCase__					, tokenizer=lowerCAmelCase__	)
												processor.save_pretrained(lowerCAmelCase__	)

								_UpperCAmelCase    :						Union[str, Any]							=      SEWForCTC(lowerCAmelCase__	)
				else:
								_UpperCAmelCase    :						List[Any]							=      SEWModel(lowerCAmelCase__	)
								feature_extractor.save_pretrained(lowerCAmelCase__	)

				recursively_load_weights(lowerCAmelCase__					, lowerCAmelCase__					, lowerCAmelCase__	)

				hf_model.save_pretrained(lowerCAmelCase__	)


if __name__ == "__main__":
					SCREAMING_SNAKE_CASE_            =			argparse.ArgumentParser()
					parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
					parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to fairseq checkpoint')
					parser.add_argument('--dict_path', default=None, type=str, help='Path to dict of fine-tuned model')
					parser.add_argument('--config_path', default=None, type=str, help='Path to hf config.json of model to convert')
					parser.add_argument(
					    '--is_finetuned', action='store_true', help='Whether the model to convert is a fine-tuned model or not'
					)
					SCREAMING_SNAKE_CASE_            =			parser.parse_args()
					convert_sew_checkpoint(
					    args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, args.is_finetuned
					)





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, 2048, 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			):


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







										'''simple docstring'''


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





						def     _UpperCAmelCase			(    self  ,   A_		):







										'''simple docstring'''


										_UpperCAmelCase    :						List[Any]							=      self.convolution(A_		)
										_UpperCAmelCase    :						Optional[int]							=      self.normalization(A_		)
										_UpperCAmelCase    :						Optional[Any]							=      self.activation(A_		)
										return hidden_state







class 			a					(		nn.Module			):


						def __init__(    self  ,   A_		):







										'''simple docstring'''


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





						def     _UpperCAmelCase			(    self  ,   A_		):







										'''simple docstring'''


										_UpperCAmelCase    :						int							=      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    :						int							=      self.embedder(A_		)
										_UpperCAmelCase    :						int							=      self.pooler(A_		)
										return embedding







class 			a					(		nn.Module			):


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







										'''simple docstring'''


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





						def     _UpperCAmelCase			(    self  ,   A_		):







										'''simple docstring'''


										_UpperCAmelCase    :						str							=      self.convolution(A_		)
										_UpperCAmelCase    :						List[str]							=      self.normalization(A_		)
										return hidden_state







class 			a					(		nn.Module			):


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







										'''simple docstring'''


										super().__init__()
										_UpperCAmelCase    :						Optional[int]							=      in_channels != out_channels or stride != 1
										_UpperCAmelCase    :						Dict							=      (
										    ResNetShortCut(A_  ,   A_  ,   stride=A_		) if should_apply_shortcut else nn.Identity()
										)
										_UpperCAmelCase    :						int							=      nn.Sequential(
										    ResNetConvLayer(A_  ,   A_  ,   stride=A_		)  ,   ResNetConvLayer(A_  ,   A_  ,   activation=A_		)  ,   )
										_UpperCAmelCase    :						Dict							=      ACTaFN[activation]





						def     _UpperCAmelCase			(    self  ,   A_		):







										'''simple docstring'''


										_UpperCAmelCase    :						Optional[Any]							=      hidden_state
										_UpperCAmelCase    :						Any							=      self.layer(A_		)
										_UpperCAmelCase    :						Optional[int]							=      self.shortcut(A_		)
										hidden_state += residual
										_UpperCAmelCase    :						Optional[int]							=      self.activation(A_		)
										return hidden_state







class 			a					(		nn.Module			):


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







										'''simple docstring'''


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





						def     _UpperCAmelCase			(    self  ,   A_		):







										'''simple docstring'''


										_UpperCAmelCase    :						List[Any]							=      hidden_state
										_UpperCAmelCase    :						List[str]							=      self.layer(A_		)
										_UpperCAmelCase    :						List[str]							=      self.shortcut(A_		)
										hidden_state += residual
										_UpperCAmelCase    :						Dict							=      self.activation(A_		)
										return hidden_state







class 			a					(		nn.Module			):


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







										'''simple docstring'''


										super().__init__()

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

										_UpperCAmelCase    :						Optional[Any]							=      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     _UpperCAmelCase			(    self  ,   A_		):







										'''simple docstring'''


										_UpperCAmelCase    :						List[Any]							=      input
										for layer in self.layers:
														_UpperCAmelCase    :						Optional[Any]							=      layer(A_		)
										return hidden_state







class 			a					(		nn.Module			):


						def __init__(    self  ,   A_		):







										'''simple docstring'''


										super().__init__()
										_UpperCAmelCase    :						Any							=      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]  ,   )		)
										_UpperCAmelCase    :						Union[str, 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     _UpperCAmelCase			(    self  ,   A_  ,   A_ = False  ,   A_ = True		):







										'''simple docstring'''


										_UpperCAmelCase    :						List[Any]							=      () if output_hidden_states else None

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

														_UpperCAmelCase    :						str							=      stage_module(A_		)

										if output_hidden_states:
														_UpperCAmelCase    :						int							=      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 			a					(		UpperCAmelCase			):
						_lowercase        =     ResNetConfig
						_lowercase        =     "resnet"
						_lowercase        =     "pixel_values"
						_lowercase        =     True


						def     _UpperCAmelCase			(    self  ,   A_		):







										'''simple docstring'''


										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     _UpperCAmelCase			(    self  ,   A_  ,   A_=False		):







										'''simple docstring'''


										if isinstance(A_  ,   A_		):
														_UpperCAmelCase    :						Optional[Any]							=      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."							,	UpperCAmelCase							,	)
class 			a					(		UpperCAmelCase			):


						def __init__(    self  ,   A_		):







										'''simple docstring'''


										super().__init__(A_		)
										_UpperCAmelCase    :						List[str]							=      config
										_UpperCAmelCase    :						Any							=      ResNetEmbeddings(A_		)
										_UpperCAmelCase    :						str							=      ResNetEncoder(A_		)
										_UpperCAmelCase    :						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     _UpperCAmelCase			(    self  ,   A_  ,   A_ = None  ,   A_ = None		):







										'''simple docstring'''


										_UpperCAmelCase    :						Optional[int]							=      (
										    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
										)
										_UpperCAmelCase    :						List[str]							=      return_dict if return_dict is not None else self.config.use_return_dict

										_UpperCAmelCase    :						List[Any]							=      self.embedder(A_		)

										_UpperCAmelCase    :						str							=      self.encoder(
										    A_  ,   output_hidden_states=A_  ,   return_dict=A_		)

										_UpperCAmelCase    :						List[Any]							=      encoder_outputs[0]

										_UpperCAmelCase    :						int							=      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    "							,	UpperCAmelCase							,	)
class 			a					(		UpperCAmelCase			):


						def __init__(    self  ,   A_		):







										'''simple docstring'''


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







										'''simple docstring'''


										_UpperCAmelCase    :						Optional[int]							=      return_dict if return_dict is not None else self.config.use_return_dict

										_UpperCAmelCase    :						Tuple							=      self.resnet(A_  ,   output_hidden_states=A_  ,   return_dict=A_		)

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

										_UpperCAmelCase    :						int							=      self.classifier(A_		)

										_UpperCAmelCase    :						Dict							=      None

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

										if not return_dict:
														_UpperCAmelCase    :						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    "							,	UpperCAmelCase							,	)
class 			a					(		UpperCAmelCase							,	UpperCAmelCase			):


						def __init__(    self  ,   A_		):







										'''simple docstring'''


										super().__init__(A_		)
										super()._init_backbone(A_		)

										_UpperCAmelCase    :						Optional[int]							=      [config.embedding_size] + config.hidden_sizes
										_UpperCAmelCase    :						str							=      ResNetEmbeddings(A_		)
										_UpperCAmelCase    :						List[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     _UpperCAmelCase			(    self  ,   A_  ,   A_ = None  ,   A_ = None		):







										'''simple docstring'''


										_UpperCAmelCase    :						Tuple							=      return_dict if return_dict is not None else self.config.use_return_dict
										_UpperCAmelCase    :						Tuple							=      (
										    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
										)

										_UpperCAmelCase    :						Tuple							=      self.embedder(A_		)

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

										_UpperCAmelCase    :						Optional[int]							=      outputs.hidden_states

										_UpperCAmelCase    :						Any							=      ()
										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=A_  ,   hidden_states=outputs.hidden_states if output_hidden_states else None  ,   attentions=A_  ,   )


