"""simple docstring"""

import argparse
import json
from pathlib import Path

import torch
import torchaudio
from datasets import load_dataset
from huggingface_hub import hf_hub_download

from transformers import ASTConfig, ASTFeatureExtractor, ASTForAudioClassification
from transformers.utils import logging


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






def 	lowercase						(		a__      :						str	)		->  str:
					_UpperCamelCase 		=      ASTConfig()

					if "10-10" in model_name:
										pass
					elif "speech-commands" in model_name:
										_UpperCamelCase 		=      128
					elif "12-12" in model_name:
										_UpperCamelCase 		=      12
										_UpperCamelCase 		=      12
					elif "14-14" in model_name:
										_UpperCamelCase 		=      14
										_UpperCamelCase 		=      14
					elif "16-16" in model_name:
										_UpperCamelCase 		=      16
										_UpperCamelCase 		=      16
					else:
										raise ValueError('''Model not supported'''	)

					_UpperCamelCase 		=      '''huggingface/label-files'''
					if "speech-commands" in model_name:
										_UpperCamelCase 		=      35
										_UpperCamelCase 		=      '''speech-commands-v2-id2label.json'''
					else:
										_UpperCamelCase 		=      527
										_UpperCamelCase 		=      '''audioset-id2label.json'''

					_UpperCamelCase 		=      json.load(open(hf_hub_download(UpperCAmelCase__					,    UpperCAmelCase__					,    repo_type='''dataset'''	)					,    '''r'''	)	)
					_UpperCamelCase 		=      {int(UpperCAmelCase__	): v for k, v in idalabel.items()}
					_UpperCamelCase 		=      idalabel
					_UpperCamelCase 		=      {v: k for k, v in idalabel.items()}

					return config






def 	lowercase						(		a__      :						Any	)		->  Union[str, Any]:



					if "module.v" in name:
										_UpperCamelCase 		=      name.replace('''module.v'''					,    '''audio_spectrogram_transformer'''	)
					if "cls_token" in name:
										_UpperCamelCase 		=      name.replace('''cls_token'''					,    '''embeddings.cls_token'''	)
					if "dist_token" in name:
										_UpperCamelCase 		=      name.replace('''dist_token'''					,    '''embeddings.distillation_token'''	)
					if "pos_embed" in name:
										_UpperCamelCase 		=      name.replace('''pos_embed'''					,    '''embeddings.position_embeddings'''	)
					if "patch_embed.proj" in name:
										_UpperCamelCase 		=      name.replace('''patch_embed.proj'''					,    '''embeddings.patch_embeddings.projection'''	)
					# transformer blocks
					if "blocks" in name:
										_UpperCamelCase 		=      name.replace('''blocks'''					,    '''encoder.layer'''	)
					if "attn.proj" in name:
										_UpperCamelCase 		=      name.replace('''attn.proj'''					,    '''attention.output.dense'''	)
					if "attn" in name:
										_UpperCamelCase 		=      name.replace('''attn'''					,    '''attention.self'''	)
					if "norm1" in name:
										_UpperCamelCase 		=      name.replace('''norm1'''					,    '''layernorm_before'''	)
					if "norm2" in name:
										_UpperCamelCase 		=      name.replace('''norm2'''					,    '''layernorm_after'''	)
					if "mlp.fc1" in name:
										_UpperCamelCase 		=      name.replace('''mlp.fc1'''					,    '''intermediate.dense'''	)
					if "mlp.fc2" in name:
										_UpperCamelCase 		=      name.replace('''mlp.fc2'''					,    '''output.dense'''	)
					# final layernorm
					if "audio_spectrogram_transformer.norm" in name:
										_UpperCamelCase 		=      name.replace('''audio_spectrogram_transformer.norm'''					,    '''audio_spectrogram_transformer.layernorm'''	)
					# classifier head
					if "module.mlp_head.0" in name:
										_UpperCamelCase 		=      name.replace('''module.mlp_head.0'''					,    '''classifier.layernorm'''	)
					if "module.mlp_head.1" in name:
										_UpperCamelCase 		=      name.replace('''module.mlp_head.1'''					,    '''classifier.dense'''	)

					return name






def 	lowercase						(		a__      :						str					,    a__      :						List[Any]	)		->  Union[str, Any]:



					for key in orig_state_dict.copy().keys():
										_UpperCamelCase 		=      orig_state_dict.pop(UpperCAmelCase__	)

										if "qkv" in key:
															_UpperCamelCase 		=      key.split('''.'''	)
															_UpperCamelCase 		=      int(key_split[3]	)
															_UpperCamelCase 		=      config.hidden_size
															if "weight" in key:
																				_UpperCamelCase 		=      val[:dim, :]
																				_UpperCamelCase 		=      val[dim : dim * 2, :]
																				_UpperCamelCase 		=      val[-dim:, :]
															else:
																				_UpperCamelCase 		=      val[:dim]
																				_UpperCamelCase 		=      val[dim : dim * 2]
																				_UpperCamelCase 		=      val[-dim:]
										else:
															_UpperCamelCase 		=      val

					return orig_state_dict






def 	lowercase						(		a__      :						Optional[int]	)		->  List[str]:
					_UpperCamelCase 		=      [
					    '''module.v.head.weight''',
					    '''module.v.head.bias''',
					    '''module.v.head_dist.weight''',
					    '''module.v.head_dist.bias''',
					]
					for k in ignore_keys:
										state_dict.pop(UpperCAmelCase__					,    UpperCAmelCase__	)






@torch.no_grad()
def 	lowercase						(		a__      :						Union[str, Any]					,    a__      :						Tuple					,    a__      :						str=False	)		->  Dict:
					_UpperCamelCase 		=      get_audio_spectrogram_transformer_config(UpperCAmelCase__	)

					_UpperCamelCase 		=      {
					    '''ast-finetuned-audioset-10-10-0.4593''': (
					        '''https://www.dropbox.com/s/ca0b1v2nlxzyeb4/audioset_10_10_0.4593.pth?dl=1'''
					    ),
					    '''ast-finetuned-audioset-10-10-0.450''': (
					        '''https://www.dropbox.com/s/1tv0hovue1bxupk/audioset_10_10_0.4495.pth?dl=1'''
					    ),
					    '''ast-finetuned-audioset-10-10-0.448''': (
					        '''https://www.dropbox.com/s/6u5sikl4b9wo4u5/audioset_10_10_0.4483.pth?dl=1'''
					    ),
					    '''ast-finetuned-audioset-10-10-0.448-v2''': (
					        '''https://www.dropbox.com/s/kt6i0v9fvfm1mbq/audioset_10_10_0.4475.pth?dl=1'''
					    ),
					    '''ast-finetuned-audioset-12-12-0.447''': (
					        '''https://www.dropbox.com/s/snfhx3tizr4nuc8/audioset_12_12_0.4467.pth?dl=1'''
					    ),
					    '''ast-finetuned-audioset-14-14-0.443''': (
					        '''https://www.dropbox.com/s/z18s6pemtnxm4k7/audioset_14_14_0.4431.pth?dl=1'''
					    ),
					    '''ast-finetuned-audioset-16-16-0.442''': (
					        '''https://www.dropbox.com/s/mdsa4t1xmcimia6/audioset_16_16_0.4422.pth?dl=1'''
					    ),
					    '''ast-finetuned-speech-commands-v2''': (
					        '''https://www.dropbox.com/s/q0tbqpwv44pquwy/speechcommands_10_10_0.9812.pth?dl=1'''
					    ),
					}

					# load original state_dict
					_UpperCamelCase 		=      model_name_to_url[model_name]
					_UpperCamelCase 		=      torch.hub.load_state_dict_from_url(UpperCAmelCase__					,    map_location='''cpu'''	)
					# remove some keys
					remove_keys(UpperCAmelCase__	)
					# rename some keys
					_UpperCamelCase 		=      convert_state_dict(UpperCAmelCase__					,    UpperCAmelCase__	)

					# load 🤗 model
					_UpperCamelCase 		=      ASTForAudioClassification(UpperCAmelCase__	)
					model.eval()

					model.load_state_dict(UpperCAmelCase__	)

					# verify outputs on dummy input
					# source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
					_UpperCamelCase 		=      -4.267_7393 if '''speech-commands''' not in model_name else -6.84_5978
					_UpperCamelCase 		=      4.568_9974 if '''speech-commands''' not in model_name else 5.565_4526
					_UpperCamelCase 		=      1024 if '''speech-commands''' not in model_name else 128
					_UpperCamelCase 		=      ASTFeatureExtractor(mean=UpperCAmelCase__					,    std=UpperCAmelCase__					,    max_length=UpperCAmelCase__	)

					if "speech-commands" in model_name:
										_UpperCamelCase 		=      load_dataset('''speech_commands'''					,    '''v0.02'''					,    split='''validation'''	)
										_UpperCamelCase 		=      dataset[0]['''audio''']['''array''']
					else:
										_UpperCamelCase 		=      hf_hub_download(
										    repo_id='''nielsr/audio-spectogram-transformer-checkpoint'''					,    filename='''sample_audio.flac'''					,    repo_type='''dataset'''					,    )

										_UpperCamelCase		,				_UpperCamelCase 		=      torchaudio.load(UpperCAmelCase__	)
										_UpperCamelCase 		=      waveform.squeeze().numpy()

					_UpperCamelCase 		=      feature_extractor(UpperCAmelCase__					,    sampling_rate=16000					,    return_tensors='''pt'''	)

					# forward pass
					_UpperCamelCase 		=      model(**UpperCAmelCase__	)
					_UpperCamelCase 		=      outputs.logits

					if model_name == "ast-finetuned-audioset-10-10-0.4593":
										_UpperCamelCase 		=      torch.tensor([-0.8760, -7.0042, -8.6602]	)
					elif model_name == "ast-finetuned-audioset-10-10-0.450":
										_UpperCamelCase 		=      torch.tensor([-1.1986, -7.0903, -8.2718]	)
					elif model_name == "ast-finetuned-audioset-10-10-0.448":
										_UpperCamelCase 		=      torch.tensor([-2.6128, -8.0080, -9.4344]	)
					elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
										_UpperCamelCase 		=      torch.tensor([-1.5080, -7.4534, -8.8917]	)
					elif model_name == "ast-finetuned-audioset-12-12-0.447":
										_UpperCamelCase 		=      torch.tensor([-0.5050, -6.5833, -8.0843]	)
					elif model_name == "ast-finetuned-audioset-14-14-0.443":
										_UpperCamelCase 		=      torch.tensor([-0.3826, -7.0336, -8.2413]	)
					elif model_name == "ast-finetuned-audioset-16-16-0.442":
										_UpperCamelCase 		=      torch.tensor([-1.2113, -6.9101, -8.3470]	)
					elif model_name == "ast-finetuned-speech-commands-v2":
										_UpperCamelCase 		=      torch.tensor([6.1589, -8.0566, -8.7984]	)
					else:
										raise ValueError('''Unknown model name'''	)
					if not torch.allclose(logits[0, :3]					,    UpperCAmelCase__					,    atol=1e-4	):
										raise ValueError('''Logits don\'t match'''	)
					print('''Looks ok!'''	)

					if pytorch_dump_folder_path is not None:
										Path(UpperCAmelCase__	).mkdir(exist_ok=UpperCAmelCase__	)
										print(F'''Saving model {model_name} to {pytorch_dump_folder_path}'''	)
										model.save_pretrained(UpperCAmelCase__	)
										print(F'''Saving feature extractor to {pytorch_dump_folder_path}'''	)
										feature_extractor.save_pretrained(UpperCAmelCase__	)

					if push_to_hub:
										print('''Pushing model and feature extractor to the hub...'''	)
										model.push_to_hub(F'''MIT/{model_name}'''	)
										feature_extractor.push_to_hub(F'''MIT/{model_name}'''	)


if __name__ == "__main__":
		UpperCAmelCase    	=    argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    """--model_name""",
		    default="""ast-finetuned-audioset-10-10-0.4593""",
		    type=str,
		    help="""Name of the Audio Spectrogram Transformer model you\'d like to convert.""",
		)
		parser.add_argument(
		    """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
		)
		parser.add_argument(
		    """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
		)

		UpperCAmelCase    	=    parser.parse_args()
		convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





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 ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_outputs import (
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import logging
from .configuration_regnet import RegNetConfig


a_    :  Optional[int]						  =		logging.get_logger(__name__)

# General docstring
a_    :  List[str]						  =		'RegNetConfig'

# Base docstring
a_    :  Union[str, Any]						  =		'facebook/regnet-y-040'
a_    :  Optional[Any]						  =		[1, 1_0_8_8, 7, 7]

# Image classification docstring
a_    :  Dict						  =		'facebook/regnet-y-040'
a_    :  List[Any]						  =		'tabby, tabby cat'

a_    :  Union[str, Any]						  =		[
    'facebook/regnet-y-040',
    # See all regnet models at https://huggingface.co/models?filter=regnet
]





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



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






           '''simple docstring'''



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







     def     _a   (self					,      __a     ):






           '''simple docstring'''



           lowerCamelCase             =		self.convolution(__a     )
           lowerCamelCase             =		self.normalization(__a     )
           lowerCamelCase             =		self.activation(__a     )
           return hidden_state





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



     def __init__(self					,      __a     ):






           '''simple docstring'''



           super().__init__()
           lowerCamelCase             =		RegNetConvLayer(
               config.num_channels					,      config.embedding_size					,      kernel_size=3					,      stride=2					,      activation=config.hidden_act     )
           lowerCamelCase             =		config.num_channels







     def     _a   (self					,      __a     ):






           '''simple docstring'''



           lowerCamelCase             =		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             =		self.embedder(__a     )
           return hidden_state





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



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






           '''simple docstring'''



           super().__init__()
           lowerCamelCase             =		nn.Convad(__a					,      __a					,      kernel_size=1					,      stride=__a					,      bias=__a     )
           lowerCamelCase             =		nn.BatchNormad(__a     )







     def     _a   (self					,      __a     ):






           '''simple docstring'''



           lowerCamelCase             =		self.convolution(__a     )
           lowerCamelCase             =		self.normalization(__a     )
           return hidden_state





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



     def __init__(self					,      __a					,      __a     ):






           '''simple docstring'''



           super().__init__()

           lowerCamelCase             =		nn.AdaptiveAvgPoolad((1, 1)     )
           lowerCamelCase             =		nn.Sequential(
               nn.Convad(__a					,      __a					,      kernel_size=1     )					,      nn.ReLU()					,      nn.Convad(__a					,      __a					,      kernel_size=1     )					,      nn.Sigmoid()					,      )







     def     _a   (self					,      __a     ):






           '''simple docstring'''



           lowerCamelCase             =		self.pooler(__a     )
           lowerCamelCase             =		self.attention(__a     )
           lowerCamelCase             =		hidden_state * attention
           return hidden_state





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



     def __init__(self					,      __a					,      __a					,      __a					,      __a = 1     ):






           '''simple docstring'''



           super().__init__()
           lowerCamelCase             =		in_channels != out_channels or stride != 1
           lowerCamelCase             =		max(1					,      out_channels // config.groups_width     )
           lowerCamelCase             =		(
               RegNetShortCut(__a					,      __a					,      stride=__a     ) if should_apply_shortcut else nn.Identity()
           )
           lowerCamelCase             =		nn.Sequential(
               RegNetConvLayer(__a					,      __a					,      kernel_size=1					,      activation=config.hidden_act     )					,      RegNetConvLayer(__a					,      __a					,      stride=__a					,      groups=__a					,      activation=config.hidden_act     )					,      RegNetConvLayer(__a					,      __a					,      kernel_size=1					,      activation=__a     )					,      )
           lowerCamelCase             =		ACTaFN[config.hidden_act]







     def     _a   (self					,      __a     ):






           '''simple docstring'''



           lowerCamelCase             =		hidden_state
           lowerCamelCase             =		self.layer(__a     )
           lowerCamelCase             =		self.shortcut(__a     )
           hidden_state += residual
           lowerCamelCase             =		self.activation(__a     )
           return hidden_state





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



     def __init__(self					,      __a					,      __a					,      __a					,      __a = 1     ):






           '''simple docstring'''



           super().__init__()
           lowerCamelCase             =		in_channels != out_channels or stride != 1
           lowerCamelCase             =		max(1					,      out_channels // config.groups_width     )
           lowerCamelCase             =		(
               RegNetShortCut(__a					,      __a					,      stride=__a     ) if should_apply_shortcut else nn.Identity()
           )
           lowerCamelCase             =		nn.Sequential(
               RegNetConvLayer(__a					,      __a					,      kernel_size=1					,      activation=config.hidden_act     )					,      RegNetConvLayer(__a					,      __a					,      stride=__a					,      groups=__a					,      activation=config.hidden_act     )					,      RegNetSELayer(__a					,      reduced_channels=int(round(in_channels / 4     )     )     )					,      RegNetConvLayer(__a					,      __a					,      kernel_size=1					,      activation=__a     )					,      )
           lowerCamelCase             =		ACTaFN[config.hidden_act]







     def     _a   (self					,      __a     ):






           '''simple docstring'''



           lowerCamelCase             =		hidden_state
           lowerCamelCase             =		self.layer(__a     )
           lowerCamelCase             =		self.shortcut(__a     )
           hidden_state += residual
           lowerCamelCase             =		self.activation(__a     )
           return hidden_state





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



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






           '''simple docstring'''



           super().__init__()

           lowerCamelCase             =		RegNetXLayer if config.layer_type == "x" else RegNetYLayer

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







     def     _a   (self					,      __a     ):






           '''simple docstring'''



           lowerCamelCase             =		self.layers(__a     )
           return hidden_state





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



     def __init__(self					,      __a     ):






           '''simple docstring'''



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







     def     _a   (self					,      __a					,      __a = False					,      __a = True     ):






           '''simple docstring'''



           lowerCamelCase             =		() if output_hidden_states else None

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

                 lowerCamelCase             =		stage_module(__a     )

           if output_hidden_states:
                 lowerCamelCase             =		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  lowerCamelCase__	(    UpperCAmelCase_):
     """simple docstring"""



     _A      =      RegNetConfig
     _A      =      'regnet'
     _A      =      'pixel_values'
     _A      =      True







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






           '''simple docstring'''



           if isinstance(__a					,      __a     ):
                 lowerCamelCase             =		value



a_    :  Dict						  =		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 ([`RegNetConfig`]): 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_    :  Optional[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 [`~file_utils.ModelOutput`] instead of a plain tuple.\n'





@add_start_docstrings(
    'The bare RegNet model outputting raw features without any specific head on top.'							,     UpperCAmelCase_							,     )
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class  lowerCamelCase__	(    UpperCAmelCase_):
     """simple docstring"""



     def __init__(self					,      __a     ):






           '''simple docstring'''



           super().__init__(__a     )
           lowerCamelCase             =		config
           lowerCamelCase             =		RegNetEmbeddings(__a     )
           lowerCamelCase             =		RegNetEncoder(__a     )
           lowerCamelCase             =		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     _a   (self					,      __a					,      __a = None					,      __a = None     ):






           '''simple docstring'''



           lowerCamelCase             =		(
               output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
           )
           lowerCamelCase             =		return_dict if return_dict is not None else self.config.use_return_dict

           lowerCamelCase             =		self.embedder(__a     )

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

           lowerCamelCase             =		encoder_outputs[0]

           lowerCamelCase             =		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    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    '							,     UpperCAmelCase_							,     )
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class  lowerCamelCase__	(    UpperCAmelCase_):
     """simple docstring"""



     def __init__(self					,      __a     ):






           '''simple docstring'''



           super().__init__(__a     )
           lowerCamelCase             =		config.num_labels
           lowerCamelCase             =		RegNetModel(__a     )
           # classification head
           lowerCamelCase             =		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     _a   (self					,      __a = None					,      __a = None					,      __a = None					,      __a = None					,      ):






           '''simple docstring'''



           lowerCamelCase             =		return_dict if return_dict is not None else self.config.use_return_dict

           lowerCamelCase             =		self.regnet(__a					,      output_hidden_states=__a					,      return_dict=__a     )

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

           lowerCamelCase             =		self.classifier(__a     )

           lowerCamelCase             =		None

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

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

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