"""simple docstring"""






import unittest

from transformers import (
    MODEL_FOR_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
    TextGenerationPipeline,
    logging,
    pipeline,
)
from transformers.testing_utils import (
    CaptureLogger,
    is_pipeline_test,
    require_accelerate,
    require_tf,
    require_torch,
    require_torch_gpu,
    require_torch_or_tf,
)

from .test_pipelines_common import ANY
@is_pipeline_test
@require_torch_or_tf
class        snake_case (				unittest.TestCase       ):




					"""simple docstring"""

					snake_case__       					=				MODEL_FOR_CAUSAL_LM_MAPPING
					snake_case__       					=				TF_MODEL_FOR_CAUSAL_LM_MAPPING







					@require_torch
					def 	__lowerCAmelCase		(       self     :			str    ):
									UpperCAmelCase__    =							pipeline(task='text-generation'					,model='sshleifer/tiny-ctrl'					,framework='pt'    )
									# Using `do_sample=False` to force deterministic output
									UpperCAmelCase__    =							text_generator('This is a test'					,do_sample=lowerCamelCase__    )
									self.assertEqual(
									    lowerCamelCase__					,[
									        {
									            'generated_text': (
									                'This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope.'
									                ' oscope. FiliFili@@'
									            )
									        }
									    ]					,)

									UpperCAmelCase__    =							text_generator(['This is a test', 'This is a second test']    )
									self.assertEqual(
									    lowerCamelCase__					,[
									        [
									            {
									                'generated_text': (
									                    'This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope.'
									                    ' oscope. FiliFili@@'
									                )
									            }
									        ],
									        [
									            {
									                'generated_text': (
									                    'This is a second test ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy'
									                    ' oscope. oscope. FiliFili@@'
									                )
									            }
									        ],
									    ]					,)

									UpperCAmelCase__    =							text_generator('This is a test'					,do_sample=lowerCamelCase__					,num_return_sequences=2					,return_tensors=lowerCamelCase__    )
									self.assertEqual(
									    lowerCamelCase__					,[
									        {'generated_token_ids': ANY(lowerCamelCase__    )},
									        {'generated_token_ids': ANY(lowerCamelCase__    )},
									    ]					,)
									UpperCAmelCase__    =							text_generator.model.config.eos_token_id
									UpperCAmelCase__    =							'<pad>'
									UpperCAmelCase__    =							text_generator(
									    ['This is a test', 'This is a second test']					,do_sample=lowerCamelCase__					,num_return_sequences=2					,batch_size=2					,return_tensors=lowerCamelCase__					,)
									self.assertEqual(
									    lowerCamelCase__					,[
									        [
									            {'generated_token_ids': ANY(lowerCamelCase__    )},
									            {'generated_token_ids': ANY(lowerCamelCase__    )},
									        ],
									        [
									            {'generated_token_ids': ANY(lowerCamelCase__    )},
									            {'generated_token_ids': ANY(lowerCamelCase__    )},
									        ],
									    ]					,)







					@require_tf
					def 	__lowerCAmelCase		(       self     :			Tuple    ):
									UpperCAmelCase__    =							pipeline(task='text-generation'					,model='sshleifer/tiny-ctrl'					,framework='tf'    )

									# Using `do_sample=False` to force deterministic output
									UpperCAmelCase__    =							text_generator('This is a test'					,do_sample=lowerCamelCase__    )
									self.assertEqual(
									    lowerCamelCase__					,[
									        {
									            'generated_text': (
									                'This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵'
									                ' please,'
									            )
									        }
									    ]					,)

									UpperCAmelCase__    =							text_generator(['This is a test', 'This is a second test']					,do_sample=lowerCamelCase__    )
									self.assertEqual(
									    lowerCamelCase__					,[
									        [
									            {
									                'generated_text': (
									                    'This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵'
									                    ' please,'
									                )
									            }
									        ],
									        [
									            {
									                'generated_text': (
									                    'This is a second test Chieftain Chieftain prefecture prefecture prefecture Cannes Cannes'
									                    ' Cannes 閲閲Cannes Cannes Cannes 攵 please,'
									                )
									            }
									        ],
									    ]					,)







					def 	__lowerCAmelCase		(       self     :			Tuple					,lowerCamelCase__     :			int					,lowerCamelCase__     :			List[str]					,lowerCamelCase__     :			Dict    ):
									UpperCAmelCase__    =							TextGenerationPipeline(model=lowerCamelCase__					,tokenizer=lowerCamelCase__    )
									return text_generator, ["This is a test", "Another test"]







					def 	__lowerCAmelCase		(       self     :			List[Any]    ):
									UpperCAmelCase__    =							'Hello I believe in'
									UpperCAmelCase__    =							pipeline('text-generation'					,model='hf-internal-testing/tiny-random-gpt2'    )
									UpperCAmelCase__    =							text_generator(lowerCamelCase__    )
									self.assertEqual(
									    lowerCamelCase__					,[{'generated_text': 'Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe'}]					,)

									UpperCAmelCase__    =							text_generator(lowerCamelCase__					,stop_sequence=' fe'    )
									self.assertEqual(lowerCamelCase__					,[{'generated_text': 'Hello I believe in fe'}]    )







					def 	__lowerCAmelCase		(       self     :			Union[str, Any]					,lowerCamelCase__     :			List[str]					,lowerCamelCase__     :			Tuple    ):
									UpperCAmelCase__    =							text_generator.model
									UpperCAmelCase__    =							text_generator.tokenizer

									UpperCAmelCase__    =							text_generator('This is a test'    )
									self.assertEqual(lowerCamelCase__					,[{'generated_text': ANY(lowerCamelCase__    )}]    )
									self.assertTrue(outputs[0]['generated_text'].startswith('This is a test'    )    )

									UpperCAmelCase__    =							text_generator('This is a test'					,return_full_text=lowerCamelCase__    )
									self.assertEqual(lowerCamelCase__					,[{'generated_text': ANY(lowerCamelCase__    )}]    )
									self.assertNotIn('This is a test'					,outputs[0]['generated_text']    )

									UpperCAmelCase__    =							pipeline(task='text-generation'					,model=lowerCamelCase__					,tokenizer=lowerCamelCase__					,return_full_text=lowerCamelCase__    )
									UpperCAmelCase__    =							text_generator('This is a test'    )
									self.assertEqual(lowerCamelCase__					,[{'generated_text': ANY(lowerCamelCase__    )}]    )
									self.assertNotIn('This is a test'					,outputs[0]['generated_text']    )

									UpperCAmelCase__    =							text_generator('This is a test'					,return_full_text=lowerCamelCase__    )
									self.assertEqual(lowerCamelCase__					,[{'generated_text': ANY(lowerCamelCase__    )}]    )
									self.assertTrue(outputs[0]['generated_text'].startswith('This is a test'    )    )

									UpperCAmelCase__    =							text_generator(['This is great !', 'Something else']					,num_return_sequences=2					,do_sample=lowerCamelCase__    )
									self.assertEqual(
									    lowerCamelCase__					,[
									        [{'generated_text': ANY(lowerCamelCase__    )}, {'generated_text': ANY(lowerCamelCase__    )}],
									        [{'generated_text': ANY(lowerCamelCase__    )}, {'generated_text': ANY(lowerCamelCase__    )}],
									    ]					,)

									if text_generator.tokenizer.pad_token is not None:
													UpperCAmelCase__    =							text_generator(
													    ['This is great !', 'Something else']					,num_return_sequences=2					,batch_size=2					,do_sample=lowerCamelCase__    )
													self.assertEqual(
													    lowerCamelCase__					,[
													        [{'generated_text': ANY(lowerCamelCase__    )}, {'generated_text': ANY(lowerCamelCase__    )}],
													        [{'generated_text': ANY(lowerCamelCase__    )}, {'generated_text': ANY(lowerCamelCase__    )}],
													    ]					,)

									with self.assertRaises(lowerCamelCase__    ):
													UpperCAmelCase__    =							text_generator('test'					,return_full_text=lowerCamelCase__					,return_text=lowerCamelCase__    )
									with self.assertRaises(lowerCamelCase__    ):
													UpperCAmelCase__    =							text_generator('test'					,return_full_text=lowerCamelCase__					,return_tensors=lowerCamelCase__    )
									with self.assertRaises(lowerCamelCase__    ):
													UpperCAmelCase__    =							text_generator('test'					,return_text=lowerCamelCase__					,return_tensors=lowerCamelCase__    )

									# Empty prompt is slighly special
									# it requires BOS token to exist.
									# Special case for Pegasus which will always append EOS so will
									# work even without BOS.
									if (
									    text_generator.tokenizer.bos_token_id is not None
									    or "Pegasus" in tokenizer.__class__.__name__
									    or "Git" in model.__class__.__name__
									):
													UpperCAmelCase__    =							text_generator(''    )
													self.assertEqual(lowerCamelCase__					,[{'generated_text': ANY(lowerCamelCase__    )}]    )
									else:
													with self.assertRaises((ValueError, AssertionError)    ):
																	UpperCAmelCase__    =							text_generator(''    )

									if text_generator.framework == "tf":
													# TF generation does not support max_new_tokens, and it's impossible
													# to control long generation with only max_length without
													# fancy calculation, dismissing tests for now.
													return
									# We don't care about infinite range models.
									# They already work.
									# Skip this test for XGLM, since it uses sinusoidal positional embeddings which are resized on-the-fly.
									UpperCAmelCase__    =							['RwkvForCausalLM', 'XGLMForCausalLM', 'GPTNeoXForCausalLM']
									if (
									    tokenizer.model_max_length < 10_000
									    and text_generator.model.__class__.__name__ not in EXTRA_MODELS_CAN_HANDLE_LONG_INPUTS
									):
													# Handling of large generations
													with self.assertRaises((RuntimeError, IndexError, ValueError, AssertionError)    ):
																	text_generator('This is a test' * 500					,max_new_tokens=20    )

													UpperCAmelCase__    =							text_generator('This is a test' * 500					,handle_long_generation='hole'					,max_new_tokens=20    )
													# Hole strategy cannot work
													with self.assertRaises(lowerCamelCase__    ):
																	text_generator(
																	    'This is a test' * 500					,handle_long_generation='hole'					,max_new_tokens=tokenizer.model_max_length + 10					,)







					@require_torch
					@require_accelerate
					@require_torch_gpu
					def 	__lowerCAmelCase		(       self     :			List[Any]    ):
									import torch

									# Classic `model_kwargs`
									UpperCAmelCase__    =							pipeline(
									    model='hf-internal-testing/tiny-random-bloom'					,model_kwargs={'device_map': 'auto', 'torch_dtype': torch.bfloataa}					,)
									self.assertEqual(pipe.model.device					,torch.device(0    )    )
									self.assertEqual(pipe.model.lm_head.weight.dtype					,torch.bfloataa    )
									UpperCAmelCase__    =							pipe('This is a test'    )
									self.assertEqual(
									    lowerCamelCase__					,[
									        {
									            'generated_text': (
									                'This is a test test test test test test test test test test test test test test test test'
									                ' test'
									            )
									        }
									    ]					,)

									# Upgraded those two to real pipeline arguments (they just get sent for the model as they're unlikely to mean anything else.)
									UpperCAmelCase__    =							pipeline(model='hf-internal-testing/tiny-random-bloom'					,device_map='auto'					,torch_dtype=torch.bfloataa    )
									self.assertEqual(pipe.model.device					,torch.device(0    )    )
									self.assertEqual(pipe.model.lm_head.weight.dtype					,torch.bfloataa    )
									UpperCAmelCase__    =							pipe('This is a test'    )
									self.assertEqual(
									    lowerCamelCase__					,[
									        {
									            'generated_text': (
									                'This is a test test test test test test test test test test test test test test test test'
									                ' test'
									            )
									        }
									    ]					,)

									# torch_dtype will be automatically set to float32 if not provided - check: https://github.com/huggingface/transformers/pull/20602
									UpperCAmelCase__    =							pipeline(model='hf-internal-testing/tiny-random-bloom'					,device_map='auto'    )
									self.assertEqual(pipe.model.device					,torch.device(0    )    )
									self.assertEqual(pipe.model.lm_head.weight.dtype					,torch.floataa    )
									UpperCAmelCase__    =							pipe('This is a test'    )
									self.assertEqual(
									    lowerCamelCase__					,[
									        {
									            'generated_text': (
									                'This is a test test test test test test test test test test test test test test test test'
									                ' test'
									            )
									        }
									    ]					,)







					@require_torch
					@require_torch_gpu
					def 	__lowerCAmelCase		(       self     :			Tuple    ):
									import torch

									UpperCAmelCase__    =							pipeline(model='hf-internal-testing/tiny-random-bloom'					,device=0					,torch_dtype=torch.floataa    )
									pipe('This is a test'    )







					@require_torch
					@require_accelerate
					@require_torch_gpu
					def 	__lowerCAmelCase		(       self     :			Union[str, Any]    ):
									import torch

									UpperCAmelCase__    =							pipeline(model='hf-internal-testing/tiny-random-bloom'					,device_map='auto'					,torch_dtype=torch.floataa    )
									pipe('This is a test'					,do_sample=lowerCamelCase__					,top_p=0.5    )







					def 	__lowerCAmelCase		(       self     :			Optional[Any]    ):
									UpperCAmelCase__    =							'Hello world'
									UpperCAmelCase__    =							pipeline('text-generation'					,model='hf-internal-testing/tiny-random-gpt2'    )
									if text_generator.model.framework == "tf":
													UpperCAmelCase__    =							logging.get_logger('transformers.generation.tf_utils'    )
									else:
													UpperCAmelCase__    =							logging.get_logger('transformers.generation.utils'    )
									UpperCAmelCase__    =							'Both `max_new_tokens`'  # The beggining of the message to be checked in this test

									# Both are set by the user -> log warning
									with CaptureLogger(lowerCamelCase__    ) as cl:
													UpperCAmelCase__    =							text_generator(lowerCamelCase__					,max_length=10					,max_new_tokens=1    )
									self.assertIn(lowerCamelCase__					,cl.out    )

									# The user only sets one -> no warning
									with CaptureLogger(lowerCamelCase__    ) as cl:
													UpperCAmelCase__    =							text_generator(lowerCamelCase__					,max_new_tokens=1    )
									self.assertNotIn(lowerCamelCase__					,cl.out    )

									with CaptureLogger(lowerCamelCase__    ) as cl:
													UpperCAmelCase__    =							text_generator(lowerCamelCase__					,max_length=10    )
									self.assertNotIn(lowerCamelCase__					,cl.out    )




"""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()
lowerCAmelCase__						:						Optional[int]		  =			logging.get_logger(__name__)


lowerCAmelCase__						:						Optional[Any]		  =			['model.decoder.embed_positions.weights']



def 						a_		(      lowerCamelCase ):
				if "emb" in name:
								UpperCAmelCase__    =							name.replace('emb'      ,    'model.decoder.embed_tokens' )
				if "transformer" in name:
								UpperCAmelCase__    =							name.replace('transformer'      ,    'model.decoder' )
				if "cross_attention" in name:
								UpperCAmelCase__    =							name.replace('cross_attention'      ,    'encoder_attn' )
				if "linear1" in name:
								UpperCAmelCase__    =							name.replace('linear1'      ,    'fc1' )
				if "linear2" in name:
								UpperCAmelCase__    =							name.replace('linear2'      ,    'fc2' )
				if "norm1" in name:
								UpperCAmelCase__    =							name.replace('norm1'      ,    'self_attn_layer_norm' )
				if "norm_cross" in name:
								UpperCAmelCase__    =							name.replace('norm_cross'      ,    'encoder_attn_layer_norm' )
				if "norm2" in name:
								UpperCAmelCase__    =							name.replace('norm2'      ,    'final_layer_norm' )
				if "out_norm" in name:
								UpperCAmelCase__    =							name.replace('out_norm'      ,    'model.decoder.layer_norm' )
				if "linears" in name:
								UpperCAmelCase__    =							name.replace('linears'      ,    'lm_heads' )
				if "condition_provider.conditioners.description.output_proj" in name:
								UpperCAmelCase__    =							name.replace('condition_provider.conditioners.description.output_proj'      ,    'enc_to_dec_proj' )
				return name



def 						a_		(      lowerCamelCase      ,    lowerCamelCase ):
				UpperCAmelCase__    =							list(state_dict.keys() )
				UpperCAmelCase__    =							{}
				for key in keys:
								UpperCAmelCase__    =							state_dict.pop(lowerCamelCase )
								UpperCAmelCase__    =							rename_keys(lowerCamelCase )
								if "in_proj_weight" in key:
												# split fused qkv proj
												UpperCAmelCase__    =							val[:hidden_size, :]
												UpperCAmelCase__    =							val[hidden_size : 2 * hidden_size, :]
												UpperCAmelCase__    =							val[-hidden_size:, :]
								elif "enc_to_dec_proj" in key:
												UpperCAmelCase__    =							val
								else:
												UpperCAmelCase__    =							val
				return state_dict, enc_dec_proj_state_dict



def 						a_		(      lowerCamelCase ):
				if checkpoint == "small":
								# default config values
								UpperCAmelCase__    =							1_0_2_4
								UpperCAmelCase__    =							2_4
								UpperCAmelCase__    =							1_6
				elif checkpoint == "medium":
								UpperCAmelCase__    =							1_5_3_6
								UpperCAmelCase__    =							4_8
								UpperCAmelCase__    =							2_4
				elif checkpoint == "large":
								UpperCAmelCase__    =							2_0_4_8
								UpperCAmelCase__    =							4_8
								UpperCAmelCase__    =							3_2
				else:
								raise ValueError(f'''Checkpoint should be one of `[\'small\', \'medium\', \'large\']`, got {checkpoint}.''' )
				UpperCAmelCase__    =							MusicgenDecoderConfig(
				    hidden_size=lowerCamelCase      ,    ffn_dim=hidden_size * 4      ,    num_hidden_layers=lowerCamelCase      ,    num_attention_heads=lowerCamelCase      ,    )
				return config



@torch.no_grad()
def 						a_		(      lowerCamelCase      ,    lowerCamelCase=None      ,    lowerCamelCase=None      ,    lowerCamelCase="cpu" ):
				UpperCAmelCase__    =							MusicGen.get_pretrained(lowerCamelCase      ,    device=lowerCamelCase )
				UpperCAmelCase__    =							decoder_config_from_checkpoint(lowerCamelCase )

				UpperCAmelCase__    =							fairseq_model.lm.state_dict()
				UpperCAmelCase__   ,			UpperCAmelCase__    =							rename_state_dict(
				    lowerCamelCase      ,    hidden_size=decoder_config.hidden_size )

				UpperCAmelCase__    =							TaEncoderModel.from_pretrained('t5-base' )
				UpperCAmelCase__    =							EncodecModel.from_pretrained('facebook/encodec_32khz' )
				UpperCAmelCase__    =							MusicgenForCausalLM(lowerCamelCase ).eval()

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

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

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

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

				# init the composite model
				UpperCAmelCase__    =							MusicgenForConditionalGeneration(text_encoder=lowerCamelCase      ,    audio_encoder=lowerCamelCase      ,    decoder=lowerCamelCase )

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

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

				with torch.no_grad():
								UpperCAmelCase__    =							model(input_ids=lowerCamelCase      ,    decoder_input_ids=lowerCamelCase ).logits

				if logits.shape != (8, 1, 2_0_4_8):
								raise ValueError('Incorrect shape for logits' )

				# now construct the processor
				UpperCAmelCase__    =							AutoTokenizer.from_pretrained('t5-base' )
				UpperCAmelCase__    =							AutoFeatureExtractor.from_pretrained('facebook/encodec_32khz'      ,    padding_side='left' )

				UpperCAmelCase__    =							MusicgenProcessor(feature_extractor=lowerCamelCase      ,    tokenizer=lowerCamelCase )

				# set the appropriate bos/pad token ids
				UpperCAmelCase__    =							2_0_4_8
				UpperCAmelCase__    =							2_0_4_8

				# set other default generation config params
				UpperCAmelCase__    =							int(3_0 * audio_encoder.config.frame_rate )
				UpperCAmelCase__    =							True
				UpperCAmelCase__    =							3.0

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

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


if __name__ == "__main__":
						lowerCAmelCase__						:						Tuple		  =			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.'
						)

						lowerCAmelCase__						:						List[str]		  =			parser.parse_args()
						convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)



