





'''simple docstring'''



from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...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


_SCREAMING_SNAKE_CASE					=    logging.get_logger(__name__)

# General docstring
_SCREAMING_SNAKE_CASE					=    '''RegNetConfig'''

# Base docstring
_SCREAMING_SNAKE_CASE					=    '''facebook/regnet-y-040'''
_SCREAMING_SNAKE_CASE					=    [1, 1_0_8_8, 7, 7]

# Image classification docstring
_SCREAMING_SNAKE_CASE					=    '''facebook/regnet-y-040'''
_SCREAMING_SNAKE_CASE					=    '''tabby, tabby cat'''

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






class 							__lowercase       (      nn.Module						):







	'''simple docstring'''

	def __init__(self     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase = 3     ,_lowerCamelCase = 1     ,_lowerCamelCase = 1     ,_lowerCamelCase = "relu"     ,)							->							str:







				'''simple docstring'''





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


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase )							->							int:







				'''simple docstring'''





				__lowercase										=							self.convolution(_lowerCamelCase )
				__lowercase										=							self.normalization(_lowerCamelCase )
				__lowercase										=							self.activation(_lowerCamelCase )
				return hidden_state






class 							__lowercase       (      nn.Module						):







	'''simple docstring'''

	def __init__(self     ,_lowerCamelCase )							->							Union[str, Any]:







				'''simple docstring'''





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


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase )							->							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 )
				return hidden_state






class 							__lowercase       (      nn.Module						):







	'''simple docstring'''

	def __init__(self     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase = 2 )							->							Tuple:







				'''simple docstring'''





				super().__init__()
				__lowercase										=							nn.Convad(_lowerCamelCase     ,_lowerCamelCase     ,kernel_size=1     ,stride=_lowerCamelCase     ,bias=_lowerCamelCase )
				__lowercase										=							nn.BatchNormad(_lowerCamelCase )


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase )							->							Tensor:







				'''simple docstring'''





				__lowercase										=							self.convolution(_lowerCamelCase )
				__lowercase										=							self.normalization(_lowerCamelCase )
				return hidden_state






class 							__lowercase       (      nn.Module						):







	'''simple docstring'''

	def __init__(self     ,_lowerCamelCase     ,_lowerCamelCase )							->							Any:







				'''simple docstring'''





				super().__init__()

				__lowercase										=							nn.AdaptiveAvgPoolad((1, 1) )
				__lowercase										=							nn.Sequential(
				    nn.Convad(_lowerCamelCase     ,_lowerCamelCase     ,kernel_size=1 )     ,nn.ReLU()     ,nn.Convad(_lowerCamelCase     ,_lowerCamelCase     ,kernel_size=1 )     ,nn.Sigmoid()     ,)


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase )							->							Dict:







				'''simple docstring'''





				__lowercase										=							self.pooler(_lowerCamelCase )
				__lowercase										=							self.attention(_lowerCamelCase )
				__lowercase										=							hidden_state * attention
				return hidden_state






class 							__lowercase       (      nn.Module						):







	'''simple docstring'''

	def __init__(self     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase = 1 )							->							List[str]:







				'''simple docstring'''





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


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase )							->							int:







				'''simple docstring'''





				__lowercase										=							hidden_state
				__lowercase										=							self.layer(_lowerCamelCase )
				__lowercase										=							self.shortcut(_lowerCamelCase )
				hidden_state += residual
				__lowercase										=							self.activation(_lowerCamelCase )
				return hidden_state






class 							__lowercase       (      nn.Module						):







	'''simple docstring'''

	def __init__(self     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase = 1 )							->							Tuple:







				'''simple docstring'''





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


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase )							->							str:







				'''simple docstring'''





				__lowercase										=							hidden_state
				__lowercase										=							self.layer(_lowerCamelCase )
				__lowercase										=							self.shortcut(_lowerCamelCase )
				hidden_state += residual
				__lowercase										=							self.activation(_lowerCamelCase )
				return hidden_state






class 							__lowercase       (      nn.Module						):







	'''simple docstring'''

	def __init__(self     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase = 2     ,_lowerCamelCase = 2     ,)							->							List[str]:







				'''simple docstring'''





				super().__init__()

				__lowercase										=							RegNetXLayer if config.layer_type == '''x''' else RegNetYLayer

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


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase )							->							List[Any]:







				'''simple docstring'''





				__lowercase										=							self.layers(_lowerCamelCase )
				return hidden_state






class 							__lowercase       (      nn.Module						):







	'''simple docstring'''

	def __init__(self     ,_lowerCamelCase )							->							Union[str, 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(
				    RegNetStage(
				        _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(RegNetStage(_lowerCamelCase     ,_lowerCamelCase     ,_lowerCamelCase     ,depth=_lowerCamelCase ) )


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase     ,_lowerCamelCase = False     ,_lowerCamelCase = True )							->							BaseModelOutputWithNoAttention:







				'''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 							__lowercase       (      lowerCAmelCase__						):







	'''simple docstring'''

	a    :				List[str]											=       RegNetConfig
	a    :				str											=       "regnet"
	a    :				List[str]											=       "pixel_values"
	a    :				str											=       True


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase )							->							int:







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







				'''simple docstring'''





				if isinstance(_lowerCamelCase     ,_lowerCamelCase ):
							__lowercase										=							value





_SCREAMING_SNAKE_CASE					=    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 ([`RegNetConfig`]): 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.
'''

_SCREAMING_SNAKE_CASE					=    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 [`~file_utils.ModelOutput`] instead of a plain tuple.
'''






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







	'''simple docstring'''

	def __init__(self     ,_lowerCamelCase )							->							List[str]:







				'''simple docstring'''





				super().__init__(_lowerCamelCase )
				__lowercase										=							config
				__lowercase										=							RegNetEmbeddings(_lowerCamelCase )
				__lowercase										=							RegNetEncoder(_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 						_UpperCAmelCase     (self     ,_lowerCamelCase     ,_lowerCamelCase = None     ,_lowerCamelCase = None )							->							BaseModelOutputWithPoolingAndNoAttention:







				'''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(
    "\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    " ,      lowerCAmelCase__ ,      )
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class 							__lowercase       (      lowerCAmelCase__						):







	'''simple docstring'''

	def __init__(self     ,_lowerCamelCase )							->							Tuple:







				'''simple docstring'''





				super().__init__(_lowerCamelCase )
				__lowercase										=							config.num_labels
				__lowercase										=							RegNetModel(_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 						_UpperCAmelCase     (self     ,_lowerCamelCase = None     ,_lowerCamelCase = None     ,_lowerCamelCase = None     ,_lowerCamelCase = None     ,)							->							ImageClassifierOutputWithNoAttention:







				'''simple docstring'''





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

				__lowercase										=							self.regnet(_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 )












'''simple docstring'''



import argparse

import torch

from transformers import (
    UniSpeechSatConfig,
    UniSpeechSatForAudioFrameClassification,
    UniSpeechSatForSequenceClassification,
    UniSpeechSatForXVector,
    WavaVecaFeatureExtractor,
    logging,
)


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




def    _lowerCAmelCase						(				lowerCamelCase_			:  List[str]		,       lowerCamelCase_			:  List[str]		,       lowerCamelCase_			:  List[str]				):
			__lowercase										=							UniSpeechSatForSequenceClassification.from_pretrained(lowerCamelCase_		,       config=lowerCamelCase_				)
			__lowercase										=							downstream_dict['''projector.weight''']
			__lowercase										=							downstream_dict['''projector.bias''']
			__lowercase										=							downstream_dict['''model.post_net.linear.weight''']
			__lowercase										=							downstream_dict['''model.post_net.linear.bias''']
			return model




def    _lowerCAmelCase						(				lowerCamelCase_			:  str		,       lowerCamelCase_			:  Optional[int]		,       lowerCamelCase_			:  Optional[Any]				):
			__lowercase										=							UniSpeechSatForAudioFrameClassification.from_pretrained(lowerCamelCase_		,       config=lowerCamelCase_				)
			__lowercase										=							downstream_dict['''model.linear.weight''']
			__lowercase										=							downstream_dict['''model.linear.bias''']
			return model




def    _lowerCAmelCase						(				lowerCamelCase_			:  Optional[Any]		,       lowerCamelCase_			:  Dict		,       lowerCamelCase_			:  Optional[int]				):
			__lowercase										=							UniSpeechSatForXVector.from_pretrained(lowerCamelCase_		,       config=lowerCamelCase_				)
			__lowercase										=							downstream_dict['''connector.weight''']
			__lowercase										=							downstream_dict['''connector.bias''']
			for i, kernel_size in enumerate(hf_config.tdnn_kernel				):
						__lowercase										=							downstream_dict[
						    f"model.framelevel_feature_extractor.module.{i}.kernel.weight"
						]
						__lowercase										=							downstream_dict[f"model.framelevel_feature_extractor.module.{i}.kernel.bias"]

			__lowercase										=							downstream_dict['''model.utterancelevel_feature_extractor.linear1.weight''']
			__lowercase										=							downstream_dict['''model.utterancelevel_feature_extractor.linear1.bias''']
			__lowercase										=							downstream_dict['''model.utterancelevel_feature_extractor.linear2.weight''']
			__lowercase										=							downstream_dict['''model.utterancelevel_feature_extractor.linear2.bias''']
			__lowercase										=							downstream_dict['''objective.W''']
			return model




@torch.no_grad()
def    _lowerCAmelCase						(				lowerCamelCase_			:  Tuple		,       lowerCamelCase_			:  List[str]		,       lowerCamelCase_			:  Dict		,       lowerCamelCase_			:  Optional[int]				):
			__lowercase										=							torch.load(lowerCamelCase_		,       map_location='''cpu'''				)

			__lowercase										=							checkpoint['''Downstream''']

			__lowercase										=							UniSpeechSatConfig.from_pretrained(lowerCamelCase_				)
			__lowercase										=							WavaVecaFeatureExtractor.from_pretrained(
			    lowerCamelCase_		,       return_attention_mask=lowerCamelCase_		,       do_normalize=lowerCamelCase_				)

			__lowercase										=							hf_config.architectures[0]
			if arch.endswith('''ForSequenceClassification'''				):
						__lowercase										=							convert_classification(lowerCamelCase_		,       lowerCamelCase_		,       lowerCamelCase_				)
			elif arch.endswith('''ForAudioFrameClassification'''				):
						__lowercase										=							convert_diarization(lowerCamelCase_		,       lowerCamelCase_		,       lowerCamelCase_				)
			elif arch.endswith('''ForXVector'''				):
						__lowercase										=							convert_xvector(lowerCamelCase_		,       lowerCamelCase_		,       lowerCamelCase_				)
			else:
						raise NotImplementedError(f"S3PRL weights conversion is not supported for {arch}"				)

			if hf_config.use_weighted_layer_sum:
						__lowercase										=							checkpoint['''Featurizer''']['''weights''']

			hf_feature_extractor.save_pretrained(lowerCamelCase_				)
			hf_model.save_pretrained(lowerCamelCase_				)


if __name__ == "__main__":
	_SCREAMING_SNAKE_CASE					=    argparse.ArgumentParser()
	parser.add_argument(
	    '''--base_model_name''', default=None, type=str, help='''Name of the huggingface pretrained base model.'''
	)
	parser.add_argument('''--config_path''', default=None, type=str, help='''Path to the huggingface classifier config.''')
	parser.add_argument('''--checkpoint_path''', default=None, type=str, help='''Path to the s3prl checkpoint.''')
	parser.add_argument('''--model_dump_path''', default=None, type=str, help='''Path to the final converted model.''')
	_SCREAMING_SNAKE_CASE					=    parser.parse_args()
	convert_saprl_checkpoint(args.base_model_name, args.config_path, args.checkpoint_path, args.model_dump_path)





