

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


_a:     str							=      logging.get_logger(__name__)

_a:     Union[str, Any]							=      {
    """edbeeching/decision-transformer-gym-hopper-medium""": (
        """https://huggingface.co/edbeeching/decision-transformer-gym-hopper-medium/resolve/main/config.json"""
    ),
    # See all DecisionTransformer models at https://huggingface.co/models?filter=decision_transformer
}

class     __UpperCamelCase     (       lowercase    ):
	SCREAMING_SNAKE_CASE__     					=  "decision_transformer"
	SCREAMING_SNAKE_CASE__     					=  ["past_key_values"]
	SCREAMING_SNAKE_CASE__     					=  {
	    "max_position_embeddings": "n_positions",
	    "num_attention_heads": "n_head",
	    "num_hidden_layers": "n_layer",
	}




	def __init__(		self						:    str		,						lowerCAmelCase						:    Tuple=17		,						lowerCAmelCase						:    Any=4		,						lowerCAmelCase						:    Tuple=128		,						lowerCAmelCase						:    Union[str, Any]=4_096		,						lowerCAmelCase						:    str=True		,						lowerCAmelCase						:    List[Any]=1		,						lowerCAmelCase						:    List[str]=1_024		,						lowerCAmelCase						:    Union[str, Any]=3		,						lowerCAmelCase						:    Tuple=1		,						lowerCAmelCase						:    List[str]=None		,						lowerCAmelCase						:    Any="relu"		,						lowerCAmelCase						:    List[Any]=0.1		,						lowerCAmelCase						:    List[str]=0.1		,						lowerCAmelCase						:    str=0.1		,						lowerCAmelCase						:    Tuple=1e-5		,						lowerCAmelCase						:    Optional[int]=0.02		,						lowerCAmelCase						:    Optional[int]=True		,						lowerCAmelCase						:    List[Any]=True		,						lowerCAmelCase						:    Any=50_256		,						lowerCAmelCase						:    Optional[int]=50_256		,						lowerCAmelCase						:    Union[str, Any]=False		,						lowerCAmelCase						:    Optional[Any]=False		,						**lowerCAmelCase						:    int		,						):


		'''simple docstring'''







		UpperCAmelCase_												=      state_dim
		UpperCAmelCase_												=      act_dim
		UpperCAmelCase_												=      hidden_size
		UpperCAmelCase_												=      max_ep_len
		UpperCAmelCase_												=      action_tanh
		UpperCAmelCase_												=      vocab_size
		UpperCAmelCase_												=      n_positions
		UpperCAmelCase_												=      n_layer
		UpperCAmelCase_												=      n_head
		UpperCAmelCase_												=      n_inner
		UpperCAmelCase_												=      activation_function
		UpperCAmelCase_												=      resid_pdrop
		UpperCAmelCase_												=      embd_pdrop
		UpperCAmelCase_												=      attn_pdrop
		UpperCAmelCase_												=      layer_norm_epsilon
		UpperCAmelCase_												=      initializer_range
		UpperCAmelCase_												=      scale_attn_weights
		UpperCAmelCase_												=      use_cache
		UpperCAmelCase_												=      scale_attn_by_inverse_layer_idx
		UpperCAmelCase_												=      reorder_and_upcast_attn

		UpperCAmelCase_												=      bos_token_id
		UpperCAmelCase_												=      eos_token_id

		super().__init__(bos_token_id=SCREAMING_SNAKE_CASE__		,						eos_token_id=SCREAMING_SNAKE_CASE__		,						**SCREAMING_SNAKE_CASE__				)

import unittest

from transformers import JukeboxTokenizer
from transformers.testing_utils import require_torch
class       _lowerCAmelCase  (		unittest.TestCase							):

 """simple docstring"""
 SCREAMING_SNAKE_CASE_				:					List[Any]       						=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      __lowerCAmelCase  (						self			:			Tuple   ):






      """simple docstring"""



      import torch

      UpperCamelCase						    =					JukeboxTokenizer.from_pretrained('openai/jukebox-1b-lyrics'   )
      UpperCamelCase						    =					tokenizer(**self.metas   )['input_ids']
      # fmt: off
      UpperCamelCase						    =					[
          torch.tensor([[
              0, 0, 0, 71_69, 5_07, 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, 10_69, 11]]   ),
          torch.tensor([[0, 0, 0, 10_69, 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      __lowerCAmelCase  (						self			:			Union[str, Any]   ):






      """simple docstring"""



      import torch

      UpperCamelCase						    =					JukeboxTokenizer.from_pretrained('openai/jukebox-5b-lyrics'   )
      UpperCamelCase						    =					tokenizer(**self.metas   )['input_ids']
      # fmt: off
      UpperCamelCase						    =					[
          torch.tensor([[
              0, 0, 0, 10_69, 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, 10_69, 11, -1, -1, -1, -1]]   ),
          torch.tensor([[0, 0, 0, 10_69, 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]   )   )


