code: 
  '''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 			UpperCAmelCase				(					unittest.TestCase ):
  
  
  
  
  
  		"""simple docstring"""
  
  
  
  
  
  		A__  :   Union[str, Any]          =    (JukeboxModel,) if is_torch_available() else ()
  		A__  :   List[str]          =    'openai/jukebox-1b-lyrics'
  		A__  :   Dict          =    {
  		    '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
  		A__  :   Union[str, Any]          =    [
  		    1864, 1536, 1213, 1870, 1357, 1536, 519, 880, 1323, 789, 1082, 534,
  		    1000, 1445, 1105, 1130, 967, 515, 1434, 1620, 534, 1495, 283, 1445,
  		    333, 1307, 539, 1631, 1528, 375, 1434, 673, 627, 710, 778, 1883,
  		    1405, 1276, 1455, 1228
  		]
  
  		A__  :   Union[str, Any]          =    [
  		    1489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653
  		]
  
  		A__  :   Tuple          =    [
  		    1125, 1751, 697, 1776, 1141, 1476, 391, 697, 1125, 684, 867, 416,
  		    844, 1372, 1274, 717, 1274, 844, 1299, 1419, 697, 1370, 317, 1125,
  		    191, 1440, 1370, 1440, 1370, 282, 1621, 1370, 368, 349, 867, 1872,
  		    1262, 869, 1728, 747
  		]
  		A__  :   Union[str, Any]          =    [
  		    416, 416, 1125, 1125, 416, 416, 416, 416, 416, 416, 416, 416,
  		    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
  		    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
  		    416, 416, 416, 416
  		]
  
  		A__  :   Dict          =    [
  		    1755, 842, 307, 1843, 1022, 1395, 234, 1554, 806, 739, 1022, 442,
  		    616, 556, 268, 1499, 933, 457, 1440, 1837, 755, 985, 308, 902,
  		    293, 1443, 1671, 1141, 1533, 555, 1562, 1061, 287, 417, 1022, 2008,
  		    1186, 1015, 1777, 268
  		]
  		A__  :   Optional[Any]          =    [
  		    854, 842, 1353, 114, 1353, 842, 185, 842, 185, 114, 591, 842,
  		    185, 417, 185, 842, 307, 842, 591, 842, 185, 842, 307, 842,
  		    591, 842, 1353, 842, 185, 842, 591, 842, 591, 114, 591, 842,
  		    185, 842, 591, 89
  		]
  
  		A__  :   List[str]          =    [1058304, 0, 786432, 7169, 507, 76, 27, 40, 30, 76]
  
  		A__  :   Dict          =    [
  		    390, 1160, 1002, 1907, 1788, 1788, 1788, 1907, 1002, 1002, 1854, 1002,
  		    1002, 1002, 1002, 1002, 1002, 1160, 1160, 1606, 596, 596, 1160, 1002,
  		    1516, 596, 1002, 1002, 1002, 1907, 1788, 1788, 1788, 1854, 1788, 1907,
  		    1907, 1788, 596, 1626
  		]
  		A__  :   Union[str, Any]          =    [
  		    1236, 1668, 1484, 1920, 1848, 1409, 139, 864, 1828, 1272, 1599, 824,
  		    1672, 139, 555, 1484, 824, 1920, 555, 596, 1579, 1599, 1231, 1599,
  		    1637, 1407, 212, 824, 1599, 116, 1433, 824, 258, 1599, 1433, 1895,
  		    1063, 1433, 1433, 1599
  		]
  		A__  :   List[str]          =    [
  		    1684, 1873, 1119, 1189, 395, 611, 1901, 972, 890, 1337, 1392, 1927,
  		    96, 972, 672, 780, 1119, 890, 158, 771, 1073, 1927, 353, 1331,
  		    1269, 1459, 1333, 1645, 812, 1577, 1337, 606, 353, 981, 1466, 619,
  		    197, 391, 302, 1930
  		]
  		A__  :   int          =    [
  		    390, 1160, 1002, 1907, 1788, 1788, 1788, 1907, 1002, 1002, 1854, 1002,
  		    1002, 1002, 1002, 1002, 1002, 1160, 1160, 1606, 596, 596, 1160, 1002,
  		    1516, 596, 1002, 1002, 1002, 1907, 1788, 1788, 1788, 1854, 1788, 1907,
  		    1907, 1788, 596, 1626
  		]
  		A__  :   Union[str, Any]          =    [
  		    -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__  :   List[Any]          =    [
  		    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__  :   Union[str, Any]          =    [
  		    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__  :   int          =    [
  		    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 _lowercase  (    self     )				->    int:
  				_UpperCamelCase       : Dict    =      JukeboxTokenizer.from_pretrained(self.model_id     )
  				_UpperCamelCase       : Optional[int]    =      tokenizer(**self.metas     )['''input_ids''']
  				return tokens
  		@slow
  		def _lowercase  (    self     )				->    Optional[Any]:
  				_UpperCamelCase       : List[str]    =      JukeboxModel.from_pretrained(self.model_id  ,					min_duration=0     ).eval()
  				_UpperCamelCase       : List[Any]    =      self.prepare_inputs()
  
  				set_seed(0     )
  				_UpperCamelCase       : str    =      [torch.zeros(1  ,					0  ,					dtype=torch.long     ).cpu() for _ in range(3     )]
  				_UpperCamelCase       : List[Any]    =      model._sample(lowercase_  ,					lowercase_  ,					[0]  ,					sample_length=40 * model.priors[0].raw_to_tokens  ,					save_results=lowercase_     )
  				self.assertIn(zs[0][0].detach().cpu().tolist()  ,					[self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2]     )
  
  				set_seed(0     )
  				_UpperCamelCase       : Any    =      model._sample(lowercase_  ,					lowercase_  ,					[1]  ,					sample_length=40 * model.priors[1].raw_to_tokens  ,					save_results=lowercase_     )
  				self.assertIn(zs[1][0].detach().cpu().tolist()  ,					[self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]     )
  
  				set_seed(0     )
  				_UpperCamelCase       : Tuple    =      model._sample(lowercase_  ,					lowercase_  ,					[2]  ,					sample_length=40 * model.priors[2].raw_to_tokens  ,					save_results=lowercase_     )
  				self.assertIn(zs[2][0].detach().cpu().tolist()  ,					[self.EXPECTED_OUTPUT_0, self.EXPECTED_OUTPUT_0_PT_2]     )
  		@slow
  		def _lowercase  (    self     )				->    str:
  				_UpperCamelCase       : List[Any]    =      False
  				_UpperCamelCase       : int    =      JukeboxModel.from_pretrained(self.model_id  ,					min_duration=0     ).eval()
  
  				_UpperCamelCase       : Tuple    =      self.prepare_inputs()
  				set_seed(0     )
  				_UpperCamelCase       : Optional[Any]    =      [torch.zeros(1  ,					0  ,					dtype=torch.long     ) for _ in range(3     )]
  
  				_UpperCamelCase       : List[str]    =      model.priors[0]
  				_UpperCamelCase       : Optional[int]    =      0
  				_UpperCamelCase       : Any    =      top_prior.get_music_tokens_conds(lowercase_  ,					start=lowercase_  ,					end=start + top_prior.n_ctx     )
  				_UpperCamelCase       : str    =      top_prior.get_metadata(labels[0].clone()  ,					lowercase_  ,					1058304  ,					0     )
  
  				self.assertIsNone(lowercase_     )
  				self.assertListEqual(metadata.numpy()[0][:10].tolist()  ,					self.EXPECTED_Y_COND     )
  
  				_UpperCamelCase       : Dict    =      top_prior.get_cond(lowercase_  ,					lowercase_     )
  				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 _lowercase  (    self     )				->    Optional[int]:
  				_UpperCamelCase       : str    =      False
  
  				_UpperCamelCase       : Optional[Any]    =      JukeboxModel.from_pretrained(self.model_id  ,					min_duration=0     ).eval()
  				set_seed(0     )
  				_UpperCamelCase       : str    =      torch.rand((1, 5120, 1)     )
  				_UpperCamelCase       : Union[str, Any]    =      list(self.prepare_inputs()     )
  
  				_UpperCamelCase       : int    =      [model.vqvae.encode(lowercase_  ,					start_level=2  ,					bs_chunks=waveform.shape[0]     )[0], None, None]
  				_UpperCamelCase       : Optional[Any]    =      model._sample(
  				    lowercase_  ,					lowercase_  ,					sample_levels=[0]  ,					save_results=lowercase_  ,					sample_length=40 * model.priors[0].raw_to_tokens     )
  				torch.testing.assert_allclose(zs[0][0][:40]  ,					torch.tensor(self.EXPECTED_PRIMED_0     )     )
  
  				_UpperCamelCase       : List[Any]    =      torch.cat((zs[0], torch.zeros(1  ,					2048 - zs[0].shape[-1]     ))  ,					dim=-1     ).long()
  				_UpperCamelCase       : Dict    =      [upper_a, model.vqvae.encode(lowercase_  ,					start_level=1  ,					bs_chunks=waveform.shape[0]     )[0], None]
  				_UpperCamelCase       : Optional[int]    =      model._sample(
  				    lowercase_  ,					lowercase_  ,					sample_levels=[1]  ,					save_results=lowercase_  ,					sample_length=40 * model.priors[1].raw_to_tokens     )
  				torch.testing.assert_allclose(zs[1][0][:40]  ,					torch.tensor(self.EXPECTED_PRIMED_1     )     )
  
  				_UpperCamelCase       : Any    =      torch.cat((zs[1], torch.zeros(1  ,					2048 - zs[1].shape[-1]     ))  ,					dim=-1     ).long()
  				_UpperCamelCase       : Optional[int]    =      [upper_a, upper_a, model.vqvae.encode(lowercase_  ,					start_level=0  ,					bs_chunks=waveform.shape[0]     )[0]]
  				_UpperCamelCase       : List[str]    =      model._sample(
  				    lowercase_  ,					lowercase_  ,					sample_levels=[2]  ,					save_results=lowercase_  ,					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 _lowercase  (    self     )				->    Union[str, Any]:
  				_UpperCamelCase       : str    =      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       : Dict    =      model.vqvae.encode(lowercase_  ,					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       : Dict    =      model.vqvae.decode(lowercase_  ,					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__  :   Any          =    (JukeboxModel,) if is_torch_available() else ()
  		A__  :   List[Any]          =    'openai/jukebox-5b-lyrics'
  		A__  :   Optional[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
  		A__  :   List[Any]          =    [
  		    1489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    1489, 1489, 1489, 1489, 1150, 1853, 1509, 1150, 1357, 1509, 6, 1272
  		]
  		A__  :   Tuple          =    [
  		    1489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653
  		]
  
  		A__  :   List[Any]          =    [
  		    1125, 416, 1125, 1125, 1125, 1125, 1125, 416, 416, 416, 416, 416,
  		    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
  		    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
  		    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
  		    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416
  		]
  		A__  :   Tuple          =    [
  		    416, 416, 1125, 1125, 416, 416, 416, 416, 416, 416, 416, 416, 416,
  		    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
  		    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
  		    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
  		    416, 416, 416, 416, 416, 416, 416, 416
  		]
  
  		A__  :   Tuple          =    [
  		    1755, 1061, 234, 1755, 1061, 1755, 185, 290, 307, 307, 616, 616,
  		    616, 616, 616, 616, 307, 290, 417, 1755, 234, 1755, 185, 290,
  		    290, 290, 307, 616, 616, 616, 616, 616, 290, 234, 234, 1755,
  		    234, 234, 1755, 234, 185, 185, 307, 616, 616, 616, 616, 290,
  		    1755, 1755, 1755, 234, 234, 1755, 1572, 290, 307, 616, 34, 616
  		]
  		A__  :   List[Any]          =    [
  		    854, 842, 1353, 114, 1353, 842, 185, 842, 185, 114, 591, 842, 185,
  		    417, 185, 842, 307, 842, 591, 842, 185, 842, 185, 842, 591, 842,
  		    1353, 842, 185, 842, 591, 842, 591, 114, 591, 842, 185, 842, 591,
  		    89, 591, 842, 591, 842, 591, 417, 1372, 842, 1372, 842, 34, 842,
  		    185, 89, 591, 842, 185, 842, 591, 632
  		]
  
  		A__  :   Tuple          =    [
  		    1489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653
  		]
  		A__  :   List[str]          =    [
  		    1489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
  		    653, 653, 653, 653, 653, 653, 653, 1853, 1177, 1536, 1228,
  		    710, 475, 1489, 1229, 1224, 231, 1224, 252, 1434, 653, 475,
  		    1106, 1877, 1599, 1228, 1600, 1683, 1182, 1853, 475, 1864,
  		    252, 1229, 1434, 2001
  		]
  
  		A__  :   Optional[Any]          =    [
  		    1125, 1125, 416, 1125, 1125, 416, 1125, 1125, 416, 416, 1125, 416,
  		    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
  		    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
  		    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
  		    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416
  		]
  		A__  :   Dict          =    [
  		    491, 1755, 34, 1613, 1755, 417, 992, 1613, 222, 842, 1353, 1613,
  		    844, 632, 185, 1613, 844, 632, 185, 1613, 185, 842, 677, 1613,
  		    185, 114, 1353, 1613, 307, 89, 844, 1613, 307, 1332, 234, 1979,
  		    307, 89, 1353, 616, 34, 842, 185, 842, 34, 842, 185, 842,
  		    307, 114, 185, 89, 34, 1268, 185, 89, 34, 842, 185, 89
  		]
  		def _lowercase  (    self  ,					_snake_case     )				->    Optional[Any]:
  				_UpperCamelCase       : Optional[Any]    =      JukeboxTokenizer.from_pretrained(lowercase_     )
  				_UpperCamelCase       : Union[str, Any]    =      tokenizer(**self.metas     )['''input_ids''']
  				return tokens
  		@slow
  		def _lowercase  (    self     )				->    Tuple:
  				_UpperCamelCase       : List[str]    =      JukeboxModel.from_pretrained(self.model_id  ,					min_duration=0     ).eval()
  				_UpperCamelCase       : List[str]    =      self.prepare_inputs(self.model_id     )
  
  				set_seed(0     )
  				_UpperCamelCase       : Optional[Any]    =      [torch.zeros(1  ,					0  ,					dtype=torch.long     ).cpu() for _ in range(3     )]
  				_UpperCamelCase       : str    =      model._sample(lowercase_  ,					lowercase_  ,					[0]  ,					sample_length=60 * model.priors[0].raw_to_tokens  ,					save_results=lowercase_     )
  				self.assertIn(zs[0][0].detach().cpu().tolist()  ,					[self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2]     )
  
  				set_seed(0     )
  				_UpperCamelCase       : Dict    =      model._sample(lowercase_  ,					lowercase_  ,					[1]  ,					sample_length=60 * model.priors[1].raw_to_tokens  ,					save_results=lowercase_     )
  				self.assertIn(zs[1][0].detach().cpu().tolist()  ,					[self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]     )
  
  				set_seed(0     )
  				_UpperCamelCase       : int    =      model._sample(lowercase_  ,					lowercase_  ,					[2]  ,					sample_length=60 * model.priors[2].raw_to_tokens  ,					save_results=lowercase_     )
  				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 _lowercase  (    self     )				->    List[Any]:
  				_UpperCamelCase       : Tuple    =      JukeboxModel.from_pretrained(self.model_id  ,					min_duration=0     ).eval()
  				_UpperCamelCase       : Union[str, Any]    =      [i.cuda() for i in self.prepare_inputs(self.model_id     )]
  
  				set_seed(0     )
  				model.priors[0].cuda()
  				_UpperCamelCase       : List[Any]    =      [torch.zeros(1  ,					0  ,					dtype=torch.long     ).cuda() for _ in range(3     )]
  				_UpperCamelCase       : str    =      model._sample(lowercase_  ,					lowercase_  ,					[0]  ,					sample_length=60 * model.priors[0].raw_to_tokens  ,					save_results=lowercase_     )
  				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       : Tuple    =      model._sample(lowercase_  ,					lowercase_  ,					[1]  ,					sample_length=60 * model.priors[1].raw_to_tokens  ,					save_results=lowercase_     )
  				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       : List[str]    =      model._sample(lowercase_  ,					lowercase_  ,					[2]  ,					sample_length=60 * model.priors[2].raw_to_tokens  ,					save_results=lowercase_     )
  				torch.testing.assert_allclose(zs[2][0].cpu()  ,					torch.tensor(self.EXPECTED_GPU_OUTPUTS_0     )     )
  
  
  
  
  		@slow
  		def _lowercase  (    self     )				->    List[Any]:
  				_UpperCamelCase       : Any    =      '''ArthurZ/jukebox_prior_0'''
  				_UpperCamelCase       : str    =      JukeboxPrior.from_pretrained(lowercase_  ,					min_duration=0     ).eval().half().to('''cuda'''     )
  
  				_UpperCamelCase       : Union[str, Any]    =      self.prepare_inputs(lowercase_     )[0].cuda()
  				_UpperCamelCase       : int    =      model.get_metadata(lowercase_  ,					0  ,					7680  ,					0     )
  				set_seed(0     )
  				_UpperCamelCase       : List[str]    =      model.sample(1  ,					metadata=lowercase_  ,					sample_tokens=60     )
  				self.assertIn(outputs[0].cpu().tolist()  ,					[self.EXPECTED_GPU_OUTPUTS_2, self.EXPECTED_GPU_OUTPUTS_2_PT_2]     )
  
  
  
  
style_context: 
  '''simple docstring'''
  
  _UpperCAmelCase   :      Optional[int] 						=  """
  # Transformers installation
  ! pip install transformers datasets evaluate
  # To install from source instead of the last release, comment the command above and uncomment the following one.
  # ! pip install git+https://github.com/huggingface/transformers.git
  """
  
  _UpperCAmelCase   :      Union[str, Any] 						=  [{"""type""": """code""", """content""": INSTALL_CONTENT}]
  _UpperCAmelCase   :      Union[str, Any] 						=  {
      """{processor_class}""": """FakeProcessorClass""",
      """{model_class}""": """FakeModelClass""",
      """{object_class}""": """FakeObjectClass""",
  }
  
  
  
  
  
  