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


lowerCamelCase__     :		Any        = logging.get_logger(__name__)

lowerCamelCase__     :		Optional[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    _snake_case (						UpperCAmelCase_				):
 __lowerCAmelCase	:  Any	   =      'decision_transformer'
 __lowerCAmelCase	:  Optional[Any]	   =      ['past_key_values']
 __lowerCAmelCase	:  int	   =      {
     'max_position_embeddings': 'n_positions',
     'num_attention_heads': 'n_head',
     'num_hidden_layers': 'n_layer',
 }


 def __init__(       self       ,	SCREAMING_SNAKE_CASE_=17       ,	SCREAMING_SNAKE_CASE_=4       ,	SCREAMING_SNAKE_CASE_=1_28       ,	SCREAMING_SNAKE_CASE_=40_96       ,	SCREAMING_SNAKE_CASE_=True       ,	SCREAMING_SNAKE_CASE_=1       ,	SCREAMING_SNAKE_CASE_=10_24       ,	SCREAMING_SNAKE_CASE_=3       ,	SCREAMING_SNAKE_CASE_=1       ,	SCREAMING_SNAKE_CASE_=None       ,	SCREAMING_SNAKE_CASE_="relu"       ,	SCREAMING_SNAKE_CASE_=0.1       ,	SCREAMING_SNAKE_CASE_=0.1       ,	SCREAMING_SNAKE_CASE_=0.1       ,	SCREAMING_SNAKE_CASE_=1E-5       ,	SCREAMING_SNAKE_CASE_=0.0_2       ,	SCREAMING_SNAKE_CASE_=True       ,	SCREAMING_SNAKE_CASE_=True       ,	SCREAMING_SNAKE_CASE_=5_02_56       ,	SCREAMING_SNAKE_CASE_=5_02_56       ,	SCREAMING_SNAKE_CASE_=False       ,	SCREAMING_SNAKE_CASE_=False       ,	**SCREAMING_SNAKE_CASE_       ,	):


   '''simple docstring'''




   lowercase__	: Tuple     		=   state_dim
   lowercase__	: List[Any]     		=   act_dim
   lowercase__	: Optional[int]     		=   hidden_size
   lowercase__	: Tuple     		=   max_ep_len
   lowercase__	: Any     		=   action_tanh
   lowercase__	: Union[str, Any]     		=   vocab_size
   lowercase__	: Any     		=   n_positions
   lowercase__	: Any     		=   n_layer
   lowercase__	: List[Any]     		=   n_head
   lowercase__	: Optional[int]     		=   n_inner
   lowercase__	: int     		=   activation_function
   lowercase__	: Optional[int]     		=   resid_pdrop
   lowercase__	: List[str]     		=   embd_pdrop
   lowercase__	: Any     		=   attn_pdrop
   lowercase__	: str     		=   layer_norm_epsilon
   lowercase__	: Any     		=   initializer_range
   lowercase__	: Optional[int]     		=   scale_attn_weights
   lowercase__	: Optional[Any]     		=   use_cache
   lowercase__	: Union[str, Any]     		=   scale_attn_by_inverse_layer_idx
   lowercase__	: Tuple     		=   reorder_and_upcast_attn

   lowercase__	: Union[str, Any]     		=   bos_token_id
   lowercase__	: int     		=   eos_token_id

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








import shutil
import tempfile
import unittest
from unittest.mock import patch

from transformers import (
    DefaultFlowCallback,
    IntervalStrategy,
    PrinterCallback,
    ProgressCallback,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    is_torch_available,
)
from transformers.testing_utils import require_torch


if is_torch_available():
     from transformers.trainer import DEFAULT_CALLBACKS

     from .test_trainer import RegressionDataset, RegressionModelConfig, RegressionPreTrainedModel



class    _snake_case (						UpperCAmelCase_				):


 def __init__(       self):


   '''simple docstring'''




   lowercase__	: List[Any]     		=   []


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.events.append("""on_init_end""")


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.events.append("""on_train_begin""")


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.events.append("""on_train_end""")


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.events.append("""on_epoch_begin""")


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.events.append("""on_epoch_end""")


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.events.append("""on_step_begin""")


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.events.append("""on_step_end""")


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.events.append("""on_evaluate""")


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.events.append("""on_predict""")


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.events.append("""on_save""")


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.events.append("""on_log""")



 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.events.append("""on_prediction_step""")



@require_torch
class    _snake_case (						unittest.TestCase				):


 def   lowercase__  (       self):


   '''simple docstring'''




   lowercase__	: Dict     		=   tempfile.mkdtemp()


 def   lowercase__  (       self):


   '''simple docstring'''




   shutil.rmtree(self.output_dir)


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_=0       ,	SCREAMING_SNAKE_CASE_=0       ,	SCREAMING_SNAKE_CASE_=64       ,	SCREAMING_SNAKE_CASE_=64       ,	SCREAMING_SNAKE_CASE_=None       ,	SCREAMING_SNAKE_CASE_=False       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__	: Any     		=   RegressionDataset(length=SCREAMING_SNAKE_CASE_)
   lowercase__	: Optional[int]     		=   RegressionDataset(length=SCREAMING_SNAKE_CASE_)
   lowercase__	: Dict     		=   RegressionModelConfig(a=SCREAMING_SNAKE_CASE_       ,	b=SCREAMING_SNAKE_CASE_)
   lowercase__	: Any     		=   RegressionPreTrainedModel(SCREAMING_SNAKE_CASE_)

   lowercase__	: Any     		=   TrainingArguments(self.output_dir       ,	disable_tqdm=SCREAMING_SNAKE_CASE_       ,	report_to=[]       ,	**SCREAMING_SNAKE_CASE_)
   return Trainer(
       SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	train_dataset=SCREAMING_SNAKE_CASE_       ,	eval_dataset=SCREAMING_SNAKE_CASE_       ,	callbacks=SCREAMING_SNAKE_CASE_       ,	)


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.assertEqual(len(SCREAMING_SNAKE_CASE_)       ,	len(SCREAMING_SNAKE_CASE_))

   # Order doesn't matter
   lowercase__	: str     		=   sorted(SCREAMING_SNAKE_CASE_       ,	key=lambda SCREAMING_SNAKE_CASE_: cb.__name__ if isinstance(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_) else cb.__class__.__name__)
   lowercase__	: Tuple     		=   sorted(SCREAMING_SNAKE_CASE_       ,	key=lambda SCREAMING_SNAKE_CASE_: cb.__name__ if isinstance(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_) else cb.__class__.__name__)

   for cba, cba in zip(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_):
     if isinstance(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_) and isinstance(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_):
       self.assertEqual(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_)
     elif isinstance(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_) and not isinstance(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_):
       self.assertEqual(SCREAMING_SNAKE_CASE_       ,	cba.__class__)
     elif not isinstance(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_) and isinstance(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_):
       self.assertEqual(cba.__class__       ,	SCREAMING_SNAKE_CASE_)
     else:
       self.assertEqual(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_)


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__	: int     		=   ["""on_init_end""", """on_train_begin"""]
   lowercase__	: Union[str, Any]     		=   0
   lowercase__	: Union[str, Any]     		=   len(trainer.get_eval_dataloader())
   lowercase__	: Dict     		=   ["""on_prediction_step"""] * len(trainer.get_eval_dataloader()) + ["""on_log""", """on_evaluate"""]
   for _ in range(trainer.state.num_train_epochs):
     expected_events.append("""on_epoch_begin""")
     for _ in range(SCREAMING_SNAKE_CASE_):
       step += 1
       expected_events += ["on_step_begin", "on_step_end"]
       if step % trainer.args.logging_steps == 0:
         expected_events.append("""on_log""")
       if trainer.args.evaluation_strategy == IntervalStrategy.STEPS and step % trainer.args.eval_steps == 0:
         expected_events += evaluation_events.copy()
       if step % trainer.args.save_steps == 0:
         expected_events.append("""on_save""")
     expected_events.append("""on_epoch_end""")
     if trainer.args.evaluation_strategy == IntervalStrategy.EPOCH:
       expected_events += evaluation_events.copy()
   expected_events += ["on_log", "on_train_end"]
   return expected_events


 def   lowercase__  (       self):


   '''simple docstring'''




   lowercase__	: int     		=   self.get_trainer()
   lowercase__	: Union[str, Any]     		=   DEFAULT_CALLBACKS.copy() + [ProgressCallback]
   self.check_callbacks_equality(trainer.callback_handler.callbacks       ,	SCREAMING_SNAKE_CASE_)

   # Callbacks passed at init are added to the default callbacks
   lowercase__	: Any     		=   self.get_trainer(callbacks=[MyTestTrainerCallback])
   expected_callbacks.append(SCREAMING_SNAKE_CASE_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks       ,	SCREAMING_SNAKE_CASE_)

   # TrainingArguments.disable_tqdm controls if use ProgressCallback or PrinterCallback
   lowercase__	: Any     		=   self.get_trainer(disable_tqdm=SCREAMING_SNAKE_CASE_)
   lowercase__	: Tuple     		=   DEFAULT_CALLBACKS.copy() + [PrinterCallback]
   self.check_callbacks_equality(trainer.callback_handler.callbacks       ,	SCREAMING_SNAKE_CASE_)


 def   lowercase__  (       self):


   '''simple docstring'''




   lowercase__	: Any     		=   DEFAULT_CALLBACKS.copy() + [ProgressCallback]
   lowercase__	: Tuple     		=   self.get_trainer()

   # We can add, pop, or remove by class name
   trainer.remove_callback(SCREAMING_SNAKE_CASE_)
   expected_callbacks.remove(SCREAMING_SNAKE_CASE_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks       ,	SCREAMING_SNAKE_CASE_)

   lowercase__	: Optional[int]     		=   self.get_trainer()
   lowercase__	: List[Any]     		=   trainer.pop_callback(SCREAMING_SNAKE_CASE_)
   self.assertEqual(cb.__class__       ,	SCREAMING_SNAKE_CASE_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks       ,	SCREAMING_SNAKE_CASE_)

   trainer.add_callback(SCREAMING_SNAKE_CASE_)
   expected_callbacks.insert(0       ,	SCREAMING_SNAKE_CASE_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks       ,	SCREAMING_SNAKE_CASE_)

   # We can also add, pop, or remove by instance
   lowercase__	: Union[str, Any]     		=   self.get_trainer()
   lowercase__	: Optional[Any]     		=   trainer.callback_handler.callbacks[0]
   trainer.remove_callback(SCREAMING_SNAKE_CASE_)
   expected_callbacks.remove(SCREAMING_SNAKE_CASE_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks       ,	SCREAMING_SNAKE_CASE_)

   lowercase__	: str     		=   self.get_trainer()
   lowercase__	: Optional[Any]     		=   trainer.callback_handler.callbacks[0]
   lowercase__	: Union[str, Any]     		=   trainer.pop_callback(SCREAMING_SNAKE_CASE_)
   self.assertEqual(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks       ,	SCREAMING_SNAKE_CASE_)

   trainer.add_callback(SCREAMING_SNAKE_CASE_)
   expected_callbacks.insert(0       ,	SCREAMING_SNAKE_CASE_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks       ,	SCREAMING_SNAKE_CASE_)



 def   lowercase__  (       self):


   '''simple docstring'''




   import warnings

   # XXX: for now ignore scatter_gather warnings in this test since it's not relevant to what's being tested
   warnings.simplefilter(action="""ignore"""       ,	category=SCREAMING_SNAKE_CASE_)

   lowercase__	: Union[str, Any]     		=   self.get_trainer(callbacks=[MyTestTrainerCallback])
   trainer.train()
   lowercase__	: Union[str, Any]     		=   trainer.callback_handler.callbacks[-2].events
   self.assertEqual(SCREAMING_SNAKE_CASE_       ,	self.get_expected_events(SCREAMING_SNAKE_CASE_))

   # Independent log/save/eval
   lowercase__	: List[Any]     		=   self.get_trainer(callbacks=[MyTestTrainerCallback]       ,	logging_steps=5)
   trainer.train()
   lowercase__	: List[str]     		=   trainer.callback_handler.callbacks[-2].events
   self.assertEqual(SCREAMING_SNAKE_CASE_       ,	self.get_expected_events(SCREAMING_SNAKE_CASE_))

   lowercase__	: Optional[Any]     		=   self.get_trainer(callbacks=[MyTestTrainerCallback]       ,	save_steps=5)
   trainer.train()
   lowercase__	: Dict     		=   trainer.callback_handler.callbacks[-2].events
   self.assertEqual(SCREAMING_SNAKE_CASE_       ,	self.get_expected_events(SCREAMING_SNAKE_CASE_))

   lowercase__	: Any     		=   self.get_trainer(callbacks=[MyTestTrainerCallback]       ,	eval_steps=5       ,	evaluation_strategy="""steps""")
   trainer.train()
   lowercase__	: int     		=   trainer.callback_handler.callbacks[-2].events
   self.assertEqual(SCREAMING_SNAKE_CASE_       ,	self.get_expected_events(SCREAMING_SNAKE_CASE_))

   lowercase__	: Tuple     		=   self.get_trainer(callbacks=[MyTestTrainerCallback]       ,	evaluation_strategy="""epoch""")
   trainer.train()
   lowercase__	: Optional[int]     		=   trainer.callback_handler.callbacks[-2].events
   self.assertEqual(SCREAMING_SNAKE_CASE_       ,	self.get_expected_events(SCREAMING_SNAKE_CASE_))

   # A bit of everything
   lowercase__	: Any     		=   self.get_trainer(
       callbacks=[MyTestTrainerCallback]       ,	logging_steps=3       ,	save_steps=10       ,	eval_steps=5       ,	evaluation_strategy="""steps"""       ,	)
   trainer.train()
   lowercase__	: str     		=   trainer.callback_handler.callbacks[-2].events
   self.assertEqual(SCREAMING_SNAKE_CASE_       ,	self.get_expected_events(SCREAMING_SNAKE_CASE_))

   # warning should be emitted for duplicated callbacks
   with patch("""transformers.trainer_callback.logger.warning""") as warn_mock:
     lowercase__	: Dict     		=   self.get_trainer(
         callbacks=[MyTestTrainerCallback, MyTestTrainerCallback]       ,	)
     assert str(SCREAMING_SNAKE_CASE_) in warn_mock.call_args[0][0]






