


import unittest

from transformers import JukeboxTokenizer
from transformers.testing_utils import require_torch







class        UpperCamelCase__   (	unittest.TestCase   ):
					__SCREAMING_SNAKE_CASE       :  int						=   JukeboxTokenizer
					__SCREAMING_SNAKE_CASE       :  str						=   {
					    '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        ',
					}

					@require_torch
					def   UpperCAmelCase__       (     self       :   Any       ):






										'''simple docstring'''






										import torch

										lowercase_  							=			JukeboxTokenizer.from_pretrained("""openai/jukebox-1b-lyrics"""       )
										lowercase_  							=			tokenizer(**self.metas       )["""input_ids"""]
										# fmt: off
										lowercase_  							=			[
										    torch.tensor([[
										        0, 0, 0, 7_169, 507, 9, 76, 39, 31, 46, 76, 27,
										        76, 46, 44, 27, 48, 31, 38, 38, 31, 44, 76, 32,
										        44, 41, 39, 76, 27, 40, 76, 27, 40, 46, 35, 43,
										        47, 31, 76, 38, 27, 40, 30, 64, 78, 76, 76, 76,
										        76, 76, 76, 76, 76, 23, 34, 41, 76, 45, 27, 35,
										        30, 76, 71, 20, 49, 41, 76, 48, 27, 45, 46, 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, 76, 76, 76, 76, 76, 76, 76,
										        19, 46, 27, 40, 30, 76, 35, 40, 76, 46, 34, 31,
										        76, 30, 31, 45, 31, 44, 46, 63, 76, 63, 76, 63,
										        76, 63, 76, 14, 31, 27, 44, 76, 46, 34, 31, 39,
										        64, 76, 41, 40, 76, 46, 34, 31, 76, 45, 27, 40,
										        30, 64, 78, 76, 76, 76, 76, 76, 76, 76, 76, 8,
										        27, 38, 32, 76, 45, 47, 40, 37, 76, 27, 76, 45,
										        34, 27, 46, 46, 31, 44, 31, 30, 76, 48, 35, 45,
										        27, 33, 31, 76, 38, 35, 31, 45, 64, 76, 49, 34,
										        41, 45, 31, 76, 32, 44, 41, 49, 40, 64, 78, 76,
										        76, 76, 76, 76, 76, 76, 76, 1, 40, 30, 76, 49,
										        44, 35, 40, 37, 38, 31, 30, 76, 38, 35, 42, 64,
										        76, 27, 40, 30, 76, 45, 40, 31, 31, 44, 76, 41,
										        32, 76, 29, 41, 38, 30, 76, 29, 41, 39, 39, 27,
										        40, 30, 64, 78, 76, 76, 76, 76, 76, 76, 76, 76,
										        20, 31, 38, 38, 76, 46, 34, 27, 46, 76, 35, 46,
										        45, 76, 45, 29, 47, 38, 42, 46, 41, 44, 76, 49,
										        31, 38, 38, 76, 46, 34, 41, 45, 31, 76, 42, 27,
										        45, 45, 35, 41, 40, 45, 76, 44, 31, 27, 30, 78,
										        76, 76, 76, 76, 76, 76, 76, 76, 23, 34, 35, 29,
										        34, 76, 51, 31, 46, 76, 45, 47, 44, 48, 35, 48,
										        31, 64, 76, 45, 46, 27, 39, 42, 31, 30, 76, 41,
										        40, 76, 46, 34, 31, 45, 31, 76, 38, 35, 32, 31,
										        38, 31, 45, 45, 76, 46, 34, 35, 40, 33, 45, 64,
										        78, 76, 76, 76, 76, 76, 76, 76, 76, 20, 34, 31,
										        76, 34, 27, 40, 30, 76, 46, 34, 27, 46, 76, 39,
										        41, 29, 37, 31, 30, 76, 46, 34, 31, 39, 64, 76,
										        27, 40, 30, 76, 46, 34, 31, 76, 34, 31, 27, 44,
										        46, 76, 46, 34, 27, 46, 76, 32, 31, 30, 66, 78,
										        76, 76, 76, 76, 76, 76, 76, 76, 1, 40, 30, 76,
										        41, 40, 76, 46, 34, 31, 76, 42, 31, 30, 31, 45,
										        46, 27, 38, 64, 76, 46, 34, 31, 45, 31, 76, 49,
										        41, 44, 30, 45, 76, 27, 42, 42, 31, 27, 44, 65,
										        78, 76, 76, 76, 76, 76, 76, 76, 76, 13, 51, 76,
										        40, 27, 39, 31, 76, 35, 45, 76, 15, 52, 51, 39,
										        27, 40, 30, 35, 27, 45, 64, 76, 11, 35, 40, 33,
										        76, 41, 32, 76, 11, 35, 40, 33, 45, 66, 78, 76,
										        76, 76, 76, 76, 76, 76, 76, 12, 41, 41, 37, 76,
										        41, 40, 76, 39, 51, 76, 23, 41, 44, 37, 45, 64,
										        76, 51, 31, 76, 13, 35, 33, 34, 46, 51, 64, 76,
										        27, 40, 30, 76, 30, 31, 45, 42, 27, 35, 44, 67,
										        78, 76, 76, 76, 76, 76, 76, 76, 76, 14, 41, 46,
										        34, 35, 40, 33, 76, 28, 31, 45, 35, 30, 31, 76,
										        44, 31, 39, 27, 35, 40, 45, 63, 76, 18, 41, 47,
										        40, 30, 76, 46, 34, 31, 76, 30, 31, 29, 27, 51,
										        78, 76, 76, 76, 76, 76, 76, 76, 76, 15, 32, 76,
										        46, 34, 27, 46, 76, 29, 41, 38, 41, 45, 45, 27,
										        38, 76, 23, 44, 31, 29, 37, 64, 76, 28, 41, 47,
										        40, 30, 38, 31, 45, 45, 76, 27, 40, 30, 76, 28,
										        27, 44, 31, 78, 76, 76, 76, 76, 76, 76, 76, 76,
										        20, 34, 31, 76, 38, 41, 40, 31, 76, 27, 40, 30,
										        76, 38, 31, 48, 31, 38, 76, 45, 27, 40, 30, 45,
										        76, 45, 46, 44, 31, 46, 29, 34, 76, 32, 27, 44,
										        76, 27, 49, 27, 51, 78, 76, 76, 76, 76, 76, 76,
										        76, 76]]       ),
										    torch.tensor([[0, 0, 0, 1_069, 11]]       ),
										    torch.tensor([[0, 0, 0, 1_069, 11]]       ),
										]
										# fmt: on
										self.assertTrue(torch.allclose(tokens[0]							,					EXPECTED_OUTPUT[0]       )       )
										self.assertTrue(torch.allclose(tokens[1]							,					EXPECTED_OUTPUT[1]       )       )
										self.assertTrue(torch.allclose(tokens[2]							,					EXPECTED_OUTPUT[2]       )       )







					@require_torch
					def   UpperCAmelCase__       (     self       :   int       ):






										'''simple docstring'''






										import torch

										lowercase_  							=			JukeboxTokenizer.from_pretrained("""openai/jukebox-5b-lyrics"""       )
										lowercase_  							=			tokenizer(**self.metas       )["""input_ids"""]
										# fmt: off
										lowercase_  							=			[
										    torch.tensor([[
										        0, 0, 0, 1_069, 11, -1, -1, -1, -1, 9, 77, 39,
										        31, 46, 77, 27, 77, 46, 44, 27, 48, 31, 38, 38,
										        31, 44, 77, 32, 44, 41, 39, 77, 27, 40, 77, 27,
										        40, 46, 35, 43, 47, 31, 77, 38, 27, 40, 30, 64,
										        79, 77, 77, 77, 77, 77, 77, 77, 77, 23, 34, 41,
										        77, 45, 27, 35, 30, 77, 72, 20, 49, 41, 77, 48,
										        27, 45, 46, 77, 27, 40, 30, 77, 46, 44, 47, 40,
										        37, 38, 31, 45, 45, 77, 38, 31, 33, 45, 77, 41,
										        32, 77, 45, 46, 41, 40, 31, 79, 77, 77, 77, 77,
										        77, 77, 77, 77, 19, 46, 27, 40, 30, 77, 35, 40,
										        77, 46, 34, 31, 77, 30, 31, 45, 31, 44, 46, 63,
										        77, 63, 77, 63, 77, 63, 77, 14, 31, 27, 44, 77,
										        46, 34, 31, 39, 64, 77, 41, 40, 77, 46, 34, 31,
										        77, 45, 27, 40, 30, 64, 79, 77, 77, 77, 77, 77,
										        77, 77, 77, 8, 27, 38, 32, 77, 45, 47, 40, 37,
										        77, 27, 77, 45, 34, 27, 46, 46, 31, 44, 31, 30,
										        77, 48, 35, 45, 27, 33, 31, 77, 38, 35, 31, 45,
										        64, 77, 49, 34, 41, 45, 31, 77, 32, 44, 41, 49,
										        40, 64, 79, 77, 77, 77, 77, 77, 77, 77, 77, 1,
										        40, 30, 77, 49, 44, 35, 40, 37, 38, 31, 30, 77,
										        38, 35, 42, 64, 77, 27, 40, 30, 77, 45, 40, 31,
										        31, 44, 77, 41, 32, 77, 29, 41, 38, 30, 77, 29,
										        41, 39, 39, 27, 40, 30, 64, 79, 77, 77, 77, 77,
										        77, 77, 77, 77, 20, 31, 38, 38, 77, 46, 34, 27,
										        46, 77, 35, 46, 45, 77, 45, 29, 47, 38, 42, 46,
										        41, 44, 77, 49, 31, 38, 38, 77, 46, 34, 41, 45,
										        31, 77, 42, 27, 45, 45, 35, 41, 40, 45, 77, 44,
										        31, 27, 30, 79, 77, 77, 77, 77, 77, 77, 77, 77,
										        23, 34, 35, 29, 34, 77, 51, 31, 46, 77, 45, 47,
										        44, 48, 35, 48, 31, 64, 77, 45, 46, 27, 39, 42,
										        31, 30, 77, 41, 40, 77, 46, 34, 31, 45, 31, 77,
										        38, 35, 32, 31, 38, 31, 45, 45, 77, 46, 34, 35,
										        40, 33, 45, 64, 79, 77, 77, 77, 77, 77, 77, 77,
										        77, 20, 34, 31, 77, 34, 27, 40, 30, 77, 46, 34,
										        27, 46, 77, 39, 41, 29, 37, 31, 30, 77, 46, 34,
										        31, 39, 64, 77, 27, 40, 30, 77, 46, 34, 31, 77,
										        34, 31, 27, 44, 46, 77, 46, 34, 27, 46, 77, 32,
										        31, 30, 66, 79, 77, 77, 77, 77, 77, 77, 77, 77,
										        1, 40, 30, 77, 41, 40, 77, 46, 34, 31, 77, 42,
										        31, 30, 31, 45, 46, 27, 38, 64, 77, 46, 34, 31,
										        45, 31, 77, 49, 41, 44, 30, 45, 77, 27, 42, 42,
										        31, 27, 44, 65, 79, 77, 77, 77, 77, 77, 77, 77,
										        77, 13, 51, 77, 40, 27, 39, 31, 77, 35, 45, 77,
										        15, 52, 51, 39, 27, 40, 30, 35, 27, 45, 64, 77,
										        11, 35, 40, 33, 77, 41, 32, 77, 11, 35, 40, 33,
										        45, 66, 79, 77, 77, 77, 77, 77, 77, 77, 77, 12,
										        41, 41, 37, 77, 41, 40, 77, 39, 51, 77, 23, 41,
										        44, 37, 45, 64, 77, 51, 31, 77, 13, 35, 33, 34,
										        46, 51, 64, 77, 27, 40, 30, 77, 30, 31, 45, 42,
										        27, 35, 44, 67, 79, 77, 77, 77, 77, 77, 77, 77,
										        77, 14, 41, 46, 34, 35, 40, 33, 77, 28, 31, 45,
										        35, 30, 31, 77, 44, 31, 39, 27, 35, 40, 45, 63,
										        77, 18, 41, 47, 40, 30, 77, 46, 34, 31, 77, 30,
										        31, 29, 27, 51, 79, 77, 77, 77, 77, 77, 77, 77,
										        77, 15, 32, 77, 46, 34, 27, 46, 77, 29, 41, 38,
										        41, 45, 45, 27, 38, 77, 23, 44, 31, 29, 37, 64,
										        77, 28, 41, 47, 40, 30, 38, 31, 45, 45, 77, 27,
										        40, 30, 77, 28, 27, 44, 31, 79, 77, 77, 77, 77,
										        77, 77, 77, 77, 20, 34, 31, 77, 38, 41, 40, 31,
										        77, 27, 40, 30, 77, 38, 31, 48, 31, 38, 77, 45,
										        27, 40, 30, 45, 77, 45, 46, 44, 31, 46, 29, 34,
										        77, 32, 27, 44, 77, 27, 49, 27, 51, 79, 77, 77,
										        77, 77, 77, 77, 77, 77]]       ),
										    torch.tensor([[0, 0, 0, 1_069, 11, -1, -1, -1, -1]]       ),
										    torch.tensor([[0, 0, 0, 1_069, 11, -1, -1, -1, -1]]       ),
										]
										# fmt: on
										self.assertTrue(torch.allclose(tokens[0]							,					EXPECTED_OUTPUT[0]       )       )
										self.assertTrue(torch.allclose(tokens[1]							,					EXPECTED_OUTPUT[1]       )       )
										self.assertTrue(torch.allclose(tokens[2]							,					EXPECTED_OUTPUT[2]       )       )






from dataclasses import asdict, dataclass
from typing import Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging


a					  =     logging.get_logger(__name__)

# TODO Update this
a					  =     {
    'facebook/esm-1b': 'https://huggingface.co/facebook/esm-1b/resolve/main/config.json',
    # See all ESM models at https://huggingface.co/models?filter=esm
}







class        UpperCamelCase__   (	__magic_name__   ):
					__SCREAMING_SNAKE_CASE       :  Tuple						=   'esm'

					def __init__(     self       :   Optional[Any]							,					UpperCamelCase__       :   Dict=None							,					UpperCamelCase__       :   Tuple=None							,					UpperCamelCase__       :   List[Any]=None							,					UpperCamelCase__       :   List[Any]=768							,					UpperCamelCase__       :   Tuple=12							,					UpperCamelCase__       :   Optional[int]=12							,					UpperCamelCase__       :   Dict=3_072							,					UpperCamelCase__       :   Union[str, Any]=0.1							,					UpperCamelCase__       :   Dict=0.1							,					UpperCamelCase__       :   Optional[int]=1_026							,					UpperCamelCase__       :   Any=0.02							,					UpperCamelCase__       :   Dict=1e-12							,					UpperCamelCase__       :   List[str]="absolute"							,					UpperCamelCase__       :   Union[str, Any]=True							,					UpperCamelCase__       :   Tuple=None							,					UpperCamelCase__       :   Union[str, Any]=False							,					UpperCamelCase__       :   Optional[int]=False							,					UpperCamelCase__       :   Dict=None							,					UpperCamelCase__       :   str=None							,					**UpperCamelCase__       :   str							,					):






										'''simple docstring'''






										super().__init__(pad_token_id=UpperCamelCase__							,					mask_token_id=UpperCamelCase__							,					**UpperCamelCase__       )

										lowercase_  							=			vocab_size
										lowercase_  							=			hidden_size
										lowercase_  							=			num_hidden_layers
										lowercase_  							=			num_attention_heads
										lowercase_  							=			intermediate_size
										lowercase_  							=			hidden_dropout_prob
										lowercase_  							=			attention_probs_dropout_prob
										lowercase_  							=			max_position_embeddings
										lowercase_  							=			initializer_range
										lowercase_  							=			layer_norm_eps
										lowercase_  							=			position_embedding_type
										lowercase_  							=			use_cache
										lowercase_  							=			emb_layer_norm_before
										lowercase_  							=			token_dropout
										lowercase_  							=			is_folding_model
										if is_folding_model:
															if esmfold_config is None:
																				logger.info("""No esmfold_config supplied for folding model, using default values."""       )
																				lowercase_  							=			EsmFoldConfig()
															elif isinstance(UpperCamelCase__							,					UpperCamelCase__       ):
																				lowercase_  							=			EsmFoldConfig(**UpperCamelCase__       )
															lowercase_  							=			esmfold_config
															if vocab_list is None:
																				logger.warning("""No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!"""       )
																				lowercase_  							=			get_default_vocab_list()
															else:
																				lowercase_  							=			vocab_list
										else:
															lowercase_  							=			None
															lowercase_  							=			None
										if self.esmfold_config is not None and getattr(self.esmfold_config							,					"""use_esm_attn_map"""							,					UpperCamelCase__       ):
															raise ValueError("""The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!"""       )







					def   UpperCAmelCase__       (     self       :   List[Any]       ):






										'''simple docstring'''






										lowercase_  							=			super().to_dict()
										if isinstance(self.esmfold_config							,					UpperCamelCase__       ):
															lowercase_  							=			self.esmfold_config.to_dict()
										return output







@dataclass
class        UpperCamelCase__   :
					__SCREAMING_SNAKE_CASE       :  str						=   None
					__SCREAMING_SNAKE_CASE       :  bool						=   True
					__SCREAMING_SNAKE_CASE       :  bool						=   False
					__SCREAMING_SNAKE_CASE       :  bool						=   False
					__SCREAMING_SNAKE_CASE       :  bool						=   False
					__SCREAMING_SNAKE_CASE       :  float						=   0

					__SCREAMING_SNAKE_CASE       :  bool						=   True
					__SCREAMING_SNAKE_CASE       :  bool						=   False

					__SCREAMING_SNAKE_CASE       :  int						=   128
					__SCREAMING_SNAKE_CASE       :  "TrunkConfig"						=   None

					def   UpperCAmelCase__       (     self       :   Tuple       ):






										'''simple docstring'''






										if self.trunk is None:
															lowercase_  							=			TrunkConfig()
										elif isinstance(self.trunk							,					UpperCamelCase__       ):
															lowercase_  							=			TrunkConfig(**self.trunk       )







					def   UpperCAmelCase__       (     self       :   Union[str, Any]       ):






										'''simple docstring'''






										lowercase_  							=			asdict(self       )
										lowercase_  							=			self.trunk.to_dict()
										return output







@dataclass
class        UpperCamelCase__   :
					__SCREAMING_SNAKE_CASE       :  int						=   48
					__SCREAMING_SNAKE_CASE       :  int						=   1024
					__SCREAMING_SNAKE_CASE       :  int						=   128
					__SCREAMING_SNAKE_CASE       :  int						=   32
					__SCREAMING_SNAKE_CASE       :  int						=   32
					__SCREAMING_SNAKE_CASE       :  int						=   32
					__SCREAMING_SNAKE_CASE       :  float						=   0
					__SCREAMING_SNAKE_CASE       :  float						=   0
					__SCREAMING_SNAKE_CASE       :  bool						=   False
					__SCREAMING_SNAKE_CASE       :  int						=   4
					__SCREAMING_SNAKE_CASE       :  Optional[int]						=   128
					__SCREAMING_SNAKE_CASE       :  "StructureModuleConfig"						=   None

					def   UpperCAmelCase__       (     self       :   Union[str, Any]       ):






										'''simple docstring'''






										if self.structure_module is None:
															lowercase_  							=			StructureModuleConfig()
										elif isinstance(self.structure_module							,					UpperCamelCase__       ):
															lowercase_  							=			StructureModuleConfig(**self.structure_module       )

										if self.max_recycles <= 0:
															raise ValueError(F'''`max_recycles` should be positive, got {self.max_recycles}.'''       )
										if self.sequence_state_dim % self.sequence_state_dim != 0:
															raise ValueError(
															    """`sequence_state_dim` should be a round multiple of `sequence_state_dim`, got"""
															    F''' {self.sequence_state_dim} and {self.sequence_state_dim}.'''       )
										if self.pairwise_state_dim % self.pairwise_state_dim != 0:
															raise ValueError(
															    """`pairwise_state_dim` should be a round multiple of `pairwise_state_dim`, got"""
															    F''' {self.pairwise_state_dim} and {self.pairwise_state_dim}.'''       )

										lowercase_  							=			self.sequence_state_dim // self.sequence_head_width
										lowercase_  							=			self.pairwise_state_dim // self.pairwise_head_width

										if self.sequence_state_dim != sequence_num_heads * self.sequence_head_width:
															raise ValueError(
															    """`sequence_state_dim` should be equal to `sequence_num_heads * sequence_head_width, got"""
															    F''' {self.sequence_state_dim} != {sequence_num_heads} * {self.sequence_head_width}.'''       )
										if self.pairwise_state_dim != pairwise_num_heads * self.pairwise_head_width:
															raise ValueError(
															    """`pairwise_state_dim` should be equal to `pairwise_num_heads * pairwise_head_width, got"""
															    F''' {self.pairwise_state_dim} != {pairwise_num_heads} * {self.pairwise_head_width}.'''       )
										if self.pairwise_state_dim % 2 != 0:
															raise ValueError(F'''`pairwise_state_dim` should be even, got {self.pairwise_state_dim}.'''       )

										if self.dropout >= 0.4:
															raise ValueError(F'''`dropout` should not be greater than 0.4, got {self.dropout}.'''       )







					def   UpperCAmelCase__       (     self       :   Any       ):






										'''simple docstring'''






										lowercase_  							=			asdict(self       )
										lowercase_  							=			self.structure_module.to_dict()
										return output







@dataclass
class        UpperCamelCase__   :
					__SCREAMING_SNAKE_CASE       :  int						=   384
					__SCREAMING_SNAKE_CASE       :  int						=   128
					__SCREAMING_SNAKE_CASE       :  int						=   16
					__SCREAMING_SNAKE_CASE       :  int						=   128
					__SCREAMING_SNAKE_CASE       :  int						=   12
					__SCREAMING_SNAKE_CASE       :  int						=   4
					__SCREAMING_SNAKE_CASE       :  int						=   8
					__SCREAMING_SNAKE_CASE       :  float						=   0.1
					__SCREAMING_SNAKE_CASE       :  int						=   8
					__SCREAMING_SNAKE_CASE       :  int						=   1
					__SCREAMING_SNAKE_CASE       :  int						=   2
					__SCREAMING_SNAKE_CASE       :  int						=   7
					__SCREAMING_SNAKE_CASE       :  int						=   10
					__SCREAMING_SNAKE_CASE       :  float						=   1e-8
					__SCREAMING_SNAKE_CASE       :  float						=   1e5

					def   UpperCAmelCase__       (     self       :   List[str]       ):






										'''simple docstring'''






										return asdict(self       )



def 			UpperCAmelCase_				(      ):
					return (
					    "<cls>",
					    "<pad>",
					    "<eos>",
					    "<unk>",
					    "L",
					    "A",
					    "G",
					    "V",
					    "S",
					    "E",
					    "R",
					    "T",
					    "I",
					    "D",
					    "P",
					    "K",
					    "Q",
					    "N",
					    "F",
					    "Y",
					    "M",
					    "H",
					    "W",
					    "C",
					    "X",
					    "B",
					    "U",
					    "Z",
					    "O",
					    ".",
					    "-",
					    "<null_1>",
					    "<mask>",
					)


