





"""simple docstring"""
from ..utils import DummyObject, requires_backends







class 	a				(					metaclass=a__   ):
 snake_case__    	=    ['''tensorflow_text''']



 def __init__(			self							,    *_snake_case							,    **_snake_case			):







  """simple docstring"""






  requires_backends(self							,    ['tensorflow_text']			)












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

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

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

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

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

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







  """simple docstring"""






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



 @slow
 def 					UpperCamelCase__   (			self			):







  """simple docstring"""






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

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

  set_seed(0			)
  lowerCAmelCase  							=      model._sample(_snake_case							,    _snake_case							,    [1]							,    sample_length=40 * model.priors[1].raw_to_tokens							,    save_results=_snake_case			)
  self.assertIn(zs[1][0].detach().cpu().tolist()							,    [self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]			)

  set_seed(0			)
  lowerCAmelCase  							=      model._sample(_snake_case							,    _snake_case							,    [2]							,    sample_length=40 * model.priors[2].raw_to_tokens							,    save_results=_snake_case			)
  self.assertIn(zs[2][0].detach().cpu().tolist()							,    [self.EXPECTED_OUTPUT_0, self.EXPECTED_OUTPUT_0_PT_2]			)



 @slow
 def 					UpperCamelCase__   (			self			):







  """simple docstring"""






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

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

  lowerCAmelCase  							=      model.priors[0]
  lowerCAmelCase  							=      0
  lowerCAmelCase  							=      top_prior.get_music_tokens_conds(_snake_case							,    start=_snake_case							,    end=start + top_prior.n_ctx			)
  lowerCAmelCase  							=      top_prior.get_metadata(labels[0].clone()							,    _snake_case							,    1_05_83_04							,    0			)

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

  lowerCAmelCase							,lowerCAmelCase							,lowerCAmelCase  							=      top_prior.get_cond(_snake_case							,    _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 					UpperCamelCase__   (			self			):







  """simple docstring"""






  lowerCAmelCase  							=      False

  lowerCAmelCase  							=      JukeboxModel.from_pretrained(self.model_id							,    min_duration=0			).eval()
  set_seed(0			)
  lowerCAmelCase  							=      torch.rand((1, 51_20, 1)			)
  lowerCAmelCase  							=      list(self.prepare_inputs()			)

  lowerCAmelCase  							=      [model.vqvae.encode(_snake_case							,    start_level=2							,    bs_chunks=waveform.shape[0]			)[0], None, None]
  lowerCAmelCase  							=      model._sample(
      _snake_case							,    _snake_case							,    sample_levels=[0]							,    save_results=_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			)			)

  lowerCAmelCase  							=      torch.cat((zs[0], torch.zeros(1							,    20_48 - zs[0].shape[-1]			))							,    dim=-1			).long()
  lowerCAmelCase  							=      [upper_a, model.vqvae.encode(_snake_case							,    start_level=1							,    bs_chunks=waveform.shape[0]			)[0], None]
  lowerCAmelCase  							=      model._sample(
      _snake_case							,    _snake_case							,    sample_levels=[1]							,    save_results=_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			)			)

  lowerCAmelCase  							=      torch.cat((zs[1], torch.zeros(1							,    20_48 - zs[1].shape[-1]			))							,    dim=-1			).long()
  lowerCAmelCase  							=      [upper_a, upper_a, model.vqvae.encode(_snake_case							,    start_level=0							,    bs_chunks=waveform.shape[0]			)[0]]
  lowerCAmelCase  							=      model._sample(
      _snake_case							,    _snake_case							,    sample_levels=[2]							,    save_results=_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 					UpperCamelCase__   (			self			):







  """simple docstring"""






  lowerCAmelCase  							=      JukeboxModel.from_pretrained(self.model_id							,    min_duration=0			).eval()
  set_seed(0			)
  lowerCAmelCase  							=      torch.rand((1, 51_20, 1)			)
  with torch.no_grad():
   lowerCAmelCase  							=      model.vqvae.encode(_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():
   lowerCAmelCase  							=      model.vqvae.decode(_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   ):
 snake_case__    	=    (JukeboxModel,) if is_torch_available() else ()
 snake_case__    	=    '''openai/jukebox-5b-lyrics'''
 snake_case__    	=    {
     '''artist''': '''Zac Brown Band''',
     '''genres''': '''Country''',
     '''lyrics''': '''I met a traveller from an antique land,
    Who said "Two vast and trunkless legs of stone
    Stand in the desert. . . . Near them, on the sand,
    Half sunk a shattered visage lies, whose frown,
    And wrinkled lip, and sneer of cold command,
    Tell that its sculptor well those passions read
    Which yet survive, stamped on these lifeless things,
    The hand that mocked them, and the heart that fed;
    And on the pedestal, these words appear:
    My name is Ozymandias, King of Kings;
    Look on my Works, ye Mighty, and despair!
    Nothing beside remains. Round the decay
    Of that colossal Wreck, boundless and bare
    The lone and level sands stretch far away
    ''',
 }

 # fmt: off
 snake_case__    	=    [
     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
 ]
 snake_case__    	=    [
     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
 ]

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

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

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

 snake_case__    	=    [
     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
 ]
 snake_case__    	=    [
     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 					UpperCamelCase__   (			self							,    _snake_case			):







  """simple docstring"""






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



 @slow
 def 					UpperCamelCase__   (			self			):







  """simple docstring"""






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

  set_seed(0			)
  lowerCAmelCase  							=      [torch.zeros(1							,    0							,    dtype=torch.long			).cpu() for _ in range(3			)]
  lowerCAmelCase  							=      model._sample(_snake_case							,    _snake_case							,    [0]							,    sample_length=60 * model.priors[0].raw_to_tokens							,    save_results=_snake_case			)
  self.assertIn(zs[0][0].detach().cpu().tolist()							,    [self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2]			)

  set_seed(0			)
  lowerCAmelCase  							=      model._sample(_snake_case							,    _snake_case							,    [1]							,    sample_length=60 * model.priors[1].raw_to_tokens							,    save_results=_snake_case			)
  self.assertIn(zs[1][0].detach().cpu().tolist()							,    [self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]			)

  set_seed(0			)
  lowerCAmelCase  							=      model._sample(_snake_case							,    _snake_case							,    [2]							,    sample_length=60 * model.priors[2].raw_to_tokens							,    save_results=_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 					UpperCamelCase__   (			self			):







  """simple docstring"""






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

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



 @slow
 def 					UpperCamelCase__   (			self			):







  """simple docstring"""






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

  lowerCAmelCase  							=      self.prepare_inputs(_snake_case			)[0].cuda()
  lowerCAmelCase  							=      model.get_metadata(_snake_case							,    0							,    76_80							,    0			)
  set_seed(0			)
  lowerCAmelCase  							=      model.sample(1							,    metadata=_snake_case							,    sample_tokens=60			)
  self.assertIn(outputs[0].cpu().tolist()							,    [self.EXPECTED_GPU_OUTPUTS_2, self.EXPECTED_GPU_OUTPUTS_2_PT_2]			)





