
'''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):
 UpperCAmelCase__     :     Optional[int]    					=       (JukeboxModel,) if is_torch_available() else ()
 UpperCAmelCase__     :     str    					=       'openai/jukebox-1b-lyrics'
 UpperCAmelCase__     :     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
 UpperCAmelCase__     :     List[str]    					=       [
     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
 ]

 UpperCAmelCase__     :     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
 ]

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

 UpperCAmelCase__     :     str    					=       [
     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
 ]
 UpperCAmelCase__     :     Optional[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
 ]

 UpperCAmelCase__     :     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]

 UpperCAmelCase__     :     int    					=       [
     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
 ]
 UpperCAmelCase__     :     Optional[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
 ]
 UpperCAmelCase__     :     int    					=       [
     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
 ]
 UpperCAmelCase__     :     Optional[int]    					=       [
     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
 ]
 UpperCAmelCase__     :     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
 ]
 UpperCAmelCase__     :     Dict    					=       [
     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
 ]
 UpperCAmelCase__     :     List[str]    					=       [
     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
 ]
 UpperCAmelCase__     :     Optional[Any]    					=       [
     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 						lowercase_    (      self	:Tuple      )					->      List[Any]:







       '''simple docstring'''



       __A													=      JukeboxTokenizer.from_pretrained(self.model_id      )
       __A													=      tokenizer(**self.metas      )['input_ids']
       return tokens






 @slow
 def 						lowercase_    (      self	:Tuple      )					->      int:







       '''simple docstring'''



       __A													=      JukeboxModel.from_pretrained(self.model_id						,					min_duration=0      ).eval()
       __A													=      self.prepare_inputs()

       set_seed(0      )
       __A													=      [torch.zeros(1						,					0						,					dtype=torch.long      ).cpu() for _ in range(3      )]
       __A													=      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      )
       __A													=      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      )
       __A													=      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 						lowercase_    (      self	:Any      )					->      str:







       '''simple docstring'''



       __A													=      False
       __A													=      JukeboxModel.from_pretrained(self.model_id						,					min_duration=0      ).eval()

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

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

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

       __A						,   __A						,   __A													=      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 						lowercase_    (      self	:Dict      )					->      Union[str, Any]:







       '''simple docstring'''



       __A													=      False

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

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

       __A													=      torch.cat((zs[0], torch.zeros(1						,					2_048 - zs[0].shape[-1]      ))						,					dim=-1      ).long()
       __A													=      [upper_a, model.vqvae.encode(_A						,					start_level=1						,					bs_chunks=waveform.shape[0]      )[0], None]
       __A													=      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      )      )

       __A													=      torch.cat((zs[1], torch.zeros(1						,					2_048 - zs[1].shape[-1]      ))						,					dim=-1      ).long()
       __A													=      [upper_a, upper_a, model.vqvae.encode(_A						,					start_level=0						,					bs_chunks=waveform.shape[0]      )[0]]
       __A													=      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 						lowercase_    (      self	:Union[str, Any]      )					->      str:







       '''simple docstring'''



       __A													=      JukeboxModel.from_pretrained(self.model_id						,					min_duration=0      ).eval()
       set_seed(0      )
       __A													=      torch.rand((1, 5_120, 1)      )
       with torch.no_grad():
             __A													=      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():
             __A													=      model.vqvae.decode(_A						,					start_level=2						,					bs_chunks=x.shape[0]      )
       torch.testing.assert_allclose(x[0, :40, 0]						,					torch.tensor(self.EXPECTED_VQVAE_DECODE      )						,					atol=1E-4						,					rtol=1E-4      )


@require_torch
class 			UpperCamelCase__	(     unittest.TestCase):
 UpperCAmelCase__     :     int    					=       (JukeboxModel,) if is_torch_available() else ()
 UpperCAmelCase__     :     Optional[int]    					=       'openai/jukebox-5b-lyrics'
 UpperCAmelCase__     :     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
 UpperCAmelCase__     :     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
 ]
 UpperCAmelCase__     :     Dict    					=       [
     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
 ]

 UpperCAmelCase__     :     Dict    					=       [
     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
 ]
 UpperCAmelCase__     :     Optional[int]    					=       [
     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
 ]

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

 UpperCAmelCase__     :     List[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,
     6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3
 ]
 UpperCAmelCase__     :     Optional[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
 ]

 UpperCAmelCase__     :     List[str]    					=       [
     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
 ]
 UpperCAmelCase__     :     List[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 						lowercase_    (      self	:Tuple						,					_A	:Union[str, Any]      )					->      Union[str, Any]:







       '''simple docstring'''



       __A													=      JukeboxTokenizer.from_pretrained(_A      )
       __A													=      tokenizer(**self.metas      )['input_ids']
       return tokens






 @slow
 def 						lowercase_    (      self	:int      )					->      str:







       '''simple docstring'''



       __A													=      JukeboxModel.from_pretrained(self.model_id						,					min_duration=0      ).eval()
       __A													=      self.prepare_inputs(self.model_id      )

       set_seed(0      )
       __A													=      [torch.zeros(1						,					0						,					dtype=torch.long      ).cpu() for _ in range(3      )]
       __A													=      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      )
       __A													=      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      )
       __A													=      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 						lowercase_    (      self	:Tuple      )					->      Tuple:







       '''simple docstring'''



       __A													=      JukeboxModel.from_pretrained(self.model_id						,					min_duration=0      ).eval()
       __A													=      [i.cuda() for i in self.prepare_inputs(self.model_id      )]

       set_seed(0      )
       model.priors[0].cuda()
       __A													=      [torch.zeros(1						,					0						,					dtype=torch.long      ).cuda() for _ in range(3      )]
       __A													=      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()
       __A													=      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()
       __A													=      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 						lowercase_    (      self	:Dict      )					->      Union[str, Any]:







       '''simple docstring'''



       __A													=      'ArthurZ/jukebox_prior_0'
       __A													=      JukeboxPrior.from_pretrained(_A						,					min_duration=0      ).eval().half().to('cuda'      )

       __A													=      self.prepare_inputs(_A      )[0].cuda()
       __A													=      model.get_metadata(_A						,					0						,					7_680						,					0      )
       set_seed(0      )
       __A													=      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 datasets

from .evaluate import (
    apply_no_ans_threshold,
    find_all_best_thresh,
    get_raw_scores,
    make_eval_dict,
    make_qid_to_has_ans,
    merge_eval,
)


__lowerCAmelCase									=				'''\
@inproceedings{Rajpurkar2016SQuAD10,
  title={SQuAD: 100, 000+ Questions for Machine Comprehension of Text},
  author={Pranav Rajpurkar and Jian Zhang and Konstantin Lopyrev and Percy Liang},
  booktitle={EMNLP},
  year={2016}
}
'''

__lowerCAmelCase									=				'''
This metric wrap the official scoring script for version 2 of the Stanford Question
Answering Dataset (SQuAD).

Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by
crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span,
from the corresponding reading passage, or the question might be unanswerable.

SQuAD2.0 combines the 100,000 questions in SQuAD1.1 with over 50,000 unanswerable questions
written adversarially by crowdworkers to look similar to answerable ones.
To do well on SQuAD2.0, systems must not only answer questions when possible, but also
determine when no answer is supported by the paragraph and abstain from answering.
'''

__lowerCAmelCase									=				'''
Computes SQuAD v2 scores (F1 and EM).
Args:
    predictions: List of triple for question-answers to score with the following elements:
        - the question-answer \'id\' field as given in the references (see below)
        - the text of the answer
        - the probability that the question has no answer
    references: List of question-answers dictionaries with the following key-values:
            - \'id\': id of the question-answer pair (see above),
            - \'answers\': a list of Dict {\'text\': text of the answer as a string}
    no_answer_threshold: float
        Probability threshold to decide that a question has no answer.
Returns:
    \'exact\': Exact match (the normalized answer exactly match the gold answer)
    \'f1\': The F-score of predicted tokens versus the gold answer
    \'total\': Number of score considered
    \'HasAns_exact\': Exact match (the normalized answer exactly match the gold answer)
    \'HasAns_f1\': The F-score of predicted tokens versus the gold answer
    \'HasAns_total\': Number of score considered
    \'NoAns_exact\': Exact match (the normalized answer exactly match the gold answer)
    \'NoAns_f1\': The F-score of predicted tokens versus the gold answer
    \'NoAns_total\': Number of score considered
    \'best_exact\': Best exact match (with varying threshold)
    \'best_exact_thresh\': No-answer probability threshold associated to the best exact match
    \'best_f1\': Best F1 (with varying threshold)
    \'best_f1_thresh\': No-answer probability threshold associated to the best F1
Examples:

    >>> predictions = [{\'prediction_text\': \'1976\', \'id\': \'56e10a3be3433e1400422b22\', \'no_answer_probability\': 0.}]
    >>> references = [{\'answers\': {\'answer_start\': [97], \'text\': [\'1976\']}, \'id\': \'56e10a3be3433e1400422b22\'}]
    >>> squad_v2_metric = datasets.load_metric("squad_v2")
    >>> results = squad_v2_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'exact\': 100.0, \'f1\': 100.0, \'total\': 1, \'HasAns_exact\': 100.0, \'HasAns_f1\': 100.0, \'HasAns_total\': 1, \'best_exact\': 100.0, \'best_exact_thresh\': 0.0, \'best_f1\': 100.0, \'best_f1_thresh\': 0.0}
'''





@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION			,			_KWARGS_DESCRIPTION )
class     __a     (       datasets.Metric ):
 def 						SCREAMING_SNAKE_CASE__  (  self					)   ->							Dict:




        '''simple docstring'''
        return datasets.MetricInfo(
            description=_DESCRIPTION							,					citation=_CITATION							,					inputs_description=_KWARGS_DESCRIPTION							,					features=datasets.Features(
                {
                    'predictions': {
                        'id': datasets.Value('string'					),
                        'prediction_text': datasets.Value('string'					),
                        'no_answer_probability': datasets.Value('float32'					),
                    },
                    'references': {
                        'id': datasets.Value('string'					),
                        'answers': datasets.features.Sequence(
                            {'text': datasets.Value('string'					), 'answer_start': datasets.Value('int32'					)}					),
                    },
                }					)							,					codebase_urls=['https://rajpurkar.github.io/SQuAD-explorer/']							,					reference_urls=['https://rajpurkar.github.io/SQuAD-explorer/']							,					)
 def 						SCREAMING_SNAKE_CASE__  (  self							,					lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__=1.0					)   ->							str:




        '''simple docstring'''
        lowercase__:  int         =  {p['id']: p['no_answer_probability'] for p in predictions}
        lowercase__:  int         =  [{'paragraphs': [{'qas': references}]}]
        lowercase__:  Optional[int]         =  {p['id']: p['prediction_text'] for p in predictions}

        lowercase__:  Any         =  make_qid_to_has_ans(lowerCAmelCase__					)  # maps qid to True/False
        lowercase__:  Optional[Any]         =  [k for k, v in qid_to_has_ans.items() if v]
        lowercase__:  Dict         =  [k for k, v in qid_to_has_ans.items() if not v]

        lowercase__     ,		lowercase__:  Dict         =  get_raw_scores(lowerCAmelCase__							,					lowerCAmelCase__					)
        lowercase__:  Optional[int]         =  apply_no_ans_threshold(lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__					)
        lowercase__:  List[Any]         =  apply_no_ans_threshold(lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__					)
        lowercase__:  Dict         =  make_eval_dict(lowerCAmelCase__							,					lowerCAmelCase__					)

        if has_ans_qids:
               lowercase__:  List[Any]         =  make_eval_dict(lowerCAmelCase__							,					lowerCAmelCase__							,					qid_list=lowerCAmelCase__					)
               merge_eval(lowerCAmelCase__							,					lowerCAmelCase__							,					'HasAns'					)
        if no_ans_qids:
               lowercase__:  int         =  make_eval_dict(lowerCAmelCase__							,					lowerCAmelCase__							,					qid_list=lowerCAmelCase__					)
               merge_eval(lowerCAmelCase__							,					lowerCAmelCase__							,					'NoAns'					)
        find_all_best_thresh(lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__					)
        return dict(lowerCAmelCase__					)


