

'''simple docstring'''
from ...configuration_utils import PretrainedConfig
from ...utils import logging


_snake_case     =  logging.get_logger(__name__)

_snake_case     =  {
    'alibaba-damo/mgp-str-base': 'https://huggingface.co/alibaba-damo/mgp-str-base/resolve/main/config.json',
}


class    a__  (							lowerCamelCase_  ):
   _SCREAMING_SNAKE_CASE    :     List[str]	       =							'mgp-str'


   def __init__( self			,				_UpperCamelCase=[32, 128]			,				_UpperCamelCase=4			,				_UpperCamelCase=3			,				_UpperCamelCase=27			,				_UpperCamelCase=38			,				_UpperCamelCase=50257			,				_UpperCamelCase=30522			,				_UpperCamelCase=768			,				_UpperCamelCase=12			,				_UpperCamelCase=12			,				_UpperCamelCase=4.0			,				_UpperCamelCase=True			,				_UpperCamelCase=False			,				_UpperCamelCase=1E-5			,				_UpperCamelCase=0.0			,				_UpperCamelCase=0.0			,				_UpperCamelCase=0.0			,				_UpperCamelCase=False			,				_UpperCamelCase=0.0_2			,				**_UpperCamelCase			,				):





      """simple docstring"""





      super().__init__(**_UpperCamelCase  )

      _lowercase			:      Optional[Any]			     =   image_size
      _lowercase			:      Union[str, Any]			     =   patch_size
      _lowercase			:      Union[str, Any]			     =   num_channels
      _lowercase			:      Optional[int]			     =   max_token_length
      _lowercase			:      Optional[Any]			     =   num_character_labels
      _lowercase			:      Tuple			     =   num_bpe_labels
      _lowercase			:      Dict			     =   num_wordpiece_labels
      _lowercase			:      str			     =   hidden_size
      _lowercase			:      Optional[int]			     =   num_hidden_layers
      _lowercase			:      Tuple			     =   num_attention_heads
      _lowercase			:      List[str]			     =   mlp_ratio
      _lowercase			:      Dict			     =   distilled
      _lowercase			:      int			     =   layer_norm_eps
      _lowercase			:      List[Any]			     =   drop_rate
      _lowercase			:      Tuple			     =   qkv_bias
      _lowercase			:      List[Any]			     =   attn_drop_rate
      _lowercase			:      Optional[Any]			     =   drop_path_rate
      _lowercase			:      List[str]			     =   output_aa_attentions
      _lowercase			:      Tuple			     =   initializer_range









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  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE			,			SCREAMING_SNAKE_CASE			,			SCREAMING_SNAKE_CASE							):




  @register_to_config
  def __init__(				self							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase = False							,   )  -> Tuple:
        '''simple docstring'''






        super().__init__()

        _lowercase	       =nn.Embedding(lowerCAmelCase							,   lowerCAmelCase					)

        _lowercase	       =nn.Embedding(lowerCAmelCase							,   lowerCAmelCase					)
        _lowercase	       =False

        _lowercase	       =nn.Dropout(p=lowerCAmelCase					)

        _lowercase	       =TaConfig(
            vocab_size=lowerCAmelCase							,   d_model=lowerCAmelCase							,   num_heads=lowerCAmelCase							,   d_kv=lowerCAmelCase							,   d_ff=lowerCAmelCase							,   dropout_rate=lowerCAmelCase							,   feed_forward_proj=lowerCAmelCase							,   is_decoder=lowerCAmelCase							,   is_encoder_decoder=lowerCAmelCase							,   )

        _lowercase	       =nn.ModuleList()
        for lyr_num in range(lowerCAmelCase					):
              _lowercase	       =TaBlock(lowerCAmelCase					)
              self.encoders.append(lowerCAmelCase					)

        _lowercase	       =TaLayerNorm(lowerCAmelCase					)
        _lowercase	       =nn.Dropout(p=lowerCAmelCase					)




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase					)  -> Dict:
        '''simple docstring'''






        _lowercase	       =self.token_embedder(lowerCAmelCase					)

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

        _lowercase	       =self.dropout_pre(lowerCAmelCase					)

        # inverted the attention mask
        _lowercase	       =encoder_input_tokens.size()
        _lowercase	       =self.get_extended_attention_mask(lowerCAmelCase							,   lowerCAmelCase					)

        for lyr in self.encoders:
              _lowercase	       =lyr(lowerCAmelCase							,   lowerCAmelCase					)[0]
        _lowercase	       =self.layer_norm(lowerCAmelCase					)

        return self.dropout_post(lowerCAmelCase					), encoder_inputs_mask


