






'''simple docstring'''



import unittest

from transformers import AlbertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import (
     MODEL_FOR_PRETRAINING_MAPPING,
     AlbertForMaskedLM,
     AlbertForMultipleChoice,
     AlbertForPreTraining,
     AlbertForQuestionAnswering,
     AlbertForSequenceClassification,
     AlbertForTokenClassification,
     AlbertModel,
 )
 from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST
class 				__SCREAMING_SNAKE_CASE							:




      def __init__(	self			, lowerCamelCase			, lowerCamelCase=13			, lowerCamelCase=7			, lowerCamelCase=True			, lowerCamelCase=True			, lowerCamelCase=True			, lowerCamelCase=True			, lowerCamelCase=99			, lowerCamelCase=16			, lowerCamelCase=36			, lowerCamelCase=6			, lowerCamelCase=6			, lowerCamelCase=6			, lowerCamelCase=37			, lowerCamelCase="gelu"			, lowerCamelCase=0.1			, lowerCamelCase=0.1			, lowerCamelCase=512			, lowerCamelCase=16			, lowerCamelCase=2			, lowerCamelCase=0.02			, lowerCamelCase=3			, lowerCamelCase=4			, lowerCamelCase=None			, )				->Optional[Any]:



         '''simple docstring'''
         __a					  =     parent
         __a					  =     batch_size
         __a					  =     seq_length
         __a					  =     is_training
         __a					  =     use_input_mask
         __a					  =     use_token_type_ids
         __a					  =     use_labels
         __a					  =     vocab_size
         __a					  =     embedding_size
         __a					  =     hidden_size
         __a					  =     num_hidden_layers
         __a					  =     num_hidden_groups
         __a					  =     num_attention_heads
         __a					  =     intermediate_size
         __a					  =     hidden_act
         __a					  =     hidden_dropout_prob
         __a					  =     attention_probs_dropout_prob
         __a					  =     max_position_embeddings
         __a					  =     type_vocab_size
         __a					  =     type_sequence_label_size
         __a					  =     initializer_range
         __a					  =     num_labels
         __a					  =     num_choices
         __a					  =     scope




      def 			__UpperCamelCase				(	self   )				->Optional[Any]:



         '''simple docstring'''
         __a					  =     ids_tensor([self.batch_size, self.seq_length]			, self.vocab_size   )

         __a					  =     None
         if self.use_input_mask:
            __a					  =     random_attention_mask([self.batch_size, self.seq_length]   )

         __a					  =     None
         if self.use_token_type_ids:
            __a					  =     ids_tensor([self.batch_size, self.seq_length]			, self.type_vocab_size   )

         __a					  =     None
         __a					  =     None
         __a					  =     None
         if self.use_labels:
            __a					  =     ids_tensor([self.batch_size]			, self.type_sequence_label_size   )
            __a					  =     ids_tensor([self.batch_size, self.seq_length]			, self.num_labels   )
            __a					  =     ids_tensor([self.batch_size]			, self.num_choices   )

         __a					  =     self.get_config()

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




      def 			__UpperCamelCase				(	self   )				->List[Any]:



         '''simple docstring'''
         return AlbertConfig(
             vocab_size=self.vocab_size			, hidden_size=self.hidden_size			, num_hidden_layers=self.num_hidden_layers			, num_attention_heads=self.num_attention_heads			, intermediate_size=self.intermediate_size			, hidden_act=self.hidden_act			, hidden_dropout_prob=self.hidden_dropout_prob			, attention_probs_dropout_prob=self.attention_probs_dropout_prob			, max_position_embeddings=self.max_position_embeddings			, type_vocab_size=self.type_vocab_size			, initializer_range=self.initializer_range			, num_hidden_groups=self.num_hidden_groups			, )




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )				->Any:



         '''simple docstring'''
         __a					  =     AlbertModel(config=_A   )
         model.to(_A   )
         model.eval()
         __a					  =     model(_A			, attention_mask=_A			, token_type_ids=_A   )
         __a					  =     model(_A			, token_type_ids=_A   )
         __a					  =     model(_A   )
         self.parent.assertEqual(result.last_hidden_state.shape			, (self.batch_size, self.seq_length, self.hidden_size)   )
         self.parent.assertEqual(result.pooler_output.shape			, (self.batch_size, self.hidden_size)   )




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )				->Optional[Any]:



         '''simple docstring'''
         __a					  =     AlbertForPreTraining(config=_A   )
         model.to(_A   )
         model.eval()
         __a					  =     model(
             _A			, attention_mask=_A			, token_type_ids=_A			, labels=_A			, sentence_order_label=_A			, )
         self.parent.assertEqual(result.prediction_logits.shape			, (self.batch_size, self.seq_length, self.vocab_size)   )
         self.parent.assertEqual(result.sop_logits.shape			, (self.batch_size, config.num_labels)   )




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )				->Union[str, Any]:



         '''simple docstring'''
         __a					  =     AlbertForMaskedLM(config=_A   )
         model.to(_A   )
         model.eval()
         __a					  =     model(_A			, attention_mask=_A			, token_type_ids=_A			, labels=_A   )
         self.parent.assertEqual(result.logits.shape			, (self.batch_size, self.seq_length, self.vocab_size)   )




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )				->str:



         '''simple docstring'''
         __a					  =     AlbertForQuestionAnswering(config=_A   )
         model.to(_A   )
         model.eval()
         __a					  =     model(
             _A			, attention_mask=_A			, token_type_ids=_A			, start_positions=_A			, end_positions=_A			, )
         self.parent.assertEqual(result.start_logits.shape			, (self.batch_size, self.seq_length)   )
         self.parent.assertEqual(result.end_logits.shape			, (self.batch_size, self.seq_length)   )




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )				->Tuple:



         '''simple docstring'''
         __a					  =     self.num_labels
         __a					  =     AlbertForSequenceClassification(_A   )
         model.to(_A   )
         model.eval()
         __a					  =     model(_A			, attention_mask=_A			, token_type_ids=_A			, labels=_A   )
         self.parent.assertEqual(result.logits.shape			, (self.batch_size, self.num_labels)   )




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )				->Optional[Any]:



         '''simple docstring'''
         __a					  =     self.num_labels
         __a					  =     AlbertForTokenClassification(config=_A   )
         model.to(_A   )
         model.eval()
         __a					  =     model(_A			, attention_mask=_A			, token_type_ids=_A			, labels=_A   )
         self.parent.assertEqual(result.logits.shape			, (self.batch_size, self.seq_length, self.num_labels)   )




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase   )				->int:



         '''simple docstring'''
         __a					  =     self.num_choices
         __a					  =     AlbertForMultipleChoice(config=_A   )
         model.to(_A   )
         model.eval()
         __a					  =     input_ids.unsqueeze(1   ).expand(-1			, self.num_choices			, -1   ).contiguous()
         __a					  =     token_type_ids.unsqueeze(1   ).expand(-1			, self.num_choices			, -1   ).contiguous()
         __a					  =     input_mask.unsqueeze(1   ).expand(-1			, self.num_choices			, -1   ).contiguous()
         __a					  =     model(
             _A			, attention_mask=_A			, token_type_ids=_A			, labels=_A			, )
         self.parent.assertEqual(result.logits.shape			, (self.batch_size, self.num_choices)   )






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



         '''simple docstring'''
         __a					  =     self.prepare_config_and_inputs()
         (
             __a						
         )					  =     config_and_inputs
         __a					  =     {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
         return config, inputs_dict







@require_torch
class 				__SCREAMING_SNAKE_CASE							(				__SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE   ,   unittest.TestCase  ):
      __a	       =(
          (
              AlbertModel,
              AlbertForPreTraining,
              AlbertForMaskedLM,
              AlbertForMultipleChoice,
              AlbertForSequenceClassification,
              AlbertForTokenClassification,
              AlbertForQuestionAnswering,
          )
          if is_torch_available()
          else ()
      )
      __a	       =(
          {
              """feature-extraction""": AlbertModel,
              """fill-mask""": AlbertForMaskedLM,
              """question-answering""": AlbertForQuestionAnswering,
              """text-classification""": AlbertForSequenceClassification,
              """token-classification""": AlbertForTokenClassification,
              """zero-shot""": AlbertForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      __a	       =True




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase=False   )				->Optional[Any]:



         '''simple docstring'''
         __a					  =     super()._prepare_for_class(_A			, _A			, return_labels=_A   )

         if return_labels:
            if model_class in get_values(_A   ):
               __a					  =     torch.zeros(
                   (self.model_tester.batch_size, self.model_tester.seq_length)			, dtype=torch.long			, device=_A   )
               __a					  =     torch.zeros(
                   self.model_tester.batch_size			, dtype=torch.long			, device=_A   )
         return inputs_dict




      def 			__UpperCamelCase				(	self   )				->str:



         '''simple docstring'''
         __a					  =     AlbertModelTester(self   )
         __a					  =     ConfigTester(self			, config_class=_A			, hidden_size=37   )




      def 			__UpperCamelCase				(	self   )				->Any:



         '''simple docstring'''
         self.config_tester.run_common_tests()




      def 			__UpperCamelCase				(	self   )				->int:



         '''simple docstring'''
         __a					  =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*_A   )




      def 			__UpperCamelCase				(	self   )				->str:



         '''simple docstring'''
         __a					  =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_pretraining(*_A   )




      def 			__UpperCamelCase				(	self   )				->Tuple:



         '''simple docstring'''
         __a					  =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*_A   )




      def 			__UpperCamelCase				(	self   )				->Optional[Any]:



         '''simple docstring'''
         __a					  =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_multiple_choice(*_A   )




      def 			__UpperCamelCase				(	self   )				->List[Any]:



         '''simple docstring'''
         __a					  =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*_A   )




      def 			__UpperCamelCase				(	self   )				->Dict:



         '''simple docstring'''
         __a					  =     self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*_A   )




      def 			__UpperCamelCase				(	self   )				->List[Any]:



         '''simple docstring'''
         __a					  =     self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
            __a					  =     type
            self.model_tester.create_and_check_model(*_A   )






      @slow
      def 			__UpperCamelCase				(	self   )				->str:



         '''simple docstring'''







         for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            __a					  =     AlbertModel.from_pretrained(_A   )
            self.assertIsNotNone(_A   )







@require_torch
class 				__SCREAMING_SNAKE_CASE							(				unittest.TestCase  ):




      @slow
      def 			__UpperCamelCase				(	self   )				->List[Any]:



         '''simple docstring'''
         __a					  =     AlbertModel.from_pretrained('albert-base-v2'   )
         __a					  =     torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]   )
         __a					  =     torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]   )
         with torch.no_grad():
            __a					  =     model(_A			, attention_mask=_A   )[0]
         __a					  =     torch.Size((1, 11, 768)   )
         self.assertEqual(output.shape			, _A   )
         __a					  =     torch.tensor(
             [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]   )

         self.assertTrue(torch.allclose(output[:, 1:4, 1:4]			, _A			, atol=1e-4   )   )






"""simple docstring"""


import argparse
from pathlib import Path
from typing import Dict, OrderedDict, Tuple

import torch
from audiocraft.models import MusicGen

from transformers import (
    AutoFeatureExtractor,
    AutoTokenizer,
    EncodecModel,
    MusicgenDecoderConfig,
    MusicgenForConditionalGeneration,
    MusicgenProcessor,
    TaEncoderModel,
)
from transformers.models.musicgen.modeling_musicgen import MusicgenForCausalLM
from transformers.utils import logging


logging.set_verbosity_info()
__magic_name__       :					List[Any]						=			logging.get_logger(__name__)


__magic_name__       :					List[str]						=			["""model.decoder.embed_positions.weights"""]

def      UpperCamelCase							(SCREAMING_SNAKE_CASE     ):
						if "emb" in name:
												UpperCamelCase				:		Tuple				   =			name.replace("""emb"""							,     """model.decoder.embed_tokens"""     )
						if "transformer" in name:
												UpperCamelCase				:		List[str]				   =			name.replace("""transformer"""							,     """model.decoder"""     )
						if "cross_attention" in name:
												UpperCamelCase				:		str				   =			name.replace("""cross_attention"""							,     """encoder_attn"""     )
						if "linear1" in name:
												UpperCamelCase				:		Optional[int]				   =			name.replace("""linear1"""							,     """fc1"""     )
						if "linear2" in name:
												UpperCamelCase				:		int				   =			name.replace("""linear2"""							,     """fc2"""     )
						if "norm1" in name:
												UpperCamelCase				:		List[Any]				   =			name.replace("""norm1"""							,     """self_attn_layer_norm"""     )
						if "norm_cross" in name:
												UpperCamelCase				:		int				   =			name.replace("""norm_cross"""							,     """encoder_attn_layer_norm"""     )
						if "norm2" in name:
												UpperCamelCase				:		str				   =			name.replace("""norm2"""							,     """final_layer_norm"""     )
						if "out_norm" in name:
												UpperCamelCase				:		Tuple				   =			name.replace("""out_norm"""							,     """model.decoder.layer_norm"""     )
						if "linears" in name:
												UpperCamelCase				:		str				   =			name.replace("""linears"""							,     """lm_heads"""     )
						if "condition_provider.conditioners.description.output_proj" in name:
												UpperCamelCase				:		Dict				   =			name.replace("""condition_provider.conditioners.description.output_proj"""							,     """enc_to_dec_proj"""     )
						return name

def      UpperCamelCase							(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE     ):
						UpperCamelCase				:		List[str]				   =			list(state_dict.keys()     )
						UpperCamelCase				:		Optional[Any]				   =			{}
						for key in keys:
												UpperCamelCase				:		Optional[int]				   =			state_dict.pop(SCREAMING_SNAKE_CASE     )
												UpperCamelCase				:		Dict				   =			rename_keys(SCREAMING_SNAKE_CASE     )
												if "in_proj_weight" in key:
																		# split fused qkv proj
																		UpperCamelCase				:		Optional[int]				   =			val[:hidden_size, :]
																		UpperCamelCase				:		List[Any]				   =			val[hidden_size : 2 * hidden_size, :]
																		UpperCamelCase				:		Dict				   =			val[-hidden_size:, :]
												elif "enc_to_dec_proj" in key:
																		UpperCamelCase				:		Any				   =			val
												else:
																		UpperCamelCase				:		str				   =			val
						return state_dict, enc_dec_proj_state_dict

def      UpperCamelCase							(SCREAMING_SNAKE_CASE     ):
						if checkpoint == "small":
												# default config values
												UpperCamelCase				:		Tuple				   =			1024
												UpperCamelCase				:		Dict				   =			24
												UpperCamelCase				:		Optional[int]				   =			16
						elif checkpoint == "medium":
												UpperCamelCase				:		Optional[Any]				   =			1536
												UpperCamelCase				:		int				   =			48
												UpperCamelCase				:		int				   =			24
						elif checkpoint == "large":
												UpperCamelCase				:		Any				   =			2048
												UpperCamelCase				:		List[Any]				   =			48
												UpperCamelCase				:		str				   =			32
						else:
												raise ValueError(f"""Checkpoint should be one of `['small', 'medium', 'large']`, got {checkpoint}."""     )
						UpperCamelCase				:		Optional[int]				   =			MusicgenDecoderConfig(
						    hidden_size=SCREAMING_SNAKE_CASE							,     ffn_dim=hidden_size * 4							,     num_hidden_layers=SCREAMING_SNAKE_CASE							,     num_attention_heads=SCREAMING_SNAKE_CASE							,     )
						return config

@torch.no_grad()
def      UpperCamelCase							(SCREAMING_SNAKE_CASE							,     SCREAMING_SNAKE_CASE=None							,     SCREAMING_SNAKE_CASE=None							,     SCREAMING_SNAKE_CASE="cpu"     ):
						UpperCamelCase				:		Optional[Any]				   =			MusicGen.get_pretrained(SCREAMING_SNAKE_CASE							,     device=SCREAMING_SNAKE_CASE     )
						UpperCamelCase				:		Tuple				   =			decoder_config_from_checkpoint(SCREAMING_SNAKE_CASE     )

						UpperCamelCase				:		Optional[Any]				   =			fairseq_model.lm.state_dict()
						UpperCamelCase   ,						UpperCamelCase				:		Optional[int]				   =			rename_state_dict(
						    SCREAMING_SNAKE_CASE							,     hidden_size=decoder_config.hidden_size     )

						UpperCamelCase				:		int				   =			TaEncoderModel.from_pretrained("""t5-base"""     )
						UpperCamelCase				:		List[str]				   =			EncodecModel.from_pretrained("""facebook/encodec_32khz"""     )
						UpperCamelCase				:		Optional[int]				   =			MusicgenForCausalLM(SCREAMING_SNAKE_CASE     ).eval()

						# load all decoder weights - expect that we'll be missing embeddings and enc-dec projection
						UpperCamelCase   ,						UpperCamelCase				:		Tuple				   =			decoder.load_state_dict(SCREAMING_SNAKE_CASE							,     strict=SCREAMING_SNAKE_CASE     )

						for key in missing_keys.copy():
												if key.startswith(("""text_encoder""", """audio_encoder""")     ) or key in EXPECTED_MISSING_KEYS:
																		missing_keys.remove(SCREAMING_SNAKE_CASE     )

						if len(SCREAMING_SNAKE_CASE     ) > 0:
												raise ValueError(f"""Missing key(s) in state_dict: {missing_keys}"""     )

						if len(SCREAMING_SNAKE_CASE     ) > 0:
												raise ValueError(f"""Unexpected key(s) in state_dict: {unexpected_keys}"""     )

						# init the composite model
						UpperCamelCase				:		Tuple				   =			MusicgenForConditionalGeneration(text_encoder=SCREAMING_SNAKE_CASE							,     audio_encoder=SCREAMING_SNAKE_CASE							,     decoder=SCREAMING_SNAKE_CASE     )

						# load the pre-trained enc-dec projection (from the decoder state dict)
						model.enc_to_dec_proj.load_state_dict(SCREAMING_SNAKE_CASE     )

						# check we can do a forward pass
						UpperCamelCase				:		int				   =			torch.arange(0							,     8							,     dtype=torch.long     ).reshape(2							,     -1     )
						UpperCamelCase				:		Tuple				   =			input_ids.reshape(2 * 4							,     -1     )

						with torch.no_grad():
												UpperCamelCase				:		Any				   =			model(input_ids=SCREAMING_SNAKE_CASE							,     decoder_input_ids=SCREAMING_SNAKE_CASE     ).logits

						if logits.shape != (8, 1, 2048):
												raise ValueError("""Incorrect shape for logits"""     )

						# now construct the processor
						UpperCamelCase				:		Any				   =			AutoTokenizer.from_pretrained("""t5-base"""     )
						UpperCamelCase				:		Tuple				   =			AutoFeatureExtractor.from_pretrained("""facebook/encodec_32khz"""							,     padding_side="""left"""     )

						UpperCamelCase				:		int				   =			MusicgenProcessor(feature_extractor=SCREAMING_SNAKE_CASE							,     tokenizer=SCREAMING_SNAKE_CASE     )

						# set the appropriate bos/pad token ids
						UpperCamelCase				:		Optional[int]				   =			2048
						UpperCamelCase				:		Tuple				   =			2048

						# set other default generation config params
						UpperCamelCase				:		Optional[int]				   =			int(30 * audio_encoder.config.frame_rate     )
						UpperCamelCase				:		int				   =			True
						UpperCamelCase				:		Dict				   =			3.0

						if pytorch_dump_folder is not None:
												Path(SCREAMING_SNAKE_CASE     ).mkdir(exist_ok=SCREAMING_SNAKE_CASE     )
												logger.info(f"""Saving model {checkpoint} to {pytorch_dump_folder}"""     )
												model.save_pretrained(SCREAMING_SNAKE_CASE     )
												processor.save_pretrained(SCREAMING_SNAKE_CASE     )

						if repo_id:
												logger.info(f"""Pushing model {checkpoint} to {repo_id}"""     )
												model.push_to_hub(SCREAMING_SNAKE_CASE     )
												processor.push_to_hub(SCREAMING_SNAKE_CASE     )


if __name__ == "__main__":
		__magic_name__       :					Any						=			argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    """--checkpoint""",
		    default="""small""",
		    type=str,
		    help="""Checkpoint size of the MusicGen model you'd like to convert. Can be one of: `['small', 'medium', 'large']`.""",
		)
		parser.add_argument(
		    """--pytorch_dump_folder""",
		    required=True,
		    default=None,
		    type=str,
		    help="""Path to the output PyTorch model directory.""",
		)
		parser.add_argument(
		    """--push_to_hub""", default=None, type=str, help="""Where to upload the converted model on the 🤗 hub."""
		)
		parser.add_argument(
		    """--device""", default="""cpu""", type=str, help="""Torch device to run the conversion, either cpu or cuda."""
		)

		__magic_name__       :					Optional[Any]						=			parser.parse_args()
		convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)



