


'''simple docstring'''



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       a							(						unittest.TestCase			):


     """simple docstring"""






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

     SCREAMING_SNAKE_CASE			:					int				  =  [
         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
     ]

     SCREAMING_SNAKE_CASE			:					str				  =  [
         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
     ]
     SCREAMING_SNAKE_CASE			:					List[str]				  =  [
         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
     ]

     SCREAMING_SNAKE_CASE			:					Any				  =  [
         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
     ]
     SCREAMING_SNAKE_CASE			:					List[Any]				  =  [
         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
     ]

     SCREAMING_SNAKE_CASE			:					Union[str, Any]				  =  [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]

     SCREAMING_SNAKE_CASE			:					Union[str, Any]				  =  [
         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
     ]
     SCREAMING_SNAKE_CASE			:					int				  =  [
         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
     ]
     SCREAMING_SNAKE_CASE			:					Tuple				  =  [
         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
     ]
     SCREAMING_SNAKE_CASE			:					str				  =  [
         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
     ]
     SCREAMING_SNAKE_CASE			:					List[Any]				  =  [
         -0.04_92, -0.05_24, -0.05_65, -0.06_40, -0.06_86, -0.06_84, -0.06_77, -0.06_64,
         -0.06_05, -0.04_90, -0.03_30, -0.01_68, -0.00_83, -0.00_75, -0.00_51, 0.00_25,
         0.01_36, 0.02_61, 0.03_86, 0.04_97, 0.05_80, 0.05_99, 0.05_83, 0.06_14,
         0.07_40, 0.08_89, 0.10_23, 0.11_62, 0.12_11, 0.12_12, 0.12_51, 0.13_36,
         0.15_02, 0.16_86, 0.18_83, 0.21_48, 0.23_63, 0.24_58, 0.25_07, 0.25_31
     ]
     SCREAMING_SNAKE_CASE			:					Union[str, Any]				  =  [
         0.02_56, -0.05_44, 0.16_00, -0.00_32, 0.10_66, 0.08_25, -0.00_13, 0.34_40,
         0.02_10, 0.04_12, -0.17_77, -0.08_92, -0.01_64, 0.02_85, -0.06_13, -0.06_17,
         -0.01_37, -0.02_01, -0.01_75, 0.02_15, -0.06_27, 0.05_20, -0.07_30, 0.09_70,
         -0.01_00, 0.04_42, -0.05_86, 0.02_07, -0.00_15, -0.00_82
     ]
     SCREAMING_SNAKE_CASE			:					Any				  =  [
         0.04_15, 0.08_77, 0.00_22, -0.00_55, 0.07_51, 0.03_34, 0.03_24, -0.00_68,
         0.00_11, 0.00_17, -0.06_76, 0.06_55, -0.01_43, 0.03_99, 0.03_03, 0.07_43,
         -0.01_68, -0.03_94, -0.11_13, 0.01_24, 0.04_42, 0.02_67, -0.00_03, -0.15_36,
         -0.01_16, -0.18_37, -0.01_80, -0.10_26, -0.07_77, -0.04_56
     ]
     SCREAMING_SNAKE_CASE			:					str				  =  [
         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 					lowerCamelCase__					(							self	:						List[str]     )		->				List[Any]:
          __UpperCAmelCase      :					Dict    				=     JukeboxTokenizer.from_pretrained(self.model_id     )
          __UpperCAmelCase      :					Optional[int]    				=     tokenizer(**self.metas     )['''input_ids''']
          return tokens

     @slow
     def 					lowerCamelCase__					(							self	:						List[Any]     )		->				List[str]:
          __UpperCAmelCase      :					List[str]    				=     JukeboxModel.from_pretrained(self.model_id				,					min_duration=0     ).eval()
          __UpperCAmelCase      :					Union[str, Any]    				=     self.prepare_inputs()

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

          set_seed(0     )
          __UpperCAmelCase      :					Optional[Any]    				=     model._sample(_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE				,					[1]				,					sample_length=40 * model.priors[1].raw_to_tokens				,					save_results=_SCREAMING_SNAKE_CASE     )
          self.assertIn(zs[1][0].detach().cpu().tolist()				,					[self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]     )

          set_seed(0     )
          __UpperCAmelCase      :					List[Any]    				=     model._sample(_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE				,					[2]				,					sample_length=40 * model.priors[2].raw_to_tokens				,					save_results=_SCREAMING_SNAKE_CASE     )
          self.assertIn(zs[2][0].detach().cpu().tolist()				,					[self.EXPECTED_OUTPUT_0, self.EXPECTED_OUTPUT_0_PT_2]     )

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

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

          __UpperCAmelCase      :					Tuple    				=     model.priors[0]
          __UpperCAmelCase      :					Optional[int]    				=     0
          __UpperCAmelCase      :					str    				=     top_prior.get_music_tokens_conds(_SCREAMING_SNAKE_CASE				,					start=_SCREAMING_SNAKE_CASE				,					end=start + top_prior.n_ctx     )
          __UpperCAmelCase      :					Dict    				=     top_prior.get_metadata(labels[0].clone()				,					_SCREAMING_SNAKE_CASE				,					105_8304				,					0     )

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

          __UpperCAmelCase				,  __UpperCAmelCase				,  __UpperCAmelCase      :					Dict    				=     top_prior.get_cond(_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE     )
          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 					lowerCamelCase__					(							self	:						List[Any]     )		->				List[str]:
          __UpperCAmelCase      :					Tuple    				=     False

          __UpperCAmelCase      :					List[Any]    				=     JukeboxModel.from_pretrained(self.model_id				,					min_duration=0     ).eval()
          set_seed(0     )
          __UpperCAmelCase      :					List[str]    				=     torch.rand((1, 5120, 1)     )
          __UpperCAmelCase      :					Union[str, Any]    				=     list(self.prepare_inputs()     )

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

          __UpperCAmelCase      :					int    				=     torch.cat((zs[0], torch.zeros(1				,					2048 - zs[0].shape[-1]     ))				,					dim=-1     ).long()
          __UpperCAmelCase      :					List[Any]    				=     [upper_a, model.vqvae.encode(_SCREAMING_SNAKE_CASE				,					start_level=1				,					bs_chunks=waveform.shape[0]     )[0], None]
          __UpperCAmelCase      :					Optional[Any]    				=     model._sample(
              _SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE				,					sample_levels=[1]				,					save_results=_SCREAMING_SNAKE_CASE				,					sample_length=40 * model.priors[1].raw_to_tokens     )
          torch.testing.assert_allclose(zs[1][0][:40]				,					torch.tensor(self.EXPECTED_PRIMED_1     )     )

          __UpperCAmelCase      :					str    				=     torch.cat((zs[1], torch.zeros(1				,					2048 - zs[1].shape[-1]     ))				,					dim=-1     ).long()
          __UpperCAmelCase      :					Union[str, Any]    				=     [upper_a, upper_a, model.vqvae.encode(_SCREAMING_SNAKE_CASE				,					start_level=0				,					bs_chunks=waveform.shape[0]     )[0]]
          __UpperCAmelCase      :					str    				=     model._sample(
              _SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE				,					sample_levels=[2]				,					save_results=_SCREAMING_SNAKE_CASE				,					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 					lowerCamelCase__					(							self	:						int     )		->				Optional[int]:
          __UpperCAmelCase      :					int    				=     JukeboxModel.from_pretrained(self.model_id				,					min_duration=0     ).eval()
          set_seed(0     )
          __UpperCAmelCase      :					str    				=     torch.rand((1, 5120, 1)     )
          with torch.no_grad():
               __UpperCAmelCase      :					Optional[Any]    				=     model.vqvae.encode(_SCREAMING_SNAKE_CASE				,					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():
               __UpperCAmelCase      :					Union[str, Any]    				=     model.vqvae.decode(_SCREAMING_SNAKE_CASE				,					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       a							(						unittest.TestCase			):


     """simple docstring"""






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

     # fmt: off
     SCREAMING_SNAKE_CASE			:					Union[str, Any]				  =  [
         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
     ]
     SCREAMING_SNAKE_CASE			:					int				  =  [
         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
     ]

     SCREAMING_SNAKE_CASE			:					Optional[Any]				  =  [
         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
     ]
     SCREAMING_SNAKE_CASE			:					Tuple				  =  [
         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
     ]

     SCREAMING_SNAKE_CASE			:					Dict				  =  [
         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
     ]
     SCREAMING_SNAKE_CASE			:					List[Any]				  =  [
         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
     ]

     SCREAMING_SNAKE_CASE			:					Optional[int]				  =  [
         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
     ]
     SCREAMING_SNAKE_CASE			:					Union[str, Any]				  =  [
         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
     ]

     SCREAMING_SNAKE_CASE			:					Any				  =  [
         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
     ]
     SCREAMING_SNAKE_CASE			:					str				  =  [
         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 					lowerCamelCase__					(							self	:						Union[str, Any]				,					snake_case	:						Any     )		->				List[Any]:
          __UpperCAmelCase      :					Union[str, Any]    				=     JukeboxTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE     )
          __UpperCAmelCase      :					List[str]    				=     tokenizer(**self.metas     )['''input_ids''']
          return tokens

     @slow
     def 					lowerCamelCase__					(							self	:						Tuple     )		->				Optional[Any]:
          __UpperCAmelCase      :					str    				=     JukeboxModel.from_pretrained(self.model_id				,					min_duration=0     ).eval()
          __UpperCAmelCase      :					Optional[Any]    				=     self.prepare_inputs(self.model_id     )

          set_seed(0     )
          __UpperCAmelCase      :					Union[str, Any]    				=     [torch.zeros(1				,					0				,					dtype=torch.long     ).cpu() for _ in range(3     )]
          __UpperCAmelCase      :					List[str]    				=     model._sample(_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE				,					[0]				,					sample_length=60 * model.priors[0].raw_to_tokens				,					save_results=_SCREAMING_SNAKE_CASE     )
          self.assertIn(zs[0][0].detach().cpu().tolist()				,					[self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2]     )

          set_seed(0     )
          __UpperCAmelCase      :					Optional[int]    				=     model._sample(_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE				,					[1]				,					sample_length=60 * model.priors[1].raw_to_tokens				,					save_results=_SCREAMING_SNAKE_CASE     )
          self.assertIn(zs[1][0].detach().cpu().tolist()				,					[self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]     )

          set_seed(0     )
          __UpperCAmelCase      :					Optional[int]    				=     model._sample(_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE				,					[2]				,					sample_length=60 * model.priors[2].raw_to_tokens				,					save_results=_SCREAMING_SNAKE_CASE     )
          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 					lowerCamelCase__					(							self	:						List[str]     )		->				Dict:
          __UpperCAmelCase      :					List[Any]    				=     JukeboxModel.from_pretrained(self.model_id				,					min_duration=0     ).eval()
          __UpperCAmelCase      :					Optional[Any]    				=     [i.cuda() for i in self.prepare_inputs(self.model_id     )]

          set_seed(0     )
          model.priors[0].cuda()
          __UpperCAmelCase      :					str    				=     [torch.zeros(1				,					0				,					dtype=torch.long     ).cuda() for _ in range(3     )]
          __UpperCAmelCase      :					List[Any]    				=     model._sample(_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE				,					[0]				,					sample_length=60 * model.priors[0].raw_to_tokens				,					save_results=_SCREAMING_SNAKE_CASE     )
          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()
          __UpperCAmelCase      :					List[str]    				=     model._sample(_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE				,					[1]				,					sample_length=60 * model.priors[1].raw_to_tokens				,					save_results=_SCREAMING_SNAKE_CASE     )
          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()
          __UpperCAmelCase      :					Dict    				=     model._sample(_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE				,					[2]				,					sample_length=60 * model.priors[2].raw_to_tokens				,					save_results=_SCREAMING_SNAKE_CASE     )
          torch.testing.assert_allclose(zs[2][0].cpu()				,					torch.tensor(self.EXPECTED_GPU_OUTPUTS_0     )     )






     @slow
     def 					lowerCamelCase__					(							self	:						Optional[Any]     )		->				Optional[Any]:
          __UpperCAmelCase      :					str    				=     '''ArthurZ/jukebox_prior_0'''
          __UpperCAmelCase      :					str    				=     JukeboxPrior.from_pretrained(_SCREAMING_SNAKE_CASE				,					min_duration=0     ).eval().half().to('''cuda'''     )

          __UpperCAmelCase      :					Optional[Any]    				=     self.prepare_inputs(_SCREAMING_SNAKE_CASE     )[0].cuda()
          __UpperCAmelCase      :					Optional[int]    				=     model.get_metadata(_SCREAMING_SNAKE_CASE				,					0				,					7680				,					0     )
          set_seed(0     )
          __UpperCAmelCase      :					str    				=     model.sample(1				,					metadata=_SCREAMING_SNAKE_CASE				,					sample_tokens=60     )
          self.assertIn(outputs[0].cpu().tolist()				,					[self.EXPECTED_GPU_OUTPUTS_2, self.EXPECTED_GPU_OUTPUTS_2_PT_2]     )
'''simple docstring'''




import itertools





def 	_UpperCAmelCase       (				_UpperCamelCase       :      int = 1_00_00_00  )      -> int:
    A_      =  [1]

    for i in itertools.count(len(_UpperCamelCase  )  ):
        A_      =  0
        for j in itertools.count(1  ):
            A_      =  -1 if j % 2 == 0 else +1
            A_      =  (j * j * 3 - j) // 2
            if index > i:
                break
            item += partitions[i - index] * sign
            item %= number
            index += j
            if index > i:
                break
            item += partitions[i - index] * sign
            item %= number

        if item == 0:
            return i
        partitions.append(_UpperCamelCase  )

    return 0


if __name__ == "__main__":
  import doctest

  doctest.testmod()

  print(F"""{solution() = }""")






