






'''simple docstring'''



from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


SCREAMING_SNAKE_CASE							=						logging.get_logger(__name__)

SCREAMING_SNAKE_CASE							=						{
    'distilbert-base-uncased': 'https://huggingface.co/distilbert-base-uncased/resolve/main/config.json',
    'distilbert-base-uncased-distilled-squad': (
        'https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/config.json'
    ),
    'distilbert-base-cased': 'https://huggingface.co/distilbert-base-cased/resolve/main/config.json',
    'distilbert-base-cased-distilled-squad': (
        'https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/config.json'
    ),
    'distilbert-base-german-cased': 'https://huggingface.co/distilbert-base-german-cased/resolve/main/config.json',
    'distilbert-base-multilingual-cased': (
        'https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/config.json'
    ),
    'distilbert-base-uncased-finetuned-sst-2-english': (
        'https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/resolve/main/config.json'
    ),
}
class UpperCAmelCase_       (     __A					):




      """simple docstring"""




      UpperCamelCase_             =      '''distilbert'''
      UpperCamelCase_             =      {
          '''hidden_size''': '''dim''',
          '''num_attention_heads''': '''n_heads''',
          '''num_hidden_layers''': '''n_layers''',
      }






      def __init__(			self      :		Optional[int]  ,   UpperCAmelCase      :		str=3_0522  ,   UpperCAmelCase      :		Any=512  ,   UpperCAmelCase      :		int=False  ,   UpperCAmelCase      :		List[str]=6  ,   UpperCAmelCase      :		Optional[int]=12  ,   UpperCAmelCase      :		Dict=768  ,   UpperCAmelCase      :		str=4 * 768  ,   UpperCAmelCase      :		int=0.1  ,   UpperCAmelCase      :		Optional[int]=0.1  ,   UpperCAmelCase      :		List[Any]="gelu"  ,   UpperCAmelCase      :		List[Any]=0.0_2  ,   UpperCAmelCase      :		List[Any]=0.1  ,   UpperCAmelCase      :		Optional[int]=0.2  ,   UpperCAmelCase      :		List[Any]=0  ,   **UpperCAmelCase      :		Dict  ,   )	->	str:


       '''simple docstring'''

       lowercase       : Dict						=vocab_size
       lowercase       : Union[str, Any]						=max_position_embeddings
       lowercase       : Any						=sinusoidal_pos_embds
       lowercase       : Dict						=n_layers
       lowercase       : Optional[int]						=n_heads
       lowercase       : Optional[Any]						=dim
       lowercase       : Any						=hidden_dim
       lowercase       : Optional[int]						=dropout
       lowercase       : str						=attention_dropout
       lowercase       : Any						=activation
       lowercase       : Optional[Any]						=initializer_range
       lowercase       : Optional[int]						=qa_dropout
       lowercase       : Optional[Any]						=seq_classif_dropout
       super().__init__(**UpperCAmelCase  ,   pad_token_id=UpperCAmelCase		)





class UpperCAmelCase_       (     __A					):




      """simple docstring"""




      @property
      def 					A__   (			self      :		Tuple		)	->	Mapping[str, Mapping[int, str]]:


       '''simple docstring'''

       if self.task == "multiple-choice":
        lowercase       : Tuple						={0: '''batch''', 1: '''choice''', 2: '''sequence'''}
       else:
        lowercase       : Any						={0: '''batch''', 1: '''sequence'''}
       return OrderedDict(
           [
               ('''input_ids''', dynamic_axis),
               ('''attention_mask''', dynamic_axis),
           ]		)





"""simple docstring"""


import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import TaBlock, TaConfig, TaLayerNorm

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin



class 	__a	(   _lowerCAmelCase     ,  _lowerCAmelCase     ,  _lowerCAmelCase						):






       @register_to_config
       def __init__(	self			:				Union[str, Any]							,   UpperCAmelCase_			:				int							,   UpperCAmelCase_			:				int							,   UpperCAmelCase_			:				int							,   UpperCAmelCase_			:				float							,   UpperCAmelCase_			:				int							,   UpperCAmelCase_			:				int							,   UpperCAmelCase_			:				int							,   UpperCAmelCase_			:				int							,   UpperCAmelCase_			:				str							,   UpperCAmelCase_			:				bool = False							,   )->	str:







              """simple docstring"""
              super().__init__()

              UpperCamelCase       			=		nn.Embedding(UpperCAmelCase_							,   UpperCAmelCase_      )

              UpperCamelCase       			=		nn.Embedding(UpperCAmelCase_							,   UpperCAmelCase_      )
              UpperCamelCase       			=		False

              UpperCamelCase       			=		nn.Dropout(p=UpperCAmelCase_      )

              UpperCamelCase       			=		TaConfig(
                  vocab_size=UpperCAmelCase_							,   d_model=UpperCAmelCase_							,   num_heads=UpperCAmelCase_							,   d_kv=UpperCAmelCase_							,   d_ff=UpperCAmelCase_							,   dropout_rate=UpperCAmelCase_							,   feed_forward_proj=UpperCAmelCase_							,   is_decoder=UpperCAmelCase_							,   is_encoder_decoder=UpperCAmelCase_							,   )

              UpperCamelCase       			=		nn.ModuleList()
              for lyr_num in range(UpperCAmelCase_      ):
                     UpperCamelCase       			=		TaBlock(UpperCAmelCase_      )
                     self.encoders.append(UpperCAmelCase_      )

              UpperCamelCase       			=		TaLayerNorm(UpperCAmelCase_      )
              UpperCamelCase       			=		nn.Dropout(p=UpperCAmelCase_      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				int							,   UpperCAmelCase_			:				int							,   UpperCAmelCase_			:				str      )->	List[Any]:







              """simple docstring"""
              UpperCamelCase       			=		self.token_embedder(UpperCAmelCase_      )

              UpperCamelCase       			=		encoder_input_tokens.shape[1]
              UpperCamelCase       			=		torch.arange(UpperCAmelCase_							,   device=encoder_input_tokens.device      )
              x += self.position_encoding(UpperCAmelCase_      )

              UpperCamelCase       			=		self.dropout_pre(UpperCAmelCase_      )

              # inverted the attention mask
              UpperCamelCase       			=		encoder_input_tokens.size()
              UpperCamelCase       			=		self.get_extended_attention_mask(UpperCAmelCase_							,   UpperCAmelCase_      )

              for lyr in self.encoders:
                     UpperCamelCase       			=		lyr(UpperCAmelCase_							,   UpperCAmelCase_      )[0]
              UpperCamelCase       			=		self.layer_norm(UpperCAmelCase_      )

              return self.dropout_post(UpperCAmelCase_      ), encoder_inputs_mask





