





'''simple docstring'''




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_SCREAMING_SNAKE_CASE							     =			{
    "configuration_electra": ["ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP", "ElectraConfig", "ElectraOnnxConfig"],
    "tokenization_electra": ["ElectraTokenizer"],
}

try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _SCREAMING_SNAKE_CASE							     =			["ElectraTokenizerFast"]

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _SCREAMING_SNAKE_CASE							     =			[
     "ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST",
     "ElectraForCausalLM",
     "ElectraForMaskedLM",
     "ElectraForMultipleChoice",
     "ElectraForPreTraining",
     "ElectraForQuestionAnswering",
     "ElectraForSequenceClassification",
     "ElectraForTokenClassification",
     "ElectraModel",
     "ElectraPreTrainedModel",
     "load_tf_weights_in_electra",
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _SCREAMING_SNAKE_CASE							     =			[
     "TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST",
     "TFElectraForMaskedLM",
     "TFElectraForMultipleChoice",
     "TFElectraForPreTraining",
     "TFElectraForQuestionAnswering",
     "TFElectraForSequenceClassification",
     "TFElectraForTokenClassification",
     "TFElectraModel",
     "TFElectraPreTrainedModel",
 ]

try:
 if not is_flax_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _SCREAMING_SNAKE_CASE							     =			[
     "FlaxElectraForCausalLM",
     "FlaxElectraForMaskedLM",
     "FlaxElectraForMultipleChoice",
     "FlaxElectraForPreTraining",
     "FlaxElectraForQuestionAnswering",
     "FlaxElectraForSequenceClassification",
     "FlaxElectraForTokenClassification",
     "FlaxElectraModel",
     "FlaxElectraPreTrainedModel",
 ]


if TYPE_CHECKING:
 from .configuration_electra import ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP, ElectraConfig, ElectraOnnxConfig
 from .tokenization_electra import ElectraTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_electra_fast import ElectraTokenizerFast

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_electra import (
      ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
      ElectraForCausalLM,
      ElectraForMaskedLM,
      ElectraForMultipleChoice,
      ElectraForPreTraining,
      ElectraForQuestionAnswering,
      ElectraForSequenceClassification,
      ElectraForTokenClassification,
      ElectraModel,
      ElectraPreTrainedModel,
      load_tf_weights_in_electra,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_electra import (
      TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFElectraForMaskedLM,
      TFElectraForMultipleChoice,
      TFElectraForPreTraining,
      TFElectraForQuestionAnswering,
      TFElectraForSequenceClassification,
      TFElectraForTokenClassification,
      TFElectraModel,
      TFElectraPreTrainedModel,
  )

 try:
  if not is_flax_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_flax_electra import (
      FlaxElectraForCausalLM,
      FlaxElectraForMaskedLM,
      FlaxElectraForMultipleChoice,
      FlaxElectraForPreTraining,
      FlaxElectraForQuestionAnswering,
      FlaxElectraForSequenceClassification,
      FlaxElectraForTokenClassification,
      FlaxElectraModel,
      FlaxElectraPreTrainedModel,
  )

else:
 import sys

 _SCREAMING_SNAKE_CASE							     =			_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)








'''simple docstring'''




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


_SCREAMING_SNAKE_CASE							     =			logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE							     =			{
    "CarlCochet/trajectory-transformer-halfcheetah-medium-v2": (
        "https://huggingface.co/CarlCochet/trajectory-transformer-halfcheetah-medium-v2/resolve/main/config.json"
    ),
    # See all TrajectoryTransformer models at https://huggingface.co/models?filter=trajectory_transformer
}
class 			_lowerCAmelCase		(			A__	):






     """simple docstring"""






     snake_case_						      =   "trajectory_transformer"
     snake_case_						      =   ["past_key_values"]
     snake_case_						      =   {
         "hidden_size": "n_embd",
         "num_attention_heads": "n_head",
         "num_hidden_layers": "n_layer",
     }



     def __init__(   self							:			Optional[Any]   ,       __snake_case							:			int=1_00   ,       __snake_case							:			List[str]=5   ,       __snake_case							:			str=1   ,       __snake_case							:			Union[str, Any]=1   ,       __snake_case							:			Union[str, Any]=2_49   ,       __snake_case							:			List[Any]=6   ,       __snake_case							:			Optional[int]=17   ,       __snake_case							:			Optional[Any]=25   ,       __snake_case							:			Union[str, Any]=4   ,       __snake_case							:			List[Any]=4   ,       __snake_case							:			Optional[int]=1_28   ,       __snake_case							:			Optional[Any]=0.1   ,       __snake_case							:			Any=0.1   ,       __snake_case							:			int=0.1   ,       __snake_case							:			Union[str, Any]=0.00_06   ,       __snake_case							:			Tuple=5_12   ,       __snake_case							:			int=0.02   ,       __snake_case							:			Any=1e-12   ,       __snake_case							:			Optional[Any]=1   ,       __snake_case							:			List[Any]=True   ,       __snake_case							:			Dict=1   ,       __snake_case							:			Dict=5_02_56   ,       __snake_case							:			Union[str, Any]=5_02_56   ,       **__snake_case							:			Tuple   ,       )->	Optional[int]:
        snake_case 						=   vocab_size
        snake_case 						=   action_weight
        snake_case 						=   reward_weight
        snake_case 						=   value_weight
        snake_case 						=   max_position_embeddings
        snake_case 						=   block_size
        snake_case 						=   action_dim
        snake_case 						=   observation_dim
        snake_case 						=   transition_dim
        snake_case 						=   learning_rate
        snake_case 						=   n_layer
        snake_case 						=   n_head
        snake_case 						=   n_embd
        snake_case 						=   embd_pdrop
        snake_case 						=   attn_pdrop
        snake_case 						=   resid_pdrop
        snake_case 						=   initializer_range
        snake_case 						=   layer_norm_eps
        snake_case 						=   kaiming_initializer_range
        snake_case 						=   use_cache
        super().__init__(pad_token_id=__snake_case   ,       bos_token_id=__snake_case   ,       eos_token_id=__snake_case   ,       **__snake_case						)

