




import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import OPTConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
       import os

       # The slow tests are often failing with OOM error on GPU
       # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
       # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
       UpperCAmelCase_   							=      'platform'

       import jax
       import jax.numpy as jnp

       from transformers import FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer




def    lowerCamelCase__			(		A__			:   Dict    ,     A__			:   Optional[Any]    ,     A__			:   Any=None    ,     A__			:   Any=None      ):





       '''simple docstring'''



       if attention_mask is None:
              __lowerCamelCase  =		np.where(input_ids != config.pad_token_id    ,     1    ,     0      )
       return {
           "input_ids": input_ids,
           "attention_mask": attention_mask,
       }





@require_flax
class        lowerCamelCase__:
      def __init__(							self:					int       ,			UpperCamelCase_:					int       ,			UpperCamelCase_:					Dict=13       ,			UpperCamelCase_:					List[Any]=7       ,			UpperCamelCase_:					int=True       ,			UpperCamelCase_:					Optional[Any]=False       ,			UpperCamelCase_:					Tuple=99       ,			UpperCamelCase_:					int=16       ,			UpperCamelCase_:					Dict=2       ,			UpperCamelCase_:					Any=4       ,			UpperCamelCase_:					str=4       ,			UpperCamelCase_:					Any="gelu"       ,			UpperCamelCase_:					Tuple=0.1       ,			UpperCamelCase_:					int=0.1       ,			UpperCamelCase_:					Optional[int]=20       ,			UpperCamelCase_:					Optional[Any]=2       ,			UpperCamelCase_:					Optional[Any]=1       ,			UpperCamelCase_:					Optional[int]=0       ,			UpperCamelCase_:					Any=16       ,			UpperCamelCase_:					Union[str, Any]=16       ,			UpperCamelCase_:					str=0.02       ,			):
             __lowerCamelCase  =		parent
             __lowerCamelCase  =		batch_size
             __lowerCamelCase  =		seq_length
             __lowerCamelCase  =		is_training
             __lowerCamelCase  =		use_labels
             __lowerCamelCase  =		vocab_size
             __lowerCamelCase  =		hidden_size
             __lowerCamelCase  =		num_hidden_layers
             __lowerCamelCase  =		num_attention_heads
             __lowerCamelCase  =		intermediate_size
             __lowerCamelCase  =		hidden_act
             __lowerCamelCase  =		hidden_dropout_prob
             __lowerCamelCase  =		attention_probs_dropout_prob
             __lowerCamelCase  =		max_position_embeddings
             __lowerCamelCase  =		eos_token_id
             __lowerCamelCase  =		pad_token_id
             __lowerCamelCase  =		bos_token_id
             __lowerCamelCase  =		embed_dim
             __lowerCamelCase  =		word_embed_proj_dim
             __lowerCamelCase  =		initializer_range
             __lowerCamelCase  =		False
      def      lowerCAmelCase__		(							self:					List[Any]  ):
             __lowerCamelCase  =		np.clip(ids_tensor([self.batch_size, self.seq_length - 1]       ,			self.vocab_size  )       ,			3       ,			self.vocab_size  )
             __lowerCamelCase  =		np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)       ,			dtype=np.intaa  ))       ,			-1  )

             __lowerCamelCase  =		OPTConfig(
                 vocab_size=self.vocab_size       ,			hidden_size=self.hidden_size       ,			num_hidden_layers=self.num_hidden_layers       ,			num_attention_heads=self.num_attention_heads       ,			ffn_dim=self.intermediate_size       ,			dropout=self.hidden_dropout_prob       ,			attention_dropout=self.attention_probs_dropout_prob       ,			max_position_embeddings=self.max_position_embeddings       ,			eos_token_id=self.eos_token_id       ,			bos_token_id=self.bos_token_id       ,			pad_token_id=self.pad_token_id       ,			embed_dim=self.embed_dim       ,			is_encoder_decoder=UpperCamelCase_       ,			word_embed_proj_dim=self.word_embed_proj_dim       ,			initializer_range=self.initializer_range       ,			use_cache=UpperCamelCase_       ,			)
             __lowerCamelCase  =		prepare_opt_inputs_dict(UpperCamelCase_       ,			UpperCamelCase_  )
             return config, inputs_dict
      def      lowerCAmelCase__		(							self:					Union[str, Any]  ):
             __lowerCamelCase,  __lowerCamelCase  =		self.prepare_config_and_inputs()
             return config, inputs_dict
      def      lowerCAmelCase__		(							self:					str       ,			UpperCamelCase_:					Optional[int]       ,			UpperCamelCase_:					Union[str, Any]       ,			UpperCamelCase_:					Union[str, Any]  ):
             __lowerCamelCase  =		20
             __lowerCamelCase  =		model_class_name(UpperCamelCase_  )

             __lowerCamelCase  =		inputs_dict["""input_ids"""]
             __lowerCamelCase  =		inputs_dict["""attention_mask"""]

             __lowerCamelCase  =		model.init_cache(input_ids.shape[0]       ,			UpperCamelCase_  )
             __lowerCamelCase  =		jnp.ones((input_ids.shape[0], max_length)       ,			dtype="""i4"""  )

             __lowerCamelCase  =		jnp.broadcast_to(
                 jnp.arange(input_ids.shape[-1] - 1  )[None, :]       ,			(input_ids.shape[0], input_ids.shape[-1] - 1)       ,			)
             __lowerCamelCase  =		model(
                 input_ids[:, :-1]       ,			attention_mask=UpperCamelCase_       ,			past_key_values=UpperCamelCase_       ,			position_ids=UpperCamelCase_       ,			)

             __lowerCamelCase  =		jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,			dtype="""i4"""  )
             __lowerCamelCase  =		model(
                 input_ids[:, -1:]       ,			attention_mask=UpperCamelCase_       ,			past_key_values=outputs_cache.past_key_values       ,			position_ids=UpperCamelCase_       ,			)

             __lowerCamelCase  =		model(UpperCamelCase_  )

             __lowerCamelCase  =		np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])  )  )
             self.parent.assertTrue(diff < 1E-3       ,			msg=F'Max diff is {diff}'  )




      def      lowerCAmelCase__		(							self:					Optional[Any]       ,			UpperCamelCase_:					Any       ,			UpperCamelCase_:					Union[str, Any]       ,			UpperCamelCase_:					Optional[int]  ):
             __lowerCamelCase  =		20
             __lowerCamelCase  =		model_class_name(UpperCamelCase_  )

             __lowerCamelCase,  __lowerCamelCase  =		(
                 inputs_dict["""input_ids"""],
                 inputs_dict["""attention_mask"""],
             )

             __lowerCamelCase  =		jnp.concatenate(
                 [
                     attention_mask,
                     jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1])  ),
                 ]       ,			axis=-1       ,			)

             __lowerCamelCase  =		model.init_cache(input_ids.shape[0]       ,			UpperCamelCase_  )
             __lowerCamelCase  =		jnp.broadcast_to(
                 jnp.arange(input_ids.shape[-1] - 1  )[None, :]       ,			(input_ids.shape[0], input_ids.shape[-1] - 1)       ,			)

             __lowerCamelCase  =		model(
                 input_ids[:, :-1]       ,			attention_mask=UpperCamelCase_       ,			past_key_values=UpperCamelCase_       ,			position_ids=UpperCamelCase_       ,			)
             __lowerCamelCase  =		jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,			dtype="""i4"""  )
             __lowerCamelCase  =		model(
                 input_ids[:, -1:]       ,			past_key_values=outputs_cache.past_key_values       ,			attention_mask=UpperCamelCase_       ,			position_ids=UpperCamelCase_       ,			)

             __lowerCamelCase  =		model(UpperCamelCase_       ,			attention_mask=UpperCamelCase_  )

             __lowerCamelCase  =		np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])  )  )
             self.parent.assertTrue(diff < 1E-3       ,			msg=F'Max diff is {diff}'  )





@require_flax
class        lowerCamelCase__(   __lowerCamelCase					,	unittest.TestCase					,	__lowerCamelCase):
      UpperCAmelCase__ :      Dict       						=  (FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
      UpperCAmelCase__ :      List[Any]       						=  () if is_flax_available() else ()
      def      lowerCAmelCase__		(							self:					Optional[Any]  ):
             __lowerCamelCase  =		FlaxOPTModelTester(self  )
      def      lowerCAmelCase__		(							self:					Union[str, Any]  ):
             __lowerCamelCase,  __lowerCamelCase  =		self.model_tester.prepare_config_and_inputs()
             for model_class in self.all_model_classes:
                    self.model_tester.check_use_cache_forward(UpperCamelCase_       ,			UpperCamelCase_       ,			UpperCamelCase_  )
      def      lowerCAmelCase__		(							self:					str  ):
             __lowerCamelCase,  __lowerCamelCase  =		self.model_tester.prepare_config_and_inputs()
             for model_class in self.all_model_classes:
                    self.model_tester.check_use_cache_forward_with_attn_mask(UpperCamelCase_       ,			UpperCamelCase_       ,			UpperCamelCase_  )




      @slow
      def      lowerCAmelCase__		(							self:					List[str]  ):
             for model_class_name in self.all_model_classes:
                    __lowerCamelCase  =		model_class_name.from_pretrained("""facebook/opt-125m"""  )
                    __lowerCamelCase  =		np.ones((1, 1)  ) * model.config.eos_token_id
                    __lowerCamelCase  =		model(UpperCamelCase_  )
                    self.assertIsNotNone(UpperCamelCase_  )





@require_sentencepiece
@require_flax
class        lowerCamelCase__(   unittest.TestCase):
      @slow
      def      lowerCAmelCase__		(							self:					int  ):
             __lowerCamelCase  =		FlaxOPTModel.from_pretrained("""facebook/opt-350m"""  )
             __lowerCamelCase  =		jnp.array([[0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 4_60_78, 15_88, 2]]  )
             __lowerCamelCase  =		model(input_ids=UpperCamelCase_  ).last_hidden_state
             __lowerCamelCase  =		(1, 11, 5_12)
             self.assertEqual(output.shape       ,			UpperCamelCase_  )
             __lowerCamelCase  =		jnp.array(
                 [[-0.2867, -1.9256, -0.3062], [-1.2711, -0.1337, -0.1897], [0.4109, 0.1187, -1.3142]]  )
             self.assertTrue(jnp.allclose(output[:, :3, :3]       ,			UpperCamelCase_       ,			atol=4E-2  )  )





@require_flax
@slow
class        lowerCamelCase__(   unittest.TestCase):
      def      lowerCAmelCase__		(							self:					str  ):
             super().setUp()
             __lowerCamelCase  =		"""facebook/opt-350m"""




      def      lowerCAmelCase__		(							self:					List[Any]  ):
             __lowerCamelCase  =		FlaxOPTForCausalLM.from_pretrained(self.path_model  )
             __lowerCamelCase  =		GPTaTokenizer.from_pretrained(self.path_model  )

             __lowerCamelCase  =		[
                 """Today is a beautiful day and I want to""",
                 """In the city of""",
                 """Paris is the capital of France and""",
                 """Computers and mobile phones have taken""",
             ]
             # verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
             __lowerCamelCase  =		tokenizer(UpperCamelCase_       ,			return_tensors="""jax"""       ,			padding=UpperCamelCase_       ,			add_special_tokens=UpperCamelCase_  )
             __lowerCamelCase  =		model(inputs.input_ids       ,			attention_mask=inputs.attention_mask  )[0].mean(axis=-1  )
             __lowerCamelCase  =		jnp.array(
                 [
                     [1.3851, -13.8923, -10.5229, -10.7533, -0.2309, -10.2384, -0.5365, -9.0947, -5.1670],
                     [-4.7073, -10.6276, -3.9415, -21.5242, -0.2822, -0.2822, -0.2822, -0.2822, -0.2822],
                     [0.6247, -3.4229, -8.9179, -1.4297, -14.1650, 1.4146, -9.0218, -0.2703, -0.2703],
                     [6.4783, -1.9913, -10.7926, -2.3336, 1.5092, -0.9974, -6.8213, 1.3477, 1.3477],
                 ]  )
             self.assertTrue(jnp.allclose(UpperCamelCase_       ,			UpperCamelCase_       ,			atol=4E-2  )  )

             __lowerCamelCase  =		jax.jit(UpperCamelCase_  )
             __lowerCamelCase  =		model(inputs.input_ids       ,			attention_mask=inputs.attention_mask  )[0].mean(axis=-1  )
             self.assertTrue(jnp.allclose(UpperCamelCase_       ,			UpperCamelCase_       ,			atol=4E-2  )  )





@require_flax
@slow
class        lowerCamelCase__(   unittest.TestCase):
      @property
      def      lowerCAmelCase__		(							self:					str  ):
             return [
                 "Today is a beautiful day and I want",
                 "In the city of",
                 "Paris is the capital of France and",
                 "Computers and mobile phones have taken",
             ]
      def      lowerCAmelCase__		(							self:					List[str]  ):
             __lowerCamelCase  =		"""facebook/opt-125m"""

             __lowerCamelCase  =		[
                 """Today is a beautiful day and I want to""",
                 """In the city of New York, the city""",
                 """Paris is the capital of France and the capital""",
                 """Computers and mobile phones have taken over the""",
             ]

             __lowerCamelCase  =		[]

             __lowerCamelCase  =		FlaxOPTForCausalLM.from_pretrained(UpperCamelCase_  )
             __lowerCamelCase  =		GPTaTokenizer.from_pretrained(UpperCamelCase_  )

             for prompt in self.prompts:
                    __lowerCamelCase  =		tokenizer(UpperCamelCase_       ,			return_tensors="""jax"""  ).input_ids

                    __lowerCamelCase  =		model.generate(UpperCamelCase_       ,			max_length=10  )
                    __lowerCamelCase  =		generated_ids[0]

                    __lowerCamelCase  =		tokenizer.batch_decode(UpperCamelCase_       ,			skip_special_tokens=UpperCamelCase_  )
                    predicted_outputs += generated_string

             self.assertListEqual(UpperCamelCase_       ,			UpperCamelCase_  )
      def      lowerCAmelCase__		(							self:					List[str]  ):
             __lowerCamelCase  =		"""facebook/opt-350m"""

             __lowerCamelCase  =		[
                 """Today is a beautiful day and I want to""",
                 """In the city of San Francisco, the city""",
                 """Paris is the capital of France and the capital""",
                 """Computers and mobile phones have taken over the""",
             ]

             __lowerCamelCase  =		[]
             __lowerCamelCase  =		FlaxOPTForCausalLM.from_pretrained(UpperCamelCase_  )
             __lowerCamelCase  =		GPTaTokenizer.from_pretrained(UpperCamelCase_  )

             for prompt in self.prompts:
                    __lowerCamelCase  =		tokenizer(UpperCamelCase_       ,			return_tensors="""jax"""  ).input_ids

                    __lowerCamelCase  =		model.generate(UpperCamelCase_       ,			max_length=10  )
                    __lowerCamelCase  =		generated_ids[0]

                    __lowerCamelCase  =		tokenizer.batch_decode(UpperCamelCase_       ,			skip_special_tokens=UpperCamelCase_  )
                    predicted_outputs += generated_string

             self.assertListEqual(UpperCamelCase_       ,			UpperCamelCase_  )
      def      lowerCAmelCase__		(							self:					int  ):
             __lowerCamelCase  =		"""facebook/opt-125m"""
             __lowerCamelCase  =		[
                 """Today is a beautiful day and I want to thank""",
                 """In the city of Rome Canaver Canaver Canaver Canaver""",
             ]
             __lowerCamelCase  =		FlaxOPTForCausalLM.from_pretrained(UpperCamelCase_  )
             __lowerCamelCase  =		GPTaTokenizer.from_pretrained(UpperCamelCase_  )
             __lowerCamelCase  =		tokenizer(
                 [
                     """Today is a beautiful day and I want to""",
                     """In the city of""",
                 ]       ,			return_tensors="""jax"""       ,			padding=UpperCamelCase_       ,			)

             __lowerCamelCase  =		jax.jit(model.generate  )

             __lowerCamelCase  =		jit_generate(inputs["""input_ids"""]       ,			attention_mask=inputs["""attention_mask"""]  ).sequences

             __lowerCamelCase  =		tokenizer.batch_decode(UpperCamelCase_       ,			skip_special_tokens=UpperCamelCase_  )

             self.assertIsNotNone(UpperCamelCase_       ,			UpperCamelCase_  )




      def      lowerCAmelCase__		(							self:					Any  ):
             __lowerCamelCase  =		"""facebook/opt-350m"""

             __lowerCamelCase  =		GPTaTokenizer.from_pretrained(UpperCamelCase_  )
             __lowerCamelCase  =		FlaxOPTForCausalLM.from_pretrained(UpperCamelCase_  )

             __lowerCamelCase  =		"""left"""

             # use different length sentences to test batching
             __lowerCamelCase  =		[
                 """Hello, my dog is a little""",
                 """Today, I""",
             ]

             __lowerCamelCase  =		tokenizer(UpperCamelCase_       ,			return_tensors="""jax"""       ,			padding=UpperCamelCase_  )
             __lowerCamelCase  =		inputs["""input_ids"""]

             __lowerCamelCase  =		model.generate(input_ids=UpperCamelCase_       ,			attention_mask=inputs["""attention_mask"""]       ,			trace=UpperCamelCase_  )

             __lowerCamelCase  =		tokenizer(sentences[0]       ,			return_tensors="""jax"""  ).input_ids
             __lowerCamelCase  =		model.generate(input_ids=UpperCamelCase_  )

             __lowerCamelCase  =		inputs_non_padded.shape[-1] - inputs["""attention_mask"""][-1].sum()
             __lowerCamelCase  =		tokenizer(sentences[1]       ,			return_tensors="""jax"""  ).input_ids
             __lowerCamelCase  =		model.generate(input_ids=UpperCamelCase_       ,			max_length=model.config.max_length - num_paddings  )

             __lowerCamelCase  =		tokenizer.batch_decode(outputs[0]       ,			skip_special_tokens=UpperCamelCase_  )
             __lowerCamelCase  =		tokenizer.decode(output_non_padded[0][0]       ,			skip_special_tokens=UpperCamelCase_  )
             __lowerCamelCase  =		tokenizer.decode(output_padded[0][0]       ,			skip_special_tokens=UpperCamelCase_  )

             __lowerCamelCase  =		[
                 """Hello, my dog is a little bit of a dork.\nI'm a little bit""",
                 """Today, I was in the middle of a conversation with a friend about the""",
             ]
             self.assertListEqual(UpperCamelCase_       ,			UpperCamelCase_  )
             self.assertListEqual(UpperCamelCase_       ,			[non_padded_sentence, padded_sentence]  )







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        lowerCamelCase__(   unittest.TestCase):
      UpperCAmelCase__ :      Optional[Any]       						=  (JukeboxModel,) if is_torch_available() else ()
      UpperCAmelCase__ :      Union[str, Any]       						=  'openai/jukebox-1b-lyrics'
      UpperCAmelCase__ :      Optional[int]       						=  {
          '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__ :      Dict       						=  [
          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
      ]

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

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

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

      UpperCAmelCase__ :      List[Any]       						=  [105_8304, 0, 78_6432, 7169, 507, 76, 27, 40, 30, 76]

      UpperCAmelCase__ :      Union[str, Any]       						=  [
          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
      ]
      UpperCAmelCase__ :      List[str]       						=  [
          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
      ]
      UpperCAmelCase__ :      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
      ]
      UpperCAmelCase__ :      Union[str, Any]       						=  [
          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
      ]
      UpperCAmelCase__ :      Optional[int]       						=  [
          -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
      ]
      UpperCAmelCase__ :      Union[str, Any]       						=  [
          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
      ]
      UpperCAmelCase__ :      str       						=  [
          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
      ]
      UpperCAmelCase__ :      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      lowerCAmelCase__		(							self:					Dict  ):
             __lowerCamelCase  =		JukeboxTokenizer.from_pretrained(self.model_id  )
             __lowerCamelCase  =		tokenizer(**self.metas  )["""input_ids"""]
             return tokens
      @slow
      def      lowerCAmelCase__		(							self:					Union[str, Any]  ):
             __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(UpperCamelCase_       ,			UpperCamelCase_       ,			[0]       ,			sample_length=40 * model.priors[0].raw_to_tokens       ,			save_results=UpperCamelCase_  )
             self.assertIn(zs[0][0].detach().cpu().tolist()       ,			[self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2]  )

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

             set_seed(0  )
             __lowerCamelCase  =		model._sample(UpperCamelCase_       ,			UpperCamelCase_       ,			[2]       ,			sample_length=40 * model.priors[2].raw_to_tokens       ,			save_results=UpperCamelCase_  )
             self.assertIn(zs[2][0].detach().cpu().tolist()       ,			[self.EXPECTED_OUTPUT_0, self.EXPECTED_OUTPUT_0_PT_2]  )
      @slow
      def      lowerCAmelCase__		(							self:					Union[str, Any]  ):
             __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(UpperCamelCase_       ,			start=UpperCamelCase_       ,			end=start + top_prior.n_ctx  )
             __lowerCamelCase  =		top_prior.get_metadata(labels[0].clone()       ,			UpperCamelCase_       ,			1_05_83_04       ,			0  )

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

             __lowerCamelCase,  __lowerCamelCase,  __lowerCamelCase  =		top_prior.get_cond(UpperCamelCase_       ,			UpperCamelCase_  )
             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      lowerCAmelCase__		(							self:					List[Any]  ):
             __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(UpperCamelCase_       ,			start_level=2       ,			bs_chunks=waveform.shape[0]  )[0], None, None]
             __lowerCamelCase  =		model._sample(
                 UpperCamelCase_       ,			UpperCamelCase_       ,			sample_levels=[0]       ,			save_results=UpperCamelCase_       ,			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(UpperCamelCase_       ,			start_level=1       ,			bs_chunks=waveform.shape[0]  )[0], None]
             __lowerCamelCase  =		model._sample(
                 UpperCamelCase_       ,			UpperCamelCase_       ,			sample_levels=[1]       ,			save_results=UpperCamelCase_       ,			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(UpperCamelCase_       ,			start_level=0       ,			bs_chunks=waveform.shape[0]  )[0]]
             __lowerCamelCase  =		model._sample(
                 UpperCamelCase_       ,			UpperCamelCase_       ,			sample_levels=[2]       ,			save_results=UpperCamelCase_       ,			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      lowerCAmelCase__		(							self:					List[Any]  ):
             __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(UpperCamelCase_       ,			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(UpperCamelCase_       ,			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        lowerCamelCase__(   unittest.TestCase):
      UpperCAmelCase__ :      List[Any]       						=  (JukeboxModel,) if is_torch_available() else ()
      UpperCAmelCase__ :      Optional[int]       						=  'openai/jukebox-5b-lyrics'
      UpperCAmelCase__ :      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
      UpperCAmelCase__ :      Optional[int]       						=  [
          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
      ]
      UpperCAmelCase__ :      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, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653
      ]

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

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

      UpperCAmelCase__ :      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,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653
      ]
      UpperCAmelCase__ :      Dict       						=  [
          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
      ]

      UpperCAmelCase__ :      Optional[int]       						=  [
          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
      ]
      UpperCAmelCase__ :      str       						=  [
          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      lowerCAmelCase__		(							self:					Optional[Any]       ,			UpperCamelCase_:					Optional[int]  ):
             __lowerCamelCase  =		JukeboxTokenizer.from_pretrained(UpperCamelCase_  )
             __lowerCamelCase  =		tokenizer(**self.metas  )["""input_ids"""]
             return tokens
      @slow
      def      lowerCAmelCase__		(							self:					Tuple  ):
             __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(UpperCamelCase_       ,			UpperCamelCase_       ,			[0]       ,			sample_length=60 * model.priors[0].raw_to_tokens       ,			save_results=UpperCamelCase_  )
             self.assertIn(zs[0][0].detach().cpu().tolist()       ,			[self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2]  )

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

             set_seed(0  )
             __lowerCamelCase  =		model._sample(UpperCamelCase_       ,			UpperCamelCase_       ,			[2]       ,			sample_length=60 * model.priors[2].raw_to_tokens       ,			save_results=UpperCamelCase_  )
             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      lowerCAmelCase__		(							self:					List[Any]  ):
             __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(UpperCamelCase_       ,			UpperCamelCase_       ,			[0]       ,			sample_length=60 * model.priors[0].raw_to_tokens       ,			save_results=UpperCamelCase_  )
             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(UpperCamelCase_       ,			UpperCamelCase_       ,			[1]       ,			sample_length=60 * model.priors[1].raw_to_tokens       ,			save_results=UpperCamelCase_  )
             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(UpperCamelCase_       ,			UpperCamelCase_       ,			[2]       ,			sample_length=60 * model.priors[2].raw_to_tokens       ,			save_results=UpperCamelCase_  )
             torch.testing.assert_allclose(zs[2][0].cpu()       ,			torch.tensor(self.EXPECTED_GPU_OUTPUTS_0  )  )




      @slow
      def      lowerCAmelCase__		(							self:					Dict  ):
             __lowerCamelCase  =		"""ArthurZ/jukebox_prior_0"""
             __lowerCamelCase  =		JukeboxPrior.from_pretrained(UpperCamelCase_       ,			min_duration=0  ).eval().half().to("""cuda"""  )

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

