


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

							"""simple docstring"""


							SCREAMING_SNAKE_CASE   :					Tuple        =						(JukeboxModel,) if is_torch_available() else ()
							SCREAMING_SNAKE_CASE   :					Union[str, Any]        =						"openai/jukebox-1b-lyrics"
							SCREAMING_SNAKE_CASE   :					List[str]        =						{
							    "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   :					Optional[Any]        =						[
							    18_64, 15_36, 12_13, 18_70, 13_57, 15_36, 5_19, 8_80, 13_23, 7_89, 10_82, 5_34,
							    10_00, 14_45, 11_05, 11_30, 9_67, 5_15, 14_34, 16_20, 5_34, 14_95, 2_83, 14_45,
							    3_33, 13_07, 5_39, 16_31, 15_28, 3_75, 14_34, 6_73, 6_27, 7_10, 7_78, 18_83,
							    14_05, 12_76, 14_55, 12_28
							]

							SCREAMING_SNAKE_CASE   :					Optional[int]        =						[
							    14_89, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53
							]

							SCREAMING_SNAKE_CASE   :					Optional[Any]        =						[
							    11_25, 17_51, 6_97, 17_76, 11_41, 14_76, 3_91, 6_97, 11_25, 6_84, 8_67, 4_16,
							    8_44, 13_72, 12_74, 7_17, 12_74, 8_44, 12_99, 14_19, 6_97, 13_70, 3_17, 11_25,
							    1_91, 14_40, 13_70, 14_40, 13_70, 2_82, 16_21, 13_70, 3_68, 3_49, 8_67, 18_72,
							    12_62, 8_69, 17_28, 7_47
							]
							SCREAMING_SNAKE_CASE   :					List[str]        =						[
							    4_16, 4_16, 11_25, 11_25, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
							    4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
							    4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
							    4_16, 4_16, 4_16, 4_16
							]

							SCREAMING_SNAKE_CASE   :					List[str]        =						[
							    17_55, 8_42, 3_07, 18_43, 10_22, 13_95, 2_34, 15_54, 8_06, 7_39, 10_22, 4_42,
							    6_16, 5_56, 2_68, 14_99, 9_33, 4_57, 14_40, 18_37, 7_55, 9_85, 3_08, 9_02,
							    2_93, 14_43, 16_71, 11_41, 15_33, 5_55, 15_62, 10_61, 2_87, 4_17, 10_22, 20_08,
							    11_86, 10_15, 17_77, 2_68
							]
							SCREAMING_SNAKE_CASE   :					List[str]        =						[
							    8_54, 8_42, 13_53, 1_14, 13_53, 8_42, 1_85, 8_42, 1_85, 1_14, 5_91, 8_42,
							    1_85, 4_17, 1_85, 8_42, 3_07, 8_42, 5_91, 8_42, 1_85, 8_42, 3_07, 8_42,
							    5_91, 8_42, 13_53, 8_42, 1_85, 8_42, 5_91, 8_42, 5_91, 1_14, 5_91, 8_42,
							    1_85, 8_42, 5_91, 89
							]

							SCREAMING_SNAKE_CASE   :					Tuple        =						[1_05_83_04, 0, 78_64_32, 71_69, 5_07, 76, 27, 40, 30, 76]

							SCREAMING_SNAKE_CASE   :					str        =						[
							    3_90, 11_60, 10_02, 19_07, 17_88, 17_88, 17_88, 19_07, 10_02, 10_02, 18_54, 10_02,
							    10_02, 10_02, 10_02, 10_02, 10_02, 11_60, 11_60, 16_06, 5_96, 5_96, 11_60, 10_02,
							    15_16, 5_96, 10_02, 10_02, 10_02, 19_07, 17_88, 17_88, 17_88, 18_54, 17_88, 19_07,
							    19_07, 17_88, 5_96, 16_26
							]
							SCREAMING_SNAKE_CASE   :					Dict        =						[
							    12_36, 16_68, 14_84, 19_20, 18_48, 14_09, 1_39, 8_64, 18_28, 12_72, 15_99, 8_24,
							    16_72, 1_39, 5_55, 14_84, 8_24, 19_20, 5_55, 5_96, 15_79, 15_99, 12_31, 15_99,
							    16_37, 14_07, 2_12, 8_24, 15_99, 1_16, 14_33, 8_24, 2_58, 15_99, 14_33, 18_95,
							    10_63, 14_33, 14_33, 15_99
							]
							SCREAMING_SNAKE_CASE   :					Dict        =						[
							    16_84, 18_73, 11_19, 11_89, 3_95, 6_11, 19_01, 9_72, 8_90, 13_37, 13_92, 19_27,
							    96, 9_72, 6_72, 7_80, 11_19, 8_90, 1_58, 7_71, 10_73, 19_27, 3_53, 13_31,
							    12_69, 14_59, 13_33, 16_45, 8_12, 15_77, 13_37, 6_06, 3_53, 9_81, 14_66, 6_19,
							    1_97, 3_91, 3_02, 19_30
							]
							SCREAMING_SNAKE_CASE   :					Dict        =						[
							    3_90, 11_60, 10_02, 19_07, 17_88, 17_88, 17_88, 19_07, 10_02, 10_02, 18_54, 10_02,
							    10_02, 10_02, 10_02, 10_02, 10_02, 11_60, 11_60, 16_06, 5_96, 5_96, 11_60, 10_02,
							    15_16, 5_96, 10_02, 10_02, 10_02, 19_07, 17_88, 17_88, 17_88, 18_54, 17_88, 19_07,
							    19_07, 17_88, 5_96, 16_26
							]
							SCREAMING_SNAKE_CASE   :					List[str]        =						[
							    -0.0492, -0.0524, -0.0565, -0.0640, -0.0686, -0.0684, -0.0677, -0.0664,
							    -0.0605, -0.0490, -0.0330, -0.0168, -0.0083, -0.0075, -0.0051, 0.0025,
							    0.0136, 0.0261, 0.0386, 0.0497, 0.0580, 0.0599, 0.0583, 0.0614,
							    0.0740, 0.0889, 0.1023, 0.1162, 0.1211, 0.1212, 0.1251, 0.1336,
							    0.1502, 0.1686, 0.1883, 0.2148, 0.2363, 0.2458, 0.2507, 0.2531
							]
							SCREAMING_SNAKE_CASE   :					int        =						[
							    0.0256, -0.0544, 0.1600, -0.0032, 0.1066, 0.0825, -0.0013, 0.3440,
							    0.0210, 0.0412, -0.1777, -0.0892, -0.0164, 0.0285, -0.0613, -0.0617,
							    -0.0137, -0.0201, -0.0175, 0.0215, -0.0627, 0.0520, -0.0730, 0.0970,
							    -0.0100, 0.0442, -0.0586, 0.0207, -0.0015, -0.0082
							]
							SCREAMING_SNAKE_CASE   :					Union[str, Any]        =						[
							    0.0415, 0.0877, 0.0022, -0.0055, 0.0751, 0.0334, 0.0324, -0.0068,
							    0.0011, 0.0017, -0.0676, 0.0655, -0.0143, 0.0399, 0.0303, 0.0743,
							    -0.0168, -0.0394, -0.1113, 0.0124, 0.0442, 0.0267, -0.0003, -0.1536,
							    -0.0116, -0.1837, -0.0180, -0.1026, -0.0777, -0.0456
							]
							SCREAMING_SNAKE_CASE   :					Tuple        =						[
							    76, 27, 40, 30, 76, 46, 44, 47, 40, 37, 38, 31, 45, 45, 76, 38, 31, 33,
							    45, 76, 41, 32, 76, 45, 46, 41, 40, 31, 78, 76
							]





							def 		__magic_name__			(    self	)->	List[str]:
											_SCREAMING_SNAKE_CASE      = JukeboxTokenizer.from_pretrained(self.model_id	)
											_SCREAMING_SNAKE_CASE      = tokenizer(**self.metas	)['input_ids']
											return tokens





							@slow
							def 		__magic_name__			(    self	)->	List[str]:
											_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 		__magic_name__			(    self	)->	List[str]:
											_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_				, 1058304				, 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 		__magic_name__			(    self	)->	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, 5120, 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				, 2048 - 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				, 2048 - 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 		__magic_name__			(    self	)->	int:
											_SCREAMING_SNAKE_CASE      = JukeboxModel.from_pretrained(self.model_id				, min_duration=0	).eval()
											set_seed(0	)
											_SCREAMING_SNAKE_CASE      = torch.rand((1, 5120, 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    __lowercase  (       unittest.TestCase							):

							"""simple docstring"""


							SCREAMING_SNAKE_CASE   :					List[Any]        =						(JukeboxModel,) if is_torch_available() else ()
							SCREAMING_SNAKE_CASE   :					Union[str, Any]        =						"openai/jukebox-5b-lyrics"
							SCREAMING_SNAKE_CASE   :					List[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        =						[
							    14_89, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    14_89, 14_89, 14_89, 14_89, 11_50, 18_53, 15_09, 11_50, 13_57, 15_09, 6, 12_72
							]
							SCREAMING_SNAKE_CASE   :					Tuple        =						[
							    14_89, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53
							]

							SCREAMING_SNAKE_CASE   :					List[str]        =						[
							    11_25, 4_16, 11_25, 11_25, 11_25, 11_25, 11_25, 4_16, 4_16, 4_16, 4_16, 4_16,
							    4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
							    4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
							    4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
							    4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16
							]
							SCREAMING_SNAKE_CASE   :					Optional[Any]        =						[
							    4_16, 4_16, 11_25, 11_25, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
							    4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
							    4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
							    4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
							    4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16
							]

							SCREAMING_SNAKE_CASE   :					str        =						[
							    17_55, 10_61, 2_34, 17_55, 10_61, 17_55, 1_85, 2_90, 3_07, 3_07, 6_16, 6_16,
							    6_16, 6_16, 6_16, 6_16, 3_07, 2_90, 4_17, 17_55, 2_34, 17_55, 1_85, 2_90,
							    2_90, 2_90, 3_07, 6_16, 6_16, 6_16, 6_16, 6_16, 2_90, 2_34, 2_34, 17_55,
							    2_34, 2_34, 17_55, 2_34, 1_85, 1_85, 3_07, 6_16, 6_16, 6_16, 6_16, 2_90,
							    17_55, 17_55, 17_55, 2_34, 2_34, 17_55, 15_72, 2_90, 3_07, 6_16, 34, 6_16
							]
							SCREAMING_SNAKE_CASE   :					Dict        =						[
							    8_54, 8_42, 13_53, 1_14, 13_53, 8_42, 1_85, 8_42, 1_85, 1_14, 5_91, 8_42, 1_85,
							    4_17, 1_85, 8_42, 3_07, 8_42, 5_91, 8_42, 1_85, 8_42, 1_85, 8_42, 5_91, 8_42,
							    13_53, 8_42, 1_85, 8_42, 5_91, 8_42, 5_91, 1_14, 5_91, 8_42, 1_85, 8_42, 5_91,
							    89, 5_91, 8_42, 5_91, 8_42, 5_91, 4_17, 13_72, 8_42, 13_72, 8_42, 34, 8_42,
							    1_85, 89, 5_91, 8_42, 1_85, 8_42, 5_91, 6_32
							]

							SCREAMING_SNAKE_CASE   :					Any        =						[
							    14_89, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53
							]
							SCREAMING_SNAKE_CASE   :					Any        =						[
							    14_89, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
							    6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 18_53, 11_77, 15_36, 12_28,
							    7_10, 4_75, 14_89, 12_29, 12_24, 2_31, 12_24, 2_52, 14_34, 6_53, 4_75,
							    11_06, 18_77, 15_99, 12_28, 16_00, 16_83, 11_82, 18_53, 4_75, 18_64,
							    2_52, 12_29, 14_34, 20_01
							]

							SCREAMING_SNAKE_CASE   :					Optional[int]        =						[
							    11_25, 11_25, 4_16, 11_25, 11_25, 4_16, 11_25, 11_25, 4_16, 4_16, 11_25, 4_16,
							    4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
							    4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
							    4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
							    4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16
							]
							SCREAMING_SNAKE_CASE   :					Tuple        =						[
							    4_91, 17_55, 34, 16_13, 17_55, 4_17, 9_92, 16_13, 2_22, 8_42, 13_53, 16_13,
							    8_44, 6_32, 1_85, 16_13, 8_44, 6_32, 1_85, 16_13, 1_85, 8_42, 6_77, 16_13,
							    1_85, 1_14, 13_53, 16_13, 3_07, 89, 8_44, 16_13, 3_07, 13_32, 2_34, 19_79,
							    3_07, 89, 13_53, 6_16, 34, 8_42, 1_85, 8_42, 34, 8_42, 1_85, 8_42,
							    3_07, 1_14, 1_85, 89, 34, 12_68, 1_85, 89, 34, 8_42, 1_85, 89
							]





							def 		__magic_name__			(    self				, A_	)->	str:
											_SCREAMING_SNAKE_CASE      = JukeboxTokenizer.from_pretrained(A_	)
											_SCREAMING_SNAKE_CASE      = tokenizer(**self.metas	)['input_ids']
											return tokens





							@slow
							def 		__magic_name__			(    self	)->	List[Any]:
											_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 		__magic_name__			(    self	)->	Union[str, Any]:
											_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 		__magic_name__			(    self	)->	Union[str, Any]:
											_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				, 7680				, 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 sys
from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union

import pandas as pd
import pyarrow as pa

import datasets
import datasets.config
from datasets.features.features import require_storage_cast
from datasets.table import table_cast


if TYPE_CHECKING:
							import sqlitea

							import sqlalchemy


snake_case			:      List[str] 							=   datasets.utils.logging.get_logger(__name__)



@dataclass
class    __lowercase  (       datasets.BuilderConfig							):

							"""simple docstring"""


							SCREAMING_SNAKE_CASE   :					Union[str, "sqlalchemy.sql.Selectable"]        =						None
							SCREAMING_SNAKE_CASE   :					Union[str, "sqlalchemy.engine.Connection", "sqlalchemy.engine.Engine", "sqlite3.Connection"]        =						None
							SCREAMING_SNAKE_CASE   :					Optional[Union[str, List[str]]]        =						None
							SCREAMING_SNAKE_CASE   :					bool        =						True
							SCREAMING_SNAKE_CASE   :					Optional[Union[List, Tuple, Dict]]        =						None
							SCREAMING_SNAKE_CASE   :					Optional[Union[List, Dict]]        =						None
							SCREAMING_SNAKE_CASE   :					Optional[List[str]]        =						None
							SCREAMING_SNAKE_CASE   :					Optional[int]        =						1_00_00
							SCREAMING_SNAKE_CASE   :					Optional[datasets.Features]        =						None





							def 		__magic_name__			(    self	)->	int:
											if self.sql is None:
															raise ValueError('sql must be specified'	)
											if self.con is None:
															raise ValueError('con must be specified'	)





							def 		__magic_name__			(    self				, A_				, A_ = None				, )->	str:
											_SCREAMING_SNAKE_CASE      = config_kwargs.copy()
											# We need to stringify the Selectable object to make its hash deterministic

											# The process of stringifying is explained here: http://docs.sqlalchemy.org/en/latest/faq/sqlexpressions.html
											_SCREAMING_SNAKE_CASE      = config_kwargs['sql']
											if not isinstance(A_				, A_	):
															if datasets.config.SQLALCHEMY_AVAILABLE and "sqlalchemy" in sys.modules:
																			import sqlalchemy

																			if isinstance(A_				, sqlalchemy.sql.Selectable	):
																							_SCREAMING_SNAKE_CASE      = sqlalchemy.create_engine(config_kwargs['con'].split('://'	)[0] + '://'	)
																							_SCREAMING_SNAKE_CASE      = str(sql.compile(dialect=engine.dialect	)	)
																							_SCREAMING_SNAKE_CASE      = sql_str
																			else:
																							raise TypeError(
																							    F'''Supported types for \'sql\' are string and sqlalchemy.sql.Selectable but got {type(A_	)}: {sql}'''	)
															else:
																			raise TypeError(
																			    F'''Supported types for \'sql\' are string and sqlalchemy.sql.Selectable but got {type(A_	)}: {sql}'''	)
											_SCREAMING_SNAKE_CASE      = config_kwargs['con']
											if not isinstance(A_				, A_	):
															_SCREAMING_SNAKE_CASE      = id(A_	)
															logger.info(
															    F'''SQL connection \'con\' of type {type(A_	)} couldn\'t be hashed properly. To enable hashing, specify \'con\' as URI string instead.'''	)

											return super().create_config_id(A_				, custom_features=A_	)





							@property
							def 		__magic_name__			(    self	)->	Dict:
											_SCREAMING_SNAKE_CASE      = {
											    'index_col': self.index_col,
											    'columns': self.columns,
											    'params': self.params,
											    'coerce_float': self.coerce_float,
											    'parse_dates': self.parse_dates,
											}
											return pd_read_sql_kwargs







class    __lowercase  (       datasets.ArrowBasedBuilder							):

							"""simple docstring"""


							SCREAMING_SNAKE_CASE   :					Any        =						SqlConfig





							def 		__magic_name__			(    self	)->	Dict:
											return datasets.DatasetInfo(features=self.config.features	)





							def 		__magic_name__			(    self				, A_	)->	Optional[Any]:
											return [datasets.SplitGenerator(name=datasets.Split.TRAIN				, gen_kwargs={}	)]





							def 		__magic_name__			(    self				, A_	)->	pa.Table:
											if self.config.features is not None:
															_SCREAMING_SNAKE_CASE      = self.config.features.arrow_schema
															if all(not require_storage_cast(A_	) for feature in self.config.features.values()	):
																			# cheaper cast
																			_SCREAMING_SNAKE_CASE      = pa.Table.from_arrays([pa_table[field.name] for field in schema]				, schema=A_	)
															else:
																			# more expensive cast; allows str <-> int/float or str to Audio for example
																			_SCREAMING_SNAKE_CASE      = table_cast(A_				, A_	)
											return pa_table





							def 		__magic_name__			(    self	)->	Optional[int]:
											_SCREAMING_SNAKE_CASE      = self.config.chunksize
											_SCREAMING_SNAKE_CASE      = pd.read_sql(
											    self.config.sql				, self.config.con				, chunksize=A_				, **self.config.pd_read_sql_kwargs	)
											_SCREAMING_SNAKE_CASE      = [sql_reader] if chunksize is None else sql_reader
											for chunk_idx, df in enumerate(A_	):
															_SCREAMING_SNAKE_CASE      = pa.Table.from_pandas(A_	)
															yield chunk_idx, self._cast_table(A_	)
