






'''simple docstring'''





from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast
from ...utils import logging


__A							:							Union[str, Any]  							=		logging.get_logger(__name__)

__A							:							Tuple  							=		{
    """EleutherAI/gpt-neo-1.3B""": """https://huggingface.co/EleutherAI/gpt-neo-1.3B/resolve/main/config.json""",
    # See all GPTNeo models at https://huggingface.co/models?filter=gpt_neo
}
class 				lowercase       (							_lowerCamelCase      ):


    '''simple docstring'''

    lowerCAmelCase__      				=							"gpt_neo"
    lowerCAmelCase__      				=							["past_key_values"]
    lowerCAmelCase__      				=							{"num_attention_heads": "num_heads", "num_hidden_layers": "num_layers"}

    def __init__(					self     :     str				,     __lowerCamelCase     :     Optional[int]=50257				,     __lowerCamelCase     :     Optional[int]=2048				,     __lowerCamelCase     :     Optional[int]=2048				,     __lowerCamelCase     :     Any=24				,     __lowerCamelCase     :     List[Any]=[[["global", "local"], 12]]				,     __lowerCamelCase     :     Any=16				,     __lowerCamelCase     :     Any=None				,     __lowerCamelCase     :     Optional[Any]=256				,     __lowerCamelCase     :     List[Any]="gelu_new"				,     __lowerCamelCase     :     Any=0.0				,     __lowerCamelCase     :     List[Any]=0.0				,     __lowerCamelCase     :     Union[str, Any]=0.0				,     __lowerCamelCase     :     Any=0.1				,     __lowerCamelCase     :     Dict=1E-5				,     __lowerCamelCase     :     List[Any]=0.0_2				,     __lowerCamelCase     :     List[Any]=True				,     __lowerCamelCase     :     Optional[int]=50256				,     __lowerCamelCase     :     Tuple=50256				,     **__lowerCamelCase     :     Optional[int]				,     )							->							str:

     '''simple docstring'''




     lowerCamelCase__  =     vocab_size
     lowerCamelCase__  =     max_position_embeddings
     lowerCamelCase__  =     hidden_size
     lowerCamelCase__  =     num_layers
     lowerCamelCase__  =     num_heads
     lowerCamelCase__  =     intermediate_size
     lowerCamelCase__  =     window_size
     lowerCamelCase__  =     activation_function
     lowerCamelCase__  =     resid_dropout
     lowerCamelCase__  =     embed_dropout
     lowerCamelCase__  =     attention_dropout
     lowerCamelCase__  =     classifier_dropout
     lowerCamelCase__  =     layer_norm_epsilon
     lowerCamelCase__  =     initializer_range
     lowerCamelCase__  =     use_cache

     lowerCamelCase__  =     bos_token_id
     lowerCamelCase__  =     eos_token_id

     lowerCamelCase__  =     attention_types
     lowerCamelCase__  =     self.expand_attention_types_params(__lowerCamelCase     )

     if len(self.attention_layers     ) != self.num_layers:
      raise ValueError(
          "Configuration for convolutional module is incorrect. "
          "It is required that `len(config.attention_layers)` == `config.num_layers` "
          f'''but is `len(config.attention_layers) = {len(self.attention_layers     )}`, '''
          f'''`config.num_layers = {self.num_layers}`. '''
          "`config.attention_layers` is prepared using `config.attention_types`. "
          "Please verify the value of `config.attention_types` argument."     )

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




    @staticmethod
    def    a__				(					__lowerCamelCase     :     int     )							->							str:

     '''simple docstring'''




     lowerCamelCase__  =     []
     for item in attention_types:
      for _ in range(item[1]     ):
       attentions.extend(item[0]     )
     return attentions






def    lowerCamelCase_     (    lowercase__				,     lowercase__				,     lowercase__				,     lowercase__):
 import torch

 lowerCamelCase__  =     input.size()
 lowerCamelCase__  =     len(lowercase__)
 lowerCamelCase__  =     shape[dimension]

 lowerCamelCase__  =     torch.arange(0				,     lowercase__				,     lowercase__)
 lowerCamelCase__  =     torch.div(sizedim - size				,     lowercase__				,     rounding_mode="floor") + 1
 lowerCamelCase__  =     torch.arange(lowercase__) + low_indices[:min_length][:, None]

 lowerCamelCase__  =     [slice(lowercase__)] * rank
 lowerCamelCase__  =     indices
 lowerCamelCase__  =     input[s]

 lowerCamelCase__  =     list(range(0				,     rank + 1))
 perm.append(perm.pop(dimension + 1))

 return sliced.permute(lowercase__)






def    lowerCamelCase_     (    lowercase__				,     lowercase__):
 import torch

 lowerCamelCase__  =     torch.arange(1				,     lowercase__)
 lowerCamelCase__  =     torch.remainder(lowercase__				,     lowercase__)
 lowerCamelCase__  =     remainders == 0
 lowerCamelCase__  =     candidates[divisor_indices]
 lowerCamelCase__  =     torch.max(lowercase__)
 return largest_divisor, torch.div(lowercase__				,     lowercase__				,     rounding_mode="floor")
class 				lowercase       (							_lowerCamelCase      ):


    '''simple docstring'''

    @property
    def    a__				(					self     :     Any     )							->							Mapping[str, Mapping[int, str]]:

     '''simple docstring'''




     lowerCamelCase__  =     OrderedDict({"input_ids": {0: "batch", 1: "sequence"}}     )
     if self.use_past:
      self.fill_with_past_key_values_(__lowerCamelCase				,     direction="inputs"     )
      lowerCamelCase__  =     {0: "batch", 1: "past_sequence + sequence"}
     else:
      lowerCamelCase__  =     {0: "batch", 1: "sequence"}

     return common_inputs

    @property
    def    a__				(					self     :     Any     )							->							int:

     '''simple docstring'''




     return self._config.num_heads

    def    a__				(					self     :     Optional[Any]				,     __lowerCamelCase     :     PreTrainedTokenizer				,     __lowerCamelCase     :     int = -1				,     __lowerCamelCase     :     int = -1				,     __lowerCamelCase     :     bool = False				,     __lowerCamelCase     :     Optional[TensorType] = None				,     )							->							Mapping[str, Any]:

     '''simple docstring'''




     lowerCamelCase__  =     super(__lowerCamelCase				,     self     ).generate_dummy_inputs(
         __lowerCamelCase				,     batch_size=__lowerCamelCase				,     seq_length=__lowerCamelCase				,     is_pair=__lowerCamelCase				,     framework=__lowerCamelCase     )

     # We need to order the input in the way they appears in the forward()
     lowerCamelCase__  =     OrderedDict({"input_ids": common_inputs["input_ids"]}     )

     # Need to add the past_keys
     if self.use_past:
      if not is_torch_available():
       raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."     )
      else:
       import torch

       lowerCamelCase__  ,	lowerCamelCase__  =     common_inputs["input_ids"].shape
       # Not using the same length for past_key_values
       lowerCamelCase__  =     seqlen + 2
       lowerCamelCase__  =     (
           batch,
           self.num_attention_heads,
           past_key_values_length,
           self._config.hidden_size // self.num_attention_heads,
       )
       lowerCamelCase__  =     [
           (torch.zeros(__lowerCamelCase     ), torch.zeros(__lowerCamelCase     )) for _ in range(self.num_layers     )
       ]

     lowerCamelCase__  =     common_inputs["attention_mask"]
     if self.use_past:
      lowerCamelCase__  =     ordered_inputs["attention_mask"].dtype
      lowerCamelCase__  =     torch.cat(
          [ordered_inputs["attention_mask"], torch.ones(__lowerCamelCase				,     __lowerCamelCase				,     dtype=__lowerCamelCase     )]				,     dim=1     )

     return ordered_inputs




    @property
    def    a__				(					self     :     Optional[Any]     )							->							int:

     '''simple docstring'''




     return 13









'''simple docstring'''





import math

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from .attention_processor import Attention
from .embeddings import get_timestep_embedding
from .modeling_utils import ModelMixin
class 				lowercase       (							_lowerCamelCase ,       _lowerCamelCase      ):


    '''simple docstring'''

    @register_to_config
    def __init__(					self     :     Optional[Any]				,     __lowerCamelCase     :     int = 128				,     __lowerCamelCase     :     int = 256				,     __lowerCamelCase     :     float = 2_0_0_0.0				,     __lowerCamelCase     :     int = 768				,     __lowerCamelCase     :     int = 12				,     __lowerCamelCase     :     int = 12				,     __lowerCamelCase     :     int = 64				,     __lowerCamelCase     :     int = 2048				,     __lowerCamelCase     :     float = 0.1				,     )							->							str:

     '''simple docstring'''




     super().__init__()

     lowerCamelCase__  =     nn.Sequential(
         nn.Linear(__lowerCamelCase				,     d_model * 4				,     bias=__lowerCamelCase     )				,     nn.SiLU()				,     nn.Linear(d_model * 4				,     d_model * 4				,     bias=__lowerCamelCase     )				,     nn.SiLU()				,     )

     lowerCamelCase__  =     nn.Embedding(__lowerCamelCase				,     __lowerCamelCase     )
     lowerCamelCase__  =     False

     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     __lowerCamelCase				,     bias=__lowerCamelCase     )

     lowerCamelCase__  =     nn.Dropout(p=__lowerCamelCase     )

     lowerCamelCase__  =     nn.ModuleList()
     for lyr_num in range(__lowerCamelCase     ):
      # FiLM conditional T5 decoder
      lowerCamelCase__  =     DecoderLayer(d_model=__lowerCamelCase				,     d_kv=__lowerCamelCase				,     num_heads=__lowerCamelCase				,     d_ff=__lowerCamelCase				,     dropout_rate=__lowerCamelCase     )
      self.decoders.append(__lowerCamelCase     )

     lowerCamelCase__  =     TaLayerNorm(__lowerCamelCase     )

     lowerCamelCase__  =     nn.Dropout(p=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     __lowerCamelCase				,     bias=__lowerCamelCase     )

    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     Optional[int]     )							->							Any:

     '''simple docstring'''




     lowerCamelCase__  =     torch.mul(query_input.unsqueeze(-1     )				,     key_input.unsqueeze(-2     )     )
     return mask.unsqueeze(-3     )




    def    a__				(					self     :     Tuple				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     Union[str, Any]     )							->							Optional[int]:

     '''simple docstring'''




     lowerCamelCase__  ,	lowerCamelCase__  ,	lowerCamelCase__  =     decoder_input_tokens.shape
     assert decoder_noise_time.shape == (batch,)

     # decoder_noise_time is in [0, 1), so rescale to expected timing range.
     lowerCamelCase__  =     get_timestep_embedding(
         decoder_noise_time * self.config.max_decoder_noise_time				,     embedding_dim=self.config.d_model				,     max_period=self.config.max_decoder_noise_time				,     ).to(dtype=self.dtype     )

     lowerCamelCase__  =     self.conditioning_emb(__lowerCamelCase     ).unsqueeze(1     )

     assert conditioning_emb.shape == (batch, 1, self.config.d_model * 4)

     lowerCamelCase__  =     decoder_input_tokens.shape[1]

     # If we want to use relative positions for audio context, we can just offset
     # this sequence by the length of encodings_and_masks.
     lowerCamelCase__  =     torch.broadcast_to(
         torch.arange(__lowerCamelCase				,     device=decoder_input_tokens.device     )				,     (batch, seq_length)				,     )

     lowerCamelCase__  =     self.position_encoding(__lowerCamelCase     )

     lowerCamelCase__  =     self.continuous_inputs_projection(__lowerCamelCase     )
     inputs += position_encodings
     lowerCamelCase__  =     self.dropout(__lowerCamelCase     )

     # decoder: No padding present.
     lowerCamelCase__  =     torch.ones(
         decoder_input_tokens.shape[:2]				,     device=decoder_input_tokens.device				,     dtype=inputs.dtype     )

     # Translate encoding masks to encoder-decoder masks.
     lowerCamelCase__  =     [(x, self.encoder_decoder_mask(__lowerCamelCase				,     __lowerCamelCase     )) for x, y in encodings_and_masks]

     # cross attend style: concat encodings
     lowerCamelCase__  =     torch.cat([x[0] for x in encodings_and_encdec_masks]				,     dim=1     )
     lowerCamelCase__  =     torch.cat([x[1] for x in encodings_and_encdec_masks]				,     dim=-1     )

     for lyr in self.decoders:
      lowerCamelCase__  =     lyr(
          __lowerCamelCase				,     conditioning_emb=__lowerCamelCase				,     encoder_hidden_states=__lowerCamelCase				,     encoder_attention_mask=__lowerCamelCase				,     )[0]

     lowerCamelCase__  =     self.decoder_norm(__lowerCamelCase     )
     lowerCamelCase__  =     self.post_dropout(__lowerCamelCase     )

     lowerCamelCase__  =     self.spec_out(__lowerCamelCase     )
     return spec_out
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     Union[str, Any]				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     str=1E-6     )							->							int:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     nn.ModuleList()

     # cond self attention: layer 0
     self.layer.append(
         TaLayerSelfAttentionCond(d_model=__lowerCamelCase				,     d_kv=__lowerCamelCase				,     num_heads=__lowerCamelCase				,     dropout_rate=__lowerCamelCase     )     )

     # cross attention: layer 1
     self.layer.append(
         TaLayerCrossAttention(
             d_model=__lowerCamelCase				,     d_kv=__lowerCamelCase				,     num_heads=__lowerCamelCase				,     dropout_rate=__lowerCamelCase				,     layer_norm_epsilon=__lowerCamelCase				,     )     )

     # Film Cond MLP + dropout: last layer
     self.layer.append(
         TaLayerFFCond(d_model=__lowerCamelCase				,     d_ff=__lowerCamelCase				,     dropout_rate=__lowerCamelCase				,     layer_norm_epsilon=__lowerCamelCase     )     )




    def    a__				(					self     :     List[Any]				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     Optional[Any]=None				,     __lowerCamelCase     :     Optional[int]=None				,     __lowerCamelCase     :     int=None				,     __lowerCamelCase     :     List[Any]=None				,     __lowerCamelCase     :     List[str]=None				,     )							->							Dict:

     '''simple docstring'''




     lowerCamelCase__  =     self.layer[0](
         __lowerCamelCase				,     conditioning_emb=__lowerCamelCase				,     attention_mask=__lowerCamelCase				,     )

     if encoder_hidden_states is not None:
      lowerCamelCase__  =     torch.where(encoder_attention_mask > 0				,     0				,     -1E10     ).to(
          encoder_hidden_states.dtype     )

      lowerCamelCase__  =     self.layer[1](
          __lowerCamelCase				,     key_value_states=__lowerCamelCase				,     attention_mask=__lowerCamelCase				,     )

     # Apply Film Conditional Feed Forward layer
     lowerCamelCase__  =     self.layer[-1](__lowerCamelCase				,     __lowerCamelCase     )

     return (hidden_states,)
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     List[Any]				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     Tuple     )							->							Optional[int]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     TaLayerNorm(__lowerCamelCase     )
     lowerCamelCase__  =     TaFiLMLayer(in_features=d_model * 4				,     out_features=__lowerCamelCase     )
     lowerCamelCase__  =     Attention(query_dim=__lowerCamelCase				,     heads=__lowerCamelCase				,     dim_head=__lowerCamelCase				,     out_bias=__lowerCamelCase				,     scale_qk=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Dropout(__lowerCamelCase     )




    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     Any				,     __lowerCamelCase     :     Optional[Any]=None				,     __lowerCamelCase     :     str=None				,     )							->							List[Any]:

     '''simple docstring'''




     lowerCamelCase__  =     self.layer_norm(__lowerCamelCase     )

     if conditioning_emb is not None:
      lowerCamelCase__  =     self.FiLMLayer(__lowerCamelCase				,     __lowerCamelCase     )

     # Self-attention block
     lowerCamelCase__  =     self.attention(__lowerCamelCase     )

     lowerCamelCase__  =     hidden_states + self.dropout(__lowerCamelCase     )

     return hidden_states
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     Any				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     Any				,     __lowerCamelCase     :     Dict     )							->							Union[str, Any]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     Attention(query_dim=__lowerCamelCase				,     heads=__lowerCamelCase				,     dim_head=__lowerCamelCase				,     out_bias=__lowerCamelCase				,     scale_qk=__lowerCamelCase     )
     lowerCamelCase__  =     TaLayerNorm(__lowerCamelCase				,     eps=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Dropout(__lowerCamelCase     )




    def    a__				(					self     :     List[str]				,     __lowerCamelCase     :     Any				,     __lowerCamelCase     :     Optional[Any]=None				,     __lowerCamelCase     :     Dict=None				,     )							->							Any:

     '''simple docstring'''




     lowerCamelCase__  =     self.layer_norm(__lowerCamelCase     )
     lowerCamelCase__  =     self.attention(
         __lowerCamelCase				,     encoder_hidden_states=__lowerCamelCase				,     attention_mask=attention_mask.squeeze(1     )				,     )
     lowerCamelCase__  =     hidden_states + self.dropout(__lowerCamelCase     )
     return layer_output
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     List[str]				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     List[str]     )							->							Optional[Any]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     TaDenseGatedActDense(d_model=__lowerCamelCase				,     d_ff=__lowerCamelCase				,     dropout_rate=__lowerCamelCase     )
     lowerCamelCase__  =     TaFiLMLayer(in_features=d_model * 4				,     out_features=__lowerCamelCase     )
     lowerCamelCase__  =     TaLayerNorm(__lowerCamelCase				,     eps=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Dropout(__lowerCamelCase     )




    def    a__				(					self     :     Any				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     Union[str, Any]=None     )							->							Dict:

     '''simple docstring'''




     lowerCamelCase__  =     self.layer_norm(__lowerCamelCase     )
     if conditioning_emb is not None:
      lowerCamelCase__  =     self.film(__lowerCamelCase				,     __lowerCamelCase     )

     lowerCamelCase__  =     self.DenseReluDense(__lowerCamelCase     )
     lowerCamelCase__  =     hidden_states + self.dropout(__lowerCamelCase     )
     return hidden_states
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     List[str]				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     str     )							->							str:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     __lowerCamelCase				,     bias=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     __lowerCamelCase				,     bias=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     __lowerCamelCase				,     bias=__lowerCamelCase     )
     lowerCamelCase__  =     nn.Dropout(__lowerCamelCase     )
     lowerCamelCase__  =     NewGELUActivation()




    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     Union[str, Any]     )							->							Optional[Any]:

     '''simple docstring'''




     lowerCamelCase__  =     self.act(self.wi_a(__lowerCamelCase     )     )
     lowerCamelCase__  =     self.wi_a(__lowerCamelCase     )
     lowerCamelCase__  =     hidden_gelu * hidden_linear
     lowerCamelCase__  =     self.dropout(__lowerCamelCase     )

     lowerCamelCase__  =     self.wo(__lowerCamelCase     )
     return hidden_states
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     str				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     List[Any]=1E-6     )							->							Union[str, Any]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     nn.Parameter(torch.ones(__lowerCamelCase     )     )
     lowerCamelCase__  =     eps




    def    a__				(					self     :     Optional[int]				,     __lowerCamelCase     :     str     )							->							int:

     '''simple docstring'''




     lowerCamelCase__  =     hidden_states.to(torch.floataa     ).pow(2     ).mean(-1				,     keepdim=__lowerCamelCase     )
     lowerCamelCase__  =     hidden_states * torch.rsqrt(variance + self.variance_epsilon     )

     # convert into half-precision if necessary
     if self.weight.dtype in [torch.floataa, torch.bfloataa]:
      lowerCamelCase__  =     hidden_states.to(self.weight.dtype     )

     return self.weight * hidden_states
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def    a__				(					self     :     Tuple				,     __lowerCamelCase     :     torch.Tensor     )							->							torch.Tensor:

     '''simple docstring'''




     return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi     ) * (input + 0.0_4_4_7_1_5 * torch.pow(__lowerCamelCase				,     3.0     ))     ))
class 				lowercase       (							nn.Module      ):


    '''simple docstring'''

    def __init__(					self     :     Optional[int]				,     __lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     List[Any]     )							->							Optional[int]:

     '''simple docstring'''




     super().__init__()
     lowerCamelCase__  =     nn.Linear(__lowerCamelCase				,     out_features * 2				,     bias=__lowerCamelCase     )




    def    a__				(					self     :     str				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     str     )							->							List[str]:

     '''simple docstring'''




     lowerCamelCase__  =     self.scale_bias(__lowerCamelCase     )
     lowerCamelCase__  ,	lowerCamelCase__  =     torch.chunk(__lowerCamelCase				,     2				,     -1     )
     lowerCamelCase__  =     x * (1 + scale) + shift
     return x

