




import unittest
from unittest import skip

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow
from transformers.trainer_utils import set_seed


if is_torch_available():
		import torch

		from transformers import JukeboxModel, JukeboxPrior, JukeboxTokenizer




@require_torch
class       UpperCAmelCase_   ( unittest.TestCase							):




							'''simple docstring'''



							a__    	=  (JukeboxModel,) if is_torch_available() else ()
							a__    	=  '''openai/jukebox-1b-lyrics'''
							a__    	=  {
							    '''artist''': '''Zac Brown Band''',
							    '''genres''': '''Country''',
							    '''lyrics''': '''I met a traveller from an antique land,
    Who said "Two vast and trunkless legs of stone
    Stand in the desert. . . . Near them, on the sand,
    Half sunk a shattered visage lies, whose frown,
    And wrinkled lip, and sneer of cold command,
    Tell that its sculptor well those passions read
    Which yet survive, stamped on these lifeless things,
    The hand that mocked them, and the heart that fed;
    And on the pedestal, these words appear:
    My name is Ozymandias, King of Kings;
    Look on my Works, ye Mighty, and despair!
    Nothing beside remains. Round the decay
    Of that colossal Wreck, boundless and bare
    The lone and level sands stretch far away
    ''',
							}
							# fmt: off
							a__    	=  [
							    1_8_6_4, 1_5_3_6, 1_2_1_3, 1_8_7_0, 1_3_5_7, 1_5_3_6, 5_1_9, 8_8_0, 1_3_2_3, 7_8_9, 1_0_8_2, 5_3_4,
							    1_0_0_0, 1_4_4_5, 1_1_0_5, 1_1_3_0, 9_6_7, 5_1_5, 1_4_3_4, 1_6_2_0, 5_3_4, 1_4_9_5, 2_8_3, 1_4_4_5,
							    3_3_3, 1_3_0_7, 5_3_9, 1_6_3_1, 1_5_2_8, 3_7_5, 1_4_3_4, 6_7_3, 6_2_7, 7_1_0, 7_7_8, 1_8_8_3,
							    1_4_0_5, 1_2_7_6, 1_4_5_5, 1_2_2_8
							]

							a__    	=  [
							    1_4_8_9, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3
							]

							a__    	=  [
							    1_1_2_5, 1_7_5_1, 6_9_7, 1_7_7_6, 1_1_4_1, 1_4_7_6, 3_9_1, 6_9_7, 1_1_2_5, 6_8_4, 8_6_7, 4_1_6,
							    8_4_4, 1_3_7_2, 1_2_7_4, 7_1_7, 1_2_7_4, 8_4_4, 1_2_9_9, 1_4_1_9, 6_9_7, 1_3_7_0, 3_1_7, 1_1_2_5,
							    1_9_1, 1_4_4_0, 1_3_7_0, 1_4_4_0, 1_3_7_0, 2_8_2, 1_6_2_1, 1_3_7_0, 3_6_8, 3_4_9, 8_6_7, 1_8_7_2,
							    1_2_6_2, 8_6_9, 1_7_2_8, 7_4_7
							]
							a__    	=  [
							    4_1_6, 4_1_6, 1_1_2_5, 1_1_2_5, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
							    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
							    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
							    4_1_6, 4_1_6, 4_1_6, 4_1_6
							]

							a__    	=  [
							    1_7_5_5, 8_4_2, 3_0_7, 1_8_4_3, 1_0_2_2, 1_3_9_5, 2_3_4, 1_5_5_4, 8_0_6, 7_3_9, 1_0_2_2, 4_4_2,
							    6_1_6, 5_5_6, 2_6_8, 1_4_9_9, 9_3_3, 4_5_7, 1_4_4_0, 1_8_3_7, 7_5_5, 9_8_5, 3_0_8, 9_0_2,
							    2_9_3, 1_4_4_3, 1_6_7_1, 1_1_4_1, 1_5_3_3, 5_5_5, 1_5_6_2, 1_0_6_1, 2_8_7, 4_1_7, 1_0_2_2, 2_0_0_8,
							    1_1_8_6, 1_0_1_5, 1_7_7_7, 2_6_8
							]
							a__    	=  [
							    8_5_4, 8_4_2, 1_3_5_3, 1_1_4, 1_3_5_3, 8_4_2, 1_8_5, 8_4_2, 1_8_5, 1_1_4, 5_9_1, 8_4_2,
							    1_8_5, 4_1_7, 1_8_5, 8_4_2, 3_0_7, 8_4_2, 5_9_1, 8_4_2, 1_8_5, 8_4_2, 3_0_7, 8_4_2,
							    5_9_1, 8_4_2, 1_3_5_3, 8_4_2, 1_8_5, 8_4_2, 5_9_1, 8_4_2, 5_9_1, 1_1_4, 5_9_1, 8_4_2,
							    1_8_5, 8_4_2, 5_9_1, 8_9
							]

							a__    	=  [1_0_5_8_3_0_4, 0, 7_8_6_4_3_2, 7_1_6_9, 5_0_7, 7_6, 2_7, 4_0, 3_0, 7_6]

							a__    	=  [
							    3_9_0, 1_1_6_0, 1_0_0_2, 1_9_0_7, 1_7_8_8, 1_7_8_8, 1_7_8_8, 1_9_0_7, 1_0_0_2, 1_0_0_2, 1_8_5_4, 1_0_0_2,
							    1_0_0_2, 1_0_0_2, 1_0_0_2, 1_0_0_2, 1_0_0_2, 1_1_6_0, 1_1_6_0, 1_6_0_6, 5_9_6, 5_9_6, 1_1_6_0, 1_0_0_2,
							    1_5_1_6, 5_9_6, 1_0_0_2, 1_0_0_2, 1_0_0_2, 1_9_0_7, 1_7_8_8, 1_7_8_8, 1_7_8_8, 1_8_5_4, 1_7_8_8, 1_9_0_7,
							    1_9_0_7, 1_7_8_8, 5_9_6, 1_6_2_6
							]
							a__    	=  [
							    1_2_3_6, 1_6_6_8, 1_4_8_4, 1_9_2_0, 1_8_4_8, 1_4_0_9, 1_3_9, 8_6_4, 1_8_2_8, 1_2_7_2, 1_5_9_9, 8_2_4,
							    1_6_7_2, 1_3_9, 5_5_5, 1_4_8_4, 8_2_4, 1_9_2_0, 5_5_5, 5_9_6, 1_5_7_9, 1_5_9_9, 1_2_3_1, 1_5_9_9,
							    1_6_3_7, 1_4_0_7, 2_1_2, 8_2_4, 1_5_9_9, 1_1_6, 1_4_3_3, 8_2_4, 2_5_8, 1_5_9_9, 1_4_3_3, 1_8_9_5,
							    1_0_6_3, 1_4_3_3, 1_4_3_3, 1_5_9_9
							]
							a__    	=  [
							    1_6_8_4, 1_8_7_3, 1_1_1_9, 1_1_8_9, 3_9_5, 6_1_1, 1_9_0_1, 9_7_2, 8_9_0, 1_3_3_7, 1_3_9_2, 1_9_2_7,
							    9_6, 9_7_2, 6_7_2, 7_8_0, 1_1_1_9, 8_9_0, 1_5_8, 7_7_1, 1_0_7_3, 1_9_2_7, 3_5_3, 1_3_3_1,
							    1_2_6_9, 1_4_5_9, 1_3_3_3, 1_6_4_5, 8_1_2, 1_5_7_7, 1_3_3_7, 6_0_6, 3_5_3, 9_8_1, 1_4_6_6, 6_1_9,
							    1_9_7, 3_9_1, 3_0_2, 1_9_3_0
							]
							a__    	=  [
							    3_9_0, 1_1_6_0, 1_0_0_2, 1_9_0_7, 1_7_8_8, 1_7_8_8, 1_7_8_8, 1_9_0_7, 1_0_0_2, 1_0_0_2, 1_8_5_4, 1_0_0_2,
							    1_0_0_2, 1_0_0_2, 1_0_0_2, 1_0_0_2, 1_0_0_2, 1_1_6_0, 1_1_6_0, 1_6_0_6, 5_9_6, 5_9_6, 1_1_6_0, 1_0_0_2,
							    1_5_1_6, 5_9_6, 1_0_0_2, 1_0_0_2, 1_0_0_2, 1_9_0_7, 1_7_8_8, 1_7_8_8, 1_7_8_8, 1_8_5_4, 1_7_8_8, 1_9_0_7,
							    1_9_0_7, 1_7_8_8, 5_9_6, 1_6_2_6
							]
							a__    	=  [
							    -0.0_4_9_2, -0.0_5_2_4, -0.0_5_6_5, -0.0_6_4_0, -0.0_6_8_6, -0.0_6_8_4, -0.0_6_7_7, -0.0_6_6_4,
							    -0.0_6_0_5, -0.0_4_9_0, -0.0_3_3_0, -0.0_1_6_8, -0.0_0_8_3, -0.0_0_7_5, -0.0_0_5_1, 0.0_0_2_5,
							    0.0_1_3_6, 0.0_2_6_1, 0.0_3_8_6, 0.0_4_9_7, 0.0_5_8_0, 0.0_5_9_9, 0.0_5_8_3, 0.0_6_1_4,
							    0.0_7_4_0, 0.0_8_8_9, 0.1_0_2_3, 0.1_1_6_2, 0.1_2_1_1, 0.1_2_1_2, 0.1_2_5_1, 0.1_3_3_6,
							    0.1_5_0_2, 0.1_6_8_6, 0.1_8_8_3, 0.2_1_4_8, 0.2_3_6_3, 0.2_4_5_8, 0.2_5_0_7, 0.2_5_3_1
							]
							a__    	=  [
							    0.0_2_5_6, -0.0_5_4_4, 0.1_6_0_0, -0.0_0_3_2, 0.1_0_6_6, 0.0_8_2_5, -0.0_0_1_3, 0.3_4_4_0,
							    0.0_2_1_0, 0.0_4_1_2, -0.1_7_7_7, -0.0_8_9_2, -0.0_1_6_4, 0.0_2_8_5, -0.0_6_1_3, -0.0_6_1_7,
							    -0.0_1_3_7, -0.0_2_0_1, -0.0_1_7_5, 0.0_2_1_5, -0.0_6_2_7, 0.0_5_2_0, -0.0_7_3_0, 0.0_9_7_0,
							    -0.0_1_0_0, 0.0_4_4_2, -0.0_5_8_6, 0.0_2_0_7, -0.0_0_1_5, -0.0_0_8_2
							]
							a__    	=  [
							    0.0_4_1_5, 0.0_8_7_7, 0.0_0_2_2, -0.0_0_5_5, 0.0_7_5_1, 0.0_3_3_4, 0.0_3_2_4, -0.0_0_6_8,
							    0.0_0_1_1, 0.0_0_1_7, -0.0_6_7_6, 0.0_6_5_5, -0.0_1_4_3, 0.0_3_9_9, 0.0_3_0_3, 0.0_7_4_3,
							    -0.0_1_6_8, -0.0_3_9_4, -0.1_1_1_3, 0.0_1_2_4, 0.0_4_4_2, 0.0_2_6_7, -0.0_0_0_3, -0.1_5_3_6,
							    -0.0_1_1_6, -0.1_8_3_7, -0.0_1_8_0, -0.1_0_2_6, -0.0_7_7_7, -0.0_4_5_6
							]
							a__    	=  [
							    7_6, 2_7, 4_0, 3_0, 7_6, 4_6, 4_4, 4_7, 4_0, 3_7, 3_8, 3_1, 4_5, 4_5, 7_6, 3_8, 3_1, 3_3,
							    4_5, 7_6, 4_1, 3_2, 7_6, 4_5, 4_6, 4_1, 4_0, 3_1, 7_8, 7_6
							]





							def  _UpperCAmelCase  (       self							:						int      )      ->       List[str]:
										SCREAMING_SNAKE_CASE     =       JukeboxTokenizer.from_pretrained(self.model_id      )
										SCREAMING_SNAKE_CASE     =       tokenizer(**self.metas      )["""input_ids"""]
										return tokens





							@slow
							def  _UpperCAmelCase  (       self							:						int      )      ->       int:
										SCREAMING_SNAKE_CASE     =       JukeboxModel.from_pretrained(self.model_id							,  min_duration=0      ).eval()
										SCREAMING_SNAKE_CASE     =       self.prepare_inputs()

										set_seed(0      )
										SCREAMING_SNAKE_CASE     =       [torch.zeros(1							,  0							,  dtype=torch.long      ).cpu() for _ in range(3      )]
										SCREAMING_SNAKE_CASE     =       model._sample(a							,  a							,  [0]							,  sample_length=40 * model.priors[0].raw_to_tokens							,  save_results=a      )
										self.assertIn(zs[0][0].detach().cpu().tolist()							,  [self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2]      )

										set_seed(0      )
										SCREAMING_SNAKE_CASE     =       model._sample(a							,  a							,  [1]							,  sample_length=40 * model.priors[1].raw_to_tokens							,  save_results=a      )
										self.assertIn(zs[1][0].detach().cpu().tolist()							,  [self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]      )

										set_seed(0      )
										SCREAMING_SNAKE_CASE     =       model._sample(a							,  a							,  [2]							,  sample_length=40 * model.priors[2].raw_to_tokens							,  save_results=a      )
										self.assertIn(zs[2][0].detach().cpu().tolist()							,  [self.EXPECTED_OUTPUT_0, self.EXPECTED_OUTPUT_0_PT_2]      )





							@slow
							def  _UpperCAmelCase  (       self							:						Tuple      )      ->       Optional[Any]:
										SCREAMING_SNAKE_CASE     =       False
										SCREAMING_SNAKE_CASE     =       JukeboxModel.from_pretrained(self.model_id							,  min_duration=0      ).eval()

										SCREAMING_SNAKE_CASE     =       self.prepare_inputs()
										set_seed(0      )
										SCREAMING_SNAKE_CASE     =       [torch.zeros(1							,  0							,  dtype=torch.long      ) for _ in range(3      )]

										SCREAMING_SNAKE_CASE     =       model.priors[0]
										SCREAMING_SNAKE_CASE     =       0
										SCREAMING_SNAKE_CASE     =       top_prior.get_music_tokens_conds(a							,  start=a							,  end=start + top_prior.n_ctx      )
										SCREAMING_SNAKE_CASE     =       top_prior.get_metadata(labels[0].clone()							,  a							,  1_058_304							,  0      )

										self.assertIsNone(a      )
										self.assertListEqual(metadata.numpy()[0][:10].tolist()							,  self.EXPECTED_Y_COND      )

										SCREAMING_SNAKE_CASE  ,	SCREAMING_SNAKE_CASE  ,	SCREAMING_SNAKE_CASE     =       top_prior.get_cond(a							,  a      )
										torch.testing.assert_allclose(
										    audio_conditioning[0][0][:30].detach()							,  torch.tensor(self.EXPECTED_AUDIO_COND      )							,  atol=1E-4							,  rtol=1E-4      )
										torch.testing.assert_allclose(
										    metadata_conditioning[0][0][:30].detach()							,  torch.tensor(self.EXPECTED_META_COND      )							,  atol=1E-4							,  rtol=1E-4      )
										torch.testing.assert_allclose(
										    lyric_tokens[0, :30].detach()							,  torch.tensor(self.EXPECTED_LYRIC_COND      )							,  atol=1E-4							,  rtol=1E-4      )





							@slow
							def  _UpperCAmelCase  (       self							:						str      )      ->       List[Any]:
										SCREAMING_SNAKE_CASE     =       False

										SCREAMING_SNAKE_CASE     =       JukeboxModel.from_pretrained(self.model_id							,  min_duration=0      ).eval()
										set_seed(0      )
										SCREAMING_SNAKE_CASE     =       torch.rand((1, 5_120, 1)      )
										SCREAMING_SNAKE_CASE     =       list(self.prepare_inputs()      )

										SCREAMING_SNAKE_CASE     =       [model.vqvae.encode(a							,  start_level=2							,  bs_chunks=waveform.shape[0]      )[0], None, None]
										SCREAMING_SNAKE_CASE     =       model._sample(
										    a							,  a							,  sample_levels=[0]							,  save_results=a							,  sample_length=40 * model.priors[0].raw_to_tokens      )
										torch.testing.assert_allclose(zs[0][0][:40]							,  torch.tensor(self.EXPECTED_PRIMED_0      )      )

										SCREAMING_SNAKE_CASE     =       torch.cat((zs[0], torch.zeros(1							,  2_048 - zs[0].shape[-1]      ))							,  dim=-1      ).long()
										SCREAMING_SNAKE_CASE     =       [upper_a, model.vqvae.encode(a							,  start_level=1							,  bs_chunks=waveform.shape[0]      )[0], None]
										SCREAMING_SNAKE_CASE     =       model._sample(
										    a							,  a							,  sample_levels=[1]							,  save_results=a							,  sample_length=40 * model.priors[1].raw_to_tokens      )
										torch.testing.assert_allclose(zs[1][0][:40]							,  torch.tensor(self.EXPECTED_PRIMED_1      )      )

										SCREAMING_SNAKE_CASE     =       torch.cat((zs[1], torch.zeros(1							,  2_048 - zs[1].shape[-1]      ))							,  dim=-1      ).long()
										SCREAMING_SNAKE_CASE     =       [upper_a, upper_a, model.vqvae.encode(a							,  start_level=0							,  bs_chunks=waveform.shape[0]      )[0]]
										SCREAMING_SNAKE_CASE     =       model._sample(
										    a							,  a							,  sample_levels=[2]							,  save_results=a							,  sample_length=40 * model.priors[2].raw_to_tokens      )
										torch.testing.assert_allclose(zs[2][0][:40].cpu()							,  torch.tensor(self.EXPECTED_PRIMED_2      )      )







							@slow
							def  _UpperCAmelCase  (       self							:						int      )      ->       int:
										SCREAMING_SNAKE_CASE     =       JukeboxModel.from_pretrained(self.model_id							,  min_duration=0      ).eval()
										set_seed(0      )
										SCREAMING_SNAKE_CASE     =       torch.rand((1, 5_120, 1)      )
										with torch.no_grad():
													SCREAMING_SNAKE_CASE     =       model.vqvae.encode(a							,  start_level=2							,  bs_chunks=x.shape[0]      )
										torch.testing.assert_allclose(zs[0][0]							,  torch.tensor(self.EXPECTED_VQVAE_ENCODE      )      )

										with torch.no_grad():
													SCREAMING_SNAKE_CASE     =       model.vqvae.decode(a							,  start_level=2							,  bs_chunks=x.shape[0]      )
										torch.testing.assert_allclose(x[0, :40, 0]							,  torch.tensor(self.EXPECTED_VQVAE_DECODE      )							,  atol=1E-4							,  rtol=1E-4      )





@require_torch
class       UpperCAmelCase_   ( unittest.TestCase							):




							'''simple docstring'''



							a__    	=  (JukeboxModel,) if is_torch_available() else ()
							a__    	=  '''openai/jukebox-5b-lyrics'''
							a__    	=  {
							    '''artist''': '''Zac Brown Band''',
							    '''genres''': '''Country''',
							    '''lyrics''': '''I met a traveller from an antique land,
    Who said "Two vast and trunkless legs of stone
    Stand in the desert. . . . Near them, on the sand,
    Half sunk a shattered visage lies, whose frown,
    And wrinkled lip, and sneer of cold command,
    Tell that its sculptor well those passions read
    Which yet survive, stamped on these lifeless things,
    The hand that mocked them, and the heart that fed;
    And on the pedestal, these words appear:
    My name is Ozymandias, King of Kings;
    Look on my Works, ye Mighty, and despair!
    Nothing beside remains. Round the decay
    Of that colossal Wreck, boundless and bare
    The lone and level sands stretch far away
    ''',
							}

							# fmt: off
							a__    	=  [
							    1_4_8_9, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    1_4_8_9, 1_4_8_9, 1_4_8_9, 1_4_8_9, 1_1_5_0, 1_8_5_3, 1_5_0_9, 1_1_5_0, 1_3_5_7, 1_5_0_9, 6, 1_2_7_2
							]
							a__    	=  [
							    1_4_8_9, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3
							]

							a__    	=  [
							    1_1_2_5, 4_1_6, 1_1_2_5, 1_1_2_5, 1_1_2_5, 1_1_2_5, 1_1_2_5, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
							    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
							    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
							    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
							    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6
							]
							a__    	=  [
							    4_1_6, 4_1_6, 1_1_2_5, 1_1_2_5, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
							    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
							    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
							    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
							    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6
							]

							a__    	=  [
							    1_7_5_5, 1_0_6_1, 2_3_4, 1_7_5_5, 1_0_6_1, 1_7_5_5, 1_8_5, 2_9_0, 3_0_7, 3_0_7, 6_1_6, 6_1_6,
							    6_1_6, 6_1_6, 6_1_6, 6_1_6, 3_0_7, 2_9_0, 4_1_7, 1_7_5_5, 2_3_4, 1_7_5_5, 1_8_5, 2_9_0,
							    2_9_0, 2_9_0, 3_0_7, 6_1_6, 6_1_6, 6_1_6, 6_1_6, 6_1_6, 2_9_0, 2_3_4, 2_3_4, 1_7_5_5,
							    2_3_4, 2_3_4, 1_7_5_5, 2_3_4, 1_8_5, 1_8_5, 3_0_7, 6_1_6, 6_1_6, 6_1_6, 6_1_6, 2_9_0,
							    1_7_5_5, 1_7_5_5, 1_7_5_5, 2_3_4, 2_3_4, 1_7_5_5, 1_5_7_2, 2_9_0, 3_0_7, 6_1_6, 3_4, 6_1_6
							]
							a__    	=  [
							    8_5_4, 8_4_2, 1_3_5_3, 1_1_4, 1_3_5_3, 8_4_2, 1_8_5, 8_4_2, 1_8_5, 1_1_4, 5_9_1, 8_4_2, 1_8_5,
							    4_1_7, 1_8_5, 8_4_2, 3_0_7, 8_4_2, 5_9_1, 8_4_2, 1_8_5, 8_4_2, 1_8_5, 8_4_2, 5_9_1, 8_4_2,
							    1_3_5_3, 8_4_2, 1_8_5, 8_4_2, 5_9_1, 8_4_2, 5_9_1, 1_1_4, 5_9_1, 8_4_2, 1_8_5, 8_4_2, 5_9_1,
							    8_9, 5_9_1, 8_4_2, 5_9_1, 8_4_2, 5_9_1, 4_1_7, 1_3_7_2, 8_4_2, 1_3_7_2, 8_4_2, 3_4, 8_4_2,
							    1_8_5, 8_9, 5_9_1, 8_4_2, 1_8_5, 8_4_2, 5_9_1, 6_3_2
							]

							a__    	=  [
							    1_4_8_9, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3
							]
							a__    	=  [
							    1_4_8_9, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
							    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 1_8_5_3, 1_1_7_7, 1_5_3_6, 1_2_2_8,
							    7_1_0, 4_7_5, 1_4_8_9, 1_2_2_9, 1_2_2_4, 2_3_1, 1_2_2_4, 2_5_2, 1_4_3_4, 6_5_3, 4_7_5,
							    1_1_0_6, 1_8_7_7, 1_5_9_9, 1_2_2_8, 1_6_0_0, 1_6_8_3, 1_1_8_2, 1_8_5_3, 4_7_5, 1_8_6_4,
							    2_5_2, 1_2_2_9, 1_4_3_4, 2_0_0_1
							]

							a__    	=  [
							    1_1_2_5, 1_1_2_5, 4_1_6, 1_1_2_5, 1_1_2_5, 4_1_6, 1_1_2_5, 1_1_2_5, 4_1_6, 4_1_6, 1_1_2_5, 4_1_6,
							    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
							    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
							    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
							    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6
							]
							a__    	=  [
							    4_9_1, 1_7_5_5, 3_4, 1_6_1_3, 1_7_5_5, 4_1_7, 9_9_2, 1_6_1_3, 2_2_2, 8_4_2, 1_3_5_3, 1_6_1_3,
							    8_4_4, 6_3_2, 1_8_5, 1_6_1_3, 8_4_4, 6_3_2, 1_8_5, 1_6_1_3, 1_8_5, 8_4_2, 6_7_7, 1_6_1_3,
							    1_8_5, 1_1_4, 1_3_5_3, 1_6_1_3, 3_0_7, 8_9, 8_4_4, 1_6_1_3, 3_0_7, 1_3_3_2, 2_3_4, 1_9_7_9,
							    3_0_7, 8_9, 1_3_5_3, 6_1_6, 3_4, 8_4_2, 1_8_5, 8_4_2, 3_4, 8_4_2, 1_8_5, 8_4_2,
							    3_0_7, 1_1_4, 1_8_5, 8_9, 3_4, 1_2_6_8, 1_8_5, 8_9, 3_4, 8_4_2, 1_8_5, 8_9
							]





							def  _UpperCAmelCase  (       self							:						List[Any]							,  a							:						Optional[Any]      )      ->       int:
										SCREAMING_SNAKE_CASE     =       JukeboxTokenizer.from_pretrained(a      )
										SCREAMING_SNAKE_CASE     =       tokenizer(**self.metas      )["""input_ids"""]
										return tokens





							@slow
							def  _UpperCAmelCase  (       self							:						Optional[int]      )      ->       int:
										SCREAMING_SNAKE_CASE     =       JukeboxModel.from_pretrained(self.model_id							,  min_duration=0      ).eval()
										SCREAMING_SNAKE_CASE     =       self.prepare_inputs(self.model_id      )

										set_seed(0      )
										SCREAMING_SNAKE_CASE     =       [torch.zeros(1							,  0							,  dtype=torch.long      ).cpu() for _ in range(3      )]
										SCREAMING_SNAKE_CASE     =       model._sample(a							,  a							,  [0]							,  sample_length=60 * model.priors[0].raw_to_tokens							,  save_results=a      )
										self.assertIn(zs[0][0].detach().cpu().tolist()							,  [self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2]      )

										set_seed(0      )
										SCREAMING_SNAKE_CASE     =       model._sample(a							,  a							,  [1]							,  sample_length=60 * model.priors[1].raw_to_tokens							,  save_results=a      )
										self.assertIn(zs[1][0].detach().cpu().tolist()							,  [self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]      )

										set_seed(0      )
										SCREAMING_SNAKE_CASE     =       model._sample(a							,  a							,  [2]							,  sample_length=60 * model.priors[2].raw_to_tokens							,  save_results=a      )
										self.assertIn(zs[2][0].detach().cpu().tolist()							,  [self.EXPECTED_OUTPUT_0, self.EXPECTED_OUTPUT_0_PT_2]      )





							@slow
							@skip("""Not enough GPU memory on CI runners"""      )
							def  _UpperCAmelCase  (       self							:						Optional[Any]      )      ->       List[str]:
										SCREAMING_SNAKE_CASE     =       JukeboxModel.from_pretrained(self.model_id							,  min_duration=0      ).eval()
										SCREAMING_SNAKE_CASE     =       [i.cuda() for i in self.prepare_inputs(self.model_id      )]

										set_seed(0      )
										model.priors[0].cuda()
										SCREAMING_SNAKE_CASE     =       [torch.zeros(1							,  0							,  dtype=torch.long      ).cuda() for _ in range(3      )]
										SCREAMING_SNAKE_CASE     =       model._sample(a							,  a							,  [0]							,  sample_length=60 * model.priors[0].raw_to_tokens							,  save_results=a      )
										torch.testing.assert_allclose(zs[0][0].cpu()							,  torch.tensor(self.EXPECTED_GPU_OUTPUTS_2      )      )
										model.priors[0].cpu()

										set_seed(0      )
										model.priors[1].cuda()
										SCREAMING_SNAKE_CASE     =       model._sample(a							,  a							,  [1]							,  sample_length=60 * model.priors[1].raw_to_tokens							,  save_results=a      )
										torch.testing.assert_allclose(zs[1][0].cpu()							,  torch.tensor(self.EXPECTED_GPU_OUTPUTS_1      )      )
										model.priors[1].cpu()

										set_seed(0      )
										model.priors[2].cuda()
										SCREAMING_SNAKE_CASE     =       model._sample(a							,  a							,  [2]							,  sample_length=60 * model.priors[2].raw_to_tokens							,  save_results=a      )
										torch.testing.assert_allclose(zs[2][0].cpu()							,  torch.tensor(self.EXPECTED_GPU_OUTPUTS_0      )      )







							@slow
							def  _UpperCAmelCase  (       self							:						Optional[int]      )      ->       int:
										SCREAMING_SNAKE_CASE     =       """ArthurZ/jukebox_prior_0"""
										SCREAMING_SNAKE_CASE     =       JukeboxPrior.from_pretrained(a							,  min_duration=0      ).eval().half().to("""cuda"""      )

										SCREAMING_SNAKE_CASE     =       self.prepare_inputs(a      )[0].cuda()
										SCREAMING_SNAKE_CASE     =       model.get_metadata(a							,  0							,  7_680							,  0      )
										set_seed(0      )
										SCREAMING_SNAKE_CASE     =       model.sample(1							,  metadata=a							,  sample_tokens=60      )
										self.assertIn(outputs[0].cpu().tolist()							,  [self.EXPECTED_GPU_OUTPUTS_2, self.EXPECTED_GPU_OUTPUTS_2_PT_2]      )












import argparse

from transformers import ConvBertConfig, ConvBertModel, TFConvBertModel, load_tf_weights_in_convbert
from transformers.utils import logging


logging.set_verbosity_info()
def        lowerCamelCase_     (     SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE ):
			'''simple docstring'''
			SCREAMING_SNAKE_CASE     =       ConvBertConfig.from_json_file(SCREAMING_SNAKE_CASE )
			SCREAMING_SNAKE_CASE     =       ConvBertModel(SCREAMING_SNAKE_CASE )

			SCREAMING_SNAKE_CASE     =       load_tf_weights_in_convbert(SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE       ,						SCREAMING_SNAKE_CASE )
			model.save_pretrained(SCREAMING_SNAKE_CASE )

			SCREAMING_SNAKE_CASE     =       TFConvBertModel.from_pretrained(SCREAMING_SNAKE_CASE       ,						from_pt=SCREAMING_SNAKE_CASE )
			tf_model.save_pretrained(SCREAMING_SNAKE_CASE )


if __name__ == "__main__":
		__A    :				str    						=	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(
		    """--convbert_config_file""",
		    default=None,
		    type=str,
		    required=True,
		    help=(
		        """The config json file corresponding to the pre-trained ConvBERT model. \n"""
		        """This specifies the model architecture."""
		    ),
		)
		parser.add_argument(
		    """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
		)
		__A    :				List[Any]    						=	parser.parse_args()
		convert_orig_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.convbert_config_file, args.pytorch_dump_path)






