






'''simple docstring'''





from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


__A							:							int  							=		{"""configuration_mra""": ["""MRA_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MraConfig"""]}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __A							:							Optional[Any]  							=		[
     """MRA_PRETRAINED_MODEL_ARCHIVE_LIST""",
     """MraForMaskedLM""",
     """MraForMultipleChoice""",
     """MraForQuestionAnswering""",
     """MraForSequenceClassification""",
     """MraForTokenClassification""",
     """MraLayer""",
     """MraModel""",
     """MraPreTrainedModel""",
 ]


if TYPE_CHECKING:
 from .configuration_mra import MRA_PRETRAINED_CONFIG_ARCHIVE_MAP, MraConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_mra import (
      MRA_PRETRAINED_MODEL_ARCHIVE_LIST,
      MraForMaskedLM,
      MraForMultipleChoice,
      MraForQuestionAnswering,
      MraForSequenceClassification,
      MraForTokenClassification,
      MraLayer,
      MraModel,
      MraPreTrainedModel,
  )
else:
 import sys

 __A							:							Optional[int]  							=		_LazyModule(__name__, globals()["""__file__"""], _import_structure)









'''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

