


"""simple docstring"""






import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import AutoConfig


lowercase  		=			logging.get_logger(__name__)

lowercase  		=			{
    """facebook/musicgen-small""": """https://huggingface.co/facebook/musicgen-small/resolve/main/config.json""",
    # See all Musicgen models at https://huggingface.co/models?filter=musicgen
}


class 			SCREAMING_SNAKE_CASE_  (   _lowercase):




							'''simple docstring'''





							__magic_name__					:     Optional[Any]				  =  '''musicgen_decoder'''
							__magic_name__					:     Union[str, Any]				  =  ['''past_key_values''']




							def __init__(						self	,					lowerCamelCase__=2_048	,					lowerCamelCase__=2_048	,					lowerCamelCase__=24	,					lowerCamelCase__=4_096	,					lowerCamelCase__=16	,					lowerCamelCase__=0.0	,					lowerCamelCase__=True	,					lowerCamelCase__="gelu"	,					lowerCamelCase__=1_024	,					lowerCamelCase__=0.1	,					lowerCamelCase__=0.0	,					lowerCamelCase__=0.0	,					lowerCamelCase__=0.02	,					lowerCamelCase__=False	,					lowerCamelCase__=4	,					lowerCamelCase__=2_048	,					lowerCamelCase__=2_048	,					lowerCamelCase__=None	,					lowerCamelCase__=False	,					**lowerCamelCase__	,					)   ->	int:







								'''simple docstring'''




								snake_case__  :						Optional[int]		=      vocab_size
								snake_case__  :						Optional[Any]		=      max_position_embeddings
								snake_case__  :						List[str]		=      hidden_size
								snake_case__  :						str		=      ffn_dim
								snake_case__  :						Optional[Any]		=      num_hidden_layers
								snake_case__  :						Union[str, Any]		=      num_attention_heads
								snake_case__  :						int		=      dropout
								snake_case__  :						Union[str, Any]		=      attention_dropout
								snake_case__  :						List[Any]		=      activation_dropout
								snake_case__  :						Union[str, Any]		=      activation_function
								snake_case__  :						int		=      initializer_factor
								snake_case__  :						int		=      layerdrop
								snake_case__  :						Union[str, Any]		=      use_cache
								snake_case__  :						str		=      scale_embedding  # scale factor will be sqrt(d_model) if True
								snake_case__  :						Tuple		=      num_codebooks
								super().__init__(
								    pad_token_id=lowerCamelCase__	,					bos_token_id=lowerCamelCase__	,					eos_token_id=lowerCamelCase__	,					tie_word_embeddings=lowerCamelCase__	,					**lowerCamelCase__	,					)


class 			SCREAMING_SNAKE_CASE_  (   _lowercase):




							'''simple docstring'''





							__magic_name__					:     str				  =  '''musicgen'''
							__magic_name__					:     Union[str, Any]				  =  True




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







								'''simple docstring'''




								super().__init__(**lowerCamelCase__)
								if "text_encoder" not in kwargs or "audio_encoder" not in kwargs or "decoder" not in kwargs:
									raise ValueError("Config has to be initialized with text_encoder, audio_encoder and decoder config")

								snake_case__  :						str		=      kwargs.pop("text_encoder")
								snake_case__  :						Optional[int]		=      text_encoder_config.pop("model_type")

								snake_case__  :						Dict		=      kwargs.pop("audio_encoder")
								snake_case__  :						List[str]		=      audio_encoder_config.pop("model_type")

								snake_case__  :						str		=      kwargs.pop("decoder")

								snake_case__  :						Tuple		=      AutoConfig.for_model(lowerCamelCase__	,					**lowerCamelCase__)
								snake_case__  :						Dict		=      AutoConfig.for_model(lowerCamelCase__	,					**lowerCamelCase__)
								snake_case__  :						Optional[int]		=      MusicgenDecoderConfig(**lowerCamelCase__)
								snake_case__  :						List[str]		=      True




							@classmethod
							def        UpperCAmelCase (						cls	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					**lowerCamelCase__	,					)   ->	Optional[int]:







								'''simple docstring'''




								return cls(
								    text_encoder=text_encoder_config.to_dict()	,					audio_encoder=audio_encoder_config.to_dict()	,					decoder=decoder_config.to_dict()	,					**lowerCamelCase__	,					)







							def        UpperCAmelCase (						self)   ->	Optional[int]:







								'''simple docstring'''




								snake_case__  :						Dict		=      copy.deepcopy(self.__dict__)
								snake_case__  :						Tuple		=      self.text_encoder.to_dict()
								snake_case__  :						Any		=      self.audio_encoder.to_dict()
								snake_case__  :						int		=      self.decoder.to_dict()
								snake_case__  :						List[str]		=      self.__class__.model_type
								return output













import argparse

import torch

from transformers import ElectraConfig, ElectraForMaskedLM, ElectraForPreTraining, load_tf_weights_in_electra
from transformers.utils import logging


logging.set_verbosity_info()






def    UpperCamelCase__  (   UpperCAmelCase    ,	UpperCAmelCase    ,	UpperCAmelCase    ,	UpperCAmelCase     )     ->		str:







							"""simple docstring"""





							_a  :    Union[str, Any]			=						ElectraConfig.from_json_file(UpperCAmelCase     )
							print(F'Building PyTorch model from configuration: {config}'     )

							if discriminator_or_generator == "discriminator":
														_a  :    int			=						ElectraForPreTraining(UpperCAmelCase     )
							elif discriminator_or_generator == "generator":
														_a  :    Any			=						ElectraForMaskedLM(UpperCAmelCase     )
							else:
														raise ValueError('''The discriminator_or_generator argument should be either \'discriminator\' or \'generator\''''     )

							# Load weights from tf checkpoint
							load_tf_weights_in_electra(
							    UpperCAmelCase    ,	UpperCAmelCase    ,	UpperCAmelCase    ,	discriminator_or_generator=UpperCAmelCase     )

							# Save pytorch-model
							print(F'Save PyTorch model to {pytorch_dump_path}'     )
							torch.save(model.state_dict()    ,	UpperCAmelCase     )


if __name__ == "__main__":
						__lowerCamelCase						= argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '--tf_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
						)
						parser.add_argument(
						    '--config_file',
						    default=None,
						    type=str,
						    required=True,
						    help='The config json file corresponding to the pre-trained model. \nThis specifies the model architecture.',
						)
						parser.add_argument(
						    '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
						)
						parser.add_argument(
						    '--discriminator_or_generator',
						    default=None,
						    type=str,
						    required=True,
						    help=(
						        'Whether to export the generator or the discriminator. Should be a string, either \'discriminator\' or '
						        '\'generator\'.'
						    ),
						)
						__lowerCamelCase						= parser.parse_args()
						convert_tf_checkpoint_to_pytorch(
						    args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path, args.discriminator_or_generator
						)