
'''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"""




   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,\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		= [
       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
   ]

   A		= [
       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
   ]

   A		= [
       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
   ]
   A		= [
       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
   ]

   A		= [
       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
   ]
   A		= [
       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
   ]

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

   A		= [
       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
   ]
   A		= [
       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
   ]
   A		= [
       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
   ]
   A		= [
       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
   ]
   A		= [
       -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
   ]
   A		= [
       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
   ]
   A		= [
       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
   ]
   A		= [
       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 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Dict								=		JukeboxTokenizer.from_pretrained(self.model_id      )
    SCREAMING_SNAKE_CASE  :  int								=		tokenizer(**self.metas      )['input_ids']
    return tokens






   @slow
   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Tuple								=		JukeboxModel.from_pretrained(self.model_id    ,min_duration=0      ).eval()
    SCREAMING_SNAKE_CASE  :  int								=		self.prepare_inputs()

    set_seed(0      )
    SCREAMING_SNAKE_CASE  :  int								=		[torch.zeros(1    ,0    ,dtype=torch.long      ).cpu() for _ in range(3      )]
    SCREAMING_SNAKE_CASE  :  Optional[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      )
    SCREAMING_SNAKE_CASE  :  Dict								=		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      )
    SCREAMING_SNAKE_CASE  :  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 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Tuple								=		False
    SCREAMING_SNAKE_CASE  :  Dict								=		JukeboxModel.from_pretrained(self.model_id    ,min_duration=0      ).eval()

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

    SCREAMING_SNAKE_CASE  :  str								=		model.priors[0]
    SCREAMING_SNAKE_CASE  :  Optional[int]								=		0
    SCREAMING_SNAKE_CASE  :  Dict								=		top_prior.get_music_tokens_conds(__SCREAMING_SNAKE_CASE    ,start=__SCREAMING_SNAKE_CASE    ,end=start + top_prior.n_ctx      )
    SCREAMING_SNAKE_CASE  :  Any								=		top_prior.get_metadata(labels[0].clone()    ,__SCREAMING_SNAKE_CASE    ,1058304    ,0      )

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

    SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE  :  List[Any]								=		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 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		False

    SCREAMING_SNAKE_CASE  :  Tuple								=		JukeboxModel.from_pretrained(self.model_id    ,min_duration=0      ).eval()
    set_seed(0      )
    SCREAMING_SNAKE_CASE  :  Optional[int]								=		torch.rand((1, 5120, 1)      )
    SCREAMING_SNAKE_CASE  :  Any								=		list(self.prepare_inputs()      )

    SCREAMING_SNAKE_CASE  :  str								=		[model.vqvae.encode(__SCREAMING_SNAKE_CASE    ,start_level=2    ,bs_chunks=waveform.shape[0]      )[0], None, None]
    SCREAMING_SNAKE_CASE  :  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      )      )

    SCREAMING_SNAKE_CASE  :  Tuple								=		torch.cat((zs[0], torch.zeros(1    ,2048 - zs[0].shape[-1]      ))    ,dim=-1      ).long()
    SCREAMING_SNAKE_CASE  :  Any								=		[upper_a, model.vqvae.encode(__SCREAMING_SNAKE_CASE    ,start_level=1    ,bs_chunks=waveform.shape[0]      )[0], None]
    SCREAMING_SNAKE_CASE  :  Dict								=		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      )      )

    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		torch.cat((zs[1], torch.zeros(1    ,2048 - zs[1].shape[-1]      ))    ,dim=-1      ).long()
    SCREAMING_SNAKE_CASE  :  Dict								=		[upper_a, upper_a, model.vqvae.encode(__SCREAMING_SNAKE_CASE    ,start_level=0    ,bs_chunks=waveform.shape[0]      )[0]]
    SCREAMING_SNAKE_CASE  :  List[Any]								=		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 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Any								=		JukeboxModel.from_pretrained(self.model_id    ,min_duration=0      ).eval()
    set_seed(0      )
    SCREAMING_SNAKE_CASE  :  Dict								=		torch.rand((1, 5120, 1)      )
    with torch.no_grad():
     SCREAMING_SNAKE_CASE  :  List[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():
     SCREAMING_SNAKE_CASE  :  str								=		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"""




   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,\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		= [
       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
   ]
   A		= [
       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
   ]

   A		= [
       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
   ]
   A		= [
       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
   ]

   A		= [
       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
   ]
   A		= [
       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
   ]

   A		= [
       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
   ]
   A		= [
       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
   ]

   A		= [
       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
   ]
   A		= [
       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 					__a			(       self    ,__SCREAMING_SNAKE_CASE      ):
    SCREAMING_SNAKE_CASE  :  List[Any]								=		JukeboxTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE      )
    SCREAMING_SNAKE_CASE  :  Tuple								=		tokenizer(**self.metas      )['input_ids']
    return tokens






   @slow
   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Optional[Any]								=		JukeboxModel.from_pretrained(self.model_id    ,min_duration=0      ).eval()
    SCREAMING_SNAKE_CASE  :  Dict								=		self.prepare_inputs(self.model_id      )

    set_seed(0      )
    SCREAMING_SNAKE_CASE  :  Dict								=		[torch.zeros(1    ,0    ,dtype=torch.long      ).cpu() for _ in range(3      )]
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		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      )
    SCREAMING_SNAKE_CASE  :  Optional[Any]								=		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      )
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		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 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Any								=		JukeboxModel.from_pretrained(self.model_id    ,min_duration=0      ).eval()
    SCREAMING_SNAKE_CASE  :  Dict								=		[i.cuda() for i in self.prepare_inputs(self.model_id      )]

    set_seed(0      )
    model.priors[0].cuda()
    SCREAMING_SNAKE_CASE  :  Dict								=		[torch.zeros(1    ,0    ,dtype=torch.long      ).cuda() for _ in range(3      )]
    SCREAMING_SNAKE_CASE  :  Union[str, 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()
    SCREAMING_SNAKE_CASE  :  Tuple								=		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()
    SCREAMING_SNAKE_CASE  :  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 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Any								=		'ArthurZ/jukebox_prior_0'
    SCREAMING_SNAKE_CASE  :  Any								=		JukeboxPrior.from_pretrained(__SCREAMING_SNAKE_CASE    ,min_duration=0      ).eval().half().to('cuda'      )

    SCREAMING_SNAKE_CASE  :  Optional[Any]								=		self.prepare_inputs(__SCREAMING_SNAKE_CASE      )[0].cuda()
    SCREAMING_SNAKE_CASE  :  int								=		model.get_metadata(__SCREAMING_SNAKE_CASE    ,0    ,7680    ,0      )
    set_seed(0      )
    SCREAMING_SNAKE_CASE  :  List[Any]								=		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'''

from __future__ import annotations

import unittest

from transformers import CTRLConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers.models.ctrl.modeling_tf_ctrl import (
       TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
       TFCTRLForSequenceClassification,
       TFCTRLLMHeadModel,
       TFCTRLModel,
   )


class 		_a     (							SCREAMING_SNAKE_CASE  ):







   """simple docstring"""




   def __init__(       self    ,__SCREAMING_SNAKE_CASE    ,):
    SCREAMING_SNAKE_CASE  :  Dict								=		parent
    SCREAMING_SNAKE_CASE  :  int								=		13
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		7
    SCREAMING_SNAKE_CASE  :  List[Any]								=		True
    SCREAMING_SNAKE_CASE  :  Dict								=		True
    SCREAMING_SNAKE_CASE  :  Optional[int]								=		True
    SCREAMING_SNAKE_CASE  :  int								=		True
    SCREAMING_SNAKE_CASE  :  str								=		True
    SCREAMING_SNAKE_CASE  :  int								=		99
    SCREAMING_SNAKE_CASE  :  List[str]								=		32
    SCREAMING_SNAKE_CASE  :  List[str]								=		2
    SCREAMING_SNAKE_CASE  :  Optional[Any]								=		4
    SCREAMING_SNAKE_CASE  :  Tuple								=		37
    SCREAMING_SNAKE_CASE  :  str								=		'gelu'
    SCREAMING_SNAKE_CASE  :  Tuple								=		0.1
    SCREAMING_SNAKE_CASE  :  Any								=		0.1
    SCREAMING_SNAKE_CASE  :  List[str]								=		512
    SCREAMING_SNAKE_CASE  :  List[Any]								=		16
    SCREAMING_SNAKE_CASE  :  List[Any]								=		2
    SCREAMING_SNAKE_CASE  :  Any								=		0.02
    SCREAMING_SNAKE_CASE  :  List[str]								=		3
    SCREAMING_SNAKE_CASE  :  Any								=		4
    SCREAMING_SNAKE_CASE  :  List[Any]								=		None
    SCREAMING_SNAKE_CASE  :  Optional[int]								=		self.vocab_size - 1






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Any								=		ids_tensor([self.batch_size, self.seq_length]    ,self.vocab_size      )

    SCREAMING_SNAKE_CASE  :  Dict								=		None
    if self.use_input_mask:
     SCREAMING_SNAKE_CASE  :  List[str]								=		random_attention_mask([self.batch_size, self.seq_length]      )

    SCREAMING_SNAKE_CASE  :  List[Any]								=		None
    if self.use_token_type_ids:
     SCREAMING_SNAKE_CASE  :  str								=		ids_tensor([self.batch_size, self.seq_length]    ,self.type_vocab_size      )

    SCREAMING_SNAKE_CASE  :  int								=		None
    if self.use_mc_token_ids:
     SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		ids_tensor([self.batch_size, self.num_choices]    ,self.seq_length      )

    SCREAMING_SNAKE_CASE  :  Optional[int]								=		None
    SCREAMING_SNAKE_CASE  :  List[Any]								=		None
    SCREAMING_SNAKE_CASE  :  str								=		None
    if self.use_labels:
     SCREAMING_SNAKE_CASE  :  Any								=		ids_tensor([self.batch_size]    ,self.type_sequence_label_size      )
     SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		ids_tensor([self.batch_size, self.seq_length]    ,self.num_labels      )
     SCREAMING_SNAKE_CASE  :  str								=		ids_tensor([self.batch_size]    ,self.num_choices      )

    SCREAMING_SNAKE_CASE  :  Dict								=		CTRLConfig(
        vocab_size=self.vocab_size    ,n_embd=self.hidden_size    ,n_layer=self.num_hidden_layers    ,n_head=self.num_attention_heads    ,n_positions=self.max_position_embeddings    ,pad_token_id=self.pad_token_id    ,)

    SCREAMING_SNAKE_CASE  :  Tuple								=		ids_tensor([self.num_hidden_layers, self.num_attention_heads]    ,2      )

    return (
        config,
        input_ids,
        input_mask,
        head_mask,
        token_type_ids,
        mc_token_ids,
        sequence_labels,
        token_labels,
        choice_labels,
    )






   def 					__a			(       self    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,*__SCREAMING_SNAKE_CASE      ):
    SCREAMING_SNAKE_CASE  :  int								=		TFCTRLModel(config=__SCREAMING_SNAKE_CASE      )
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		model(__SCREAMING_SNAKE_CASE      )

    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		[input_ids, None, input_mask]  # None is the input for 'past'
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		model(__SCREAMING_SNAKE_CASE      )

    SCREAMING_SNAKE_CASE  :  Tuple								=		model(__SCREAMING_SNAKE_CASE      )

    self.parent.assertEqual(result.last_hidden_state.shape    ,(self.batch_size, self.seq_length, self.hidden_size)      )






   def 					__a			(       self    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,*__SCREAMING_SNAKE_CASE      ):
    SCREAMING_SNAKE_CASE  :  Optional[int]								=		TFCTRLLMHeadModel(config=__SCREAMING_SNAKE_CASE      )
    SCREAMING_SNAKE_CASE  :  Any								=		{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
    SCREAMING_SNAKE_CASE  :  Any								=		model(__SCREAMING_SNAKE_CASE      )
    self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.seq_length, self.vocab_size)      )






   def 					__a			(       self    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,*__SCREAMING_SNAKE_CASE      ):
    SCREAMING_SNAKE_CASE  :  Tuple								=		self.num_labels
    SCREAMING_SNAKE_CASE  :  Any								=		ids_tensor([self.batch_size]    ,self.type_sequence_label_size      )
    SCREAMING_SNAKE_CASE  :  Dict								=		{
        'input_ids': input_ids,
        'token_type_ids': token_type_ids,
        'labels': sequence_labels,
    }
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		TFCTRLForSequenceClassification(__SCREAMING_SNAKE_CASE      )
    SCREAMING_SNAKE_CASE  :  str								=		model(__SCREAMING_SNAKE_CASE      )
    self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.num_labels)      )






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  str								=		self.prepare_config_and_inputs()

    (
        (
        SCREAMING_SNAKE_CASE
    )				, (
        SCREAMING_SNAKE_CASE
    )				, (
        SCREAMING_SNAKE_CASE
    )				, (
        SCREAMING_SNAKE_CASE
    )				, (
        SCREAMING_SNAKE_CASE
    )				, (
        SCREAMING_SNAKE_CASE
    )				, (
        SCREAMING_SNAKE_CASE
    )				, (
        SCREAMING_SNAKE_CASE
    )				, (
        SCREAMING_SNAKE_CASE
    )				, 
    )  :  List[Any]								=		config_and_inputs

    SCREAMING_SNAKE_CASE  :  int								=		{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
    return config, inputs_dict



@require_tf
class 		_a     (							SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			,    unittest.TestCase  ):







   """simple docstring"""




   A		= (TFCTRLModel, TFCTRLLMHeadModel, TFCTRLForSequenceClassification) if is_tf_available() else ()
   A		= (TFCTRLLMHeadModel,) if is_tf_available() else ()
   A		= (
       {
           'feature-extraction': TFCTRLModel,
           'text-classification': TFCTRLForSequenceClassification,
           'text-generation': TFCTRLLMHeadModel,
           'zero-shot': TFCTRLForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   A		= False
   A		= False






   def 					__a			(       self    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE      ):
    if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
     # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
     # `CTRLConfig` was never used in pipeline tests, either because of a missing checkpoint or because a tiny
     # config could not be created.
     return True

    return False






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  List[Any]								=		TFCTRLModelTester(self      )
    SCREAMING_SNAKE_CASE  :  Dict								=		ConfigTester(self    ,config_class=__SCREAMING_SNAKE_CASE    ,n_embd=37      )






   def 					__a			(       self      ):
    self.config_tester.run_common_tests()






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Tuple								=		self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_ctrl_model(*__SCREAMING_SNAKE_CASE      )






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  str								=		self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_ctrl_lm_head(*__SCREAMING_SNAKE_CASE      )






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Optional[Any]								=		self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_ctrl_for_sequence_classification(*__SCREAMING_SNAKE_CASE      )






   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE  :  Tuple								=		self.model_tester.prepare_config_and_inputs_for_common()
    SCREAMING_SNAKE_CASE  :  Any								=		[TFCTRLLMHeadModel]
    SCREAMING_SNAKE_CASE  :  str								=		[TFCTRLForSequenceClassification]

    for model_class in self.all_model_classes:
     SCREAMING_SNAKE_CASE  :  List[str]								=		model_class(__SCREAMING_SNAKE_CASE      )
     model.build()  # may be needed for the get_bias() call below
     assert isinstance(model.get_input_embeddings()    ,tf.keras.layers.Layer      )

     if model_class in list_lm_models:
      SCREAMING_SNAKE_CASE  :  Dict								=		model.get_output_embeddings()
      assert isinstance(__SCREAMING_SNAKE_CASE    ,tf.keras.layers.Layer      )
      SCREAMING_SNAKE_CASE  :  Any								=		model.get_bias()
      assert isinstance(__SCREAMING_SNAKE_CASE    ,__SCREAMING_SNAKE_CASE      )
      for k, v in name.items():
       assert isinstance(__SCREAMING_SNAKE_CASE    ,tf.Variable      )
     elif model_class in list_other_models_with_output_ebd:
      SCREAMING_SNAKE_CASE  :  Any								=		model.get_output_embeddings()
      assert isinstance(__SCREAMING_SNAKE_CASE    ,tf.keras.layers.Layer      )
      SCREAMING_SNAKE_CASE  :  str								=		model.get_bias()
      assert name is None
     else:
      SCREAMING_SNAKE_CASE  :  List[Any]								=		model.get_output_embeddings()
      assert x is None
      SCREAMING_SNAKE_CASE  :  int								=		model.get_bias()
      assert name is None






   @slow
   def 					__a			(       self      ):
    for model_name in TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     SCREAMING_SNAKE_CASE  :  List[str]								=		TFCTRLModel.from_pretrained(__SCREAMING_SNAKE_CASE      )
     self.assertIsNotNone(__SCREAMING_SNAKE_CASE      )



@require_tf
class 		_a     (							unittest.TestCase  ):







   """simple docstring"""




   @slow
   def 					__a			(       self      ):
    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		TFCTRLLMHeadModel.from_pretrained('ctrl'      )
    SCREAMING_SNAKE_CASE  :  Any								=		tf.convert_to_tensor([[11859, 0, 1611, 8]]    ,dtype=tf.intaa      )  # Legal the president is
    SCREAMING_SNAKE_CASE  :  Optional[Any]								=		[
        11859,
        0,
        1611,
        8,
        5,
        150,
        26449,
        2,
        19,
        348,
        469,
        3,
        2595,
        48,
        20740,
        246533,
        246533,
        19,
        30,
        5,
    ]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

    SCREAMING_SNAKE_CASE  :  Union[str, Any]								=		model.generate(__SCREAMING_SNAKE_CASE    ,do_sample=__SCREAMING_SNAKE_CASE      )
    self.assertListEqual(output_ids[0].numpy().tolist()    ,__SCREAMING_SNAKE_CASE      )


