



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







	"""simple docstring"""






	@register_to_config
	def __init__(		self    :  int	,    __snake_case    :  int = 128	,    __snake_case    :  int = 256	,    __snake_case    :  float = 20_00.0	,    __snake_case    :  int = 768	,    __snake_case    :  int = 12	,    __snake_case    :  int = 12	,    __snake_case    :  int = 64	,    __snake_case    :  int = 2048	,    __snake_case    :  float = 0.1	,    )    ->	Optional[int]:
								super().__init__()

								UpperCAmelCase				: Any											=	nn.Sequential(
								    nn.Linear(__snake_case	,    d_model * 4	,    bias=__snake_case	)	,    nn.SiLU()	,    nn.Linear(d_model * 4	,    d_model * 4	,    bias=__snake_case	)	,    nn.SiLU()	,    )

								UpperCAmelCase				: Any											=	nn.Embedding(__snake_case	,    __snake_case	)
								UpperCAmelCase				: List[Any]											=	False

								UpperCAmelCase				: Union[str, Any]											=	nn.Linear(__snake_case	,    __snake_case	,    bias=__snake_case	)

								UpperCAmelCase				: Optional[Any]											=	nn.Dropout(p=__snake_case	)

								UpperCAmelCase				: List[Any]											=	nn.ModuleList()
								for lyr_num in range(__snake_case	):
															# FiLM conditional T5 decoder
															UpperCAmelCase				: str											=	DecoderLayer(d_model=__snake_case	,    d_kv=__snake_case	,    num_heads=__snake_case	,    d_ff=__snake_case	,    dropout_rate=__snake_case	)
															self.decoders.append(__snake_case	)

								UpperCAmelCase				: int											=	TaLayerNorm(__snake_case	)

								UpperCAmelCase				: Optional[int]											=	nn.Dropout(p=__snake_case	)
								UpperCAmelCase				: str											=	nn.Linear(__snake_case	,    __snake_case	,    bias=__snake_case	)







	def   A			(		self    :  Any	,    __snake_case    :  Optional[Any]	,    __snake_case    :  List[Any]	)    ->	Any:
								UpperCAmelCase				: Union[str, Any]											=	torch.mul(query_input.unsqueeze(-1	)	,    key_input.unsqueeze(-2	)	)
								return mask.unsqueeze(-3	)







	def   A			(		self    :  List[Any]	,    __snake_case    :  int	,    __snake_case    :  Optional[Any]	,    __snake_case    :  str	)    ->	Any:
								UpperCAmelCase							,	UpperCAmelCase							,	UpperCAmelCase				: Any											=	decoder_input_tokens.shape
								assert decoder_noise_time.shape == (batch,)

								# decoder_noise_time is in [0, 1), so rescale to expected timing range.
								UpperCAmelCase				: Any											=	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	)

								UpperCAmelCase				: Dict											=	self.conditioning_emb(__snake_case	).unsqueeze(1	)

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

								UpperCAmelCase				: int											=	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.
								UpperCAmelCase				: Tuple											=	torch.broadcast_to(
								    torch.arange(__snake_case	,    device=decoder_input_tokens.device	)	,    (batch, seq_length)	,    )

								UpperCAmelCase				: str											=	self.position_encoding(__snake_case	)

								UpperCAmelCase				: List[str]											=	self.continuous_inputs_projection(__snake_case	)
								inputs += position_encodings
								UpperCAmelCase				: Optional[Any]											=	self.dropout(__snake_case	)

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

								# Translate encoding masks to encoder-decoder masks.
								UpperCAmelCase				: Tuple											=	[(x, self.encoder_decoder_mask(__snake_case	,    __snake_case	)) for x, y in encodings_and_masks]

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

								for lyr in self.decoders:
															UpperCAmelCase				: Tuple											=	lyr(
															    __snake_case	,    conditioning_emb=__snake_case	,    encoder_hidden_states=__snake_case	,    encoder_attention_mask=__snake_case	,    )[0]

								UpperCAmelCase				: int											=	self.decoder_norm(__snake_case	)
								UpperCAmelCase				: Optional[Any]											=	self.post_dropout(__snake_case	)

								UpperCAmelCase				: str											=	self.spec_out(__snake_case	)
								return spec_out

class        SCREAMING_SNAKE_CASE(							nn.Module    ):







	"""simple docstring"""






	def __init__(		self    :  Any	,    __snake_case    :  Any	,    __snake_case    :  Optional[Any]	,    __snake_case    :  List[Any]	,    __snake_case    :  str	,    __snake_case    :  Any	,    __snake_case    :  Tuple=1E-6	)    ->	int:
								super().__init__()
								UpperCAmelCase				: Optional[int]											=	nn.ModuleList()

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

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

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







	def   A			(		self    :  Union[str, Any]	,    __snake_case    :  Dict	,    __snake_case    :  List[str]=None	,    __snake_case    :  Dict=None	,    __snake_case    :  Dict=None	,    __snake_case    :  List[str]=None	,    __snake_case    :  Optional[int]=None	,    )    ->	Tuple:
								UpperCAmelCase				: Dict											=	self.layer[0](
								    __snake_case	,    conditioning_emb=__snake_case	,    attention_mask=__snake_case	,    )

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

															UpperCAmelCase				: int											=	self.layer[1](
															    __snake_case	,    key_value_states=__snake_case	,    attention_mask=__snake_case	,    )

								# Apply Film Conditional Feed Forward layer
								UpperCAmelCase				: str											=	self.layer[-1](__snake_case	,    __snake_case	)

								return (hidden_states,)

class        SCREAMING_SNAKE_CASE(							nn.Module    ):







	"""simple docstring"""






	def __init__(		self    :  Optional[int]	,    __snake_case    :  Optional[int]	,    __snake_case    :  Tuple	,    __snake_case    :  Tuple	,    __snake_case    :  str	)    ->	Optional[int]:
								super().__init__()
								UpperCAmelCase				: Any											=	TaLayerNorm(__snake_case	)
								UpperCAmelCase				: Tuple											=	TaFiLMLayer(in_features=d_model * 4	,    out_features=__snake_case	)
								UpperCAmelCase				: List[str]											=	Attention(query_dim=__snake_case	,    heads=__snake_case	,    dim_head=__snake_case	,    out_bias=__snake_case	,    scale_qk=__snake_case	)
								UpperCAmelCase				: List[Any]											=	nn.Dropout(__snake_case	)







	def   A			(		self    :  List[Any]	,    __snake_case    :  Optional[Any]	,    __snake_case    :  Union[str, Any]=None	,    __snake_case    :  Dict=None	,    )    ->	Optional[Any]:
								# pre_self_attention_layer_norm
								UpperCAmelCase				: Optional[Any]											=	self.layer_norm(__snake_case	)

								if conditioning_emb is not None:
															UpperCAmelCase				: Any											=	self.FiLMLayer(__snake_case	,    __snake_case	)

								# Self-attention block
								UpperCAmelCase				: Optional[Any]											=	self.attention(__snake_case	)

								UpperCAmelCase				: Union[str, Any]											=	hidden_states + self.dropout(__snake_case	)

								return hidden_states

class        SCREAMING_SNAKE_CASE(							nn.Module    ):







	"""simple docstring"""






	def __init__(		self    :  Optional[int]	,    __snake_case    :  Union[str, Any]	,    __snake_case    :  Union[str, Any]	,    __snake_case    :  Any	,    __snake_case    :  List[Any]	,    __snake_case    :  List[str]	)    ->	str:
								super().__init__()
								UpperCAmelCase				: Dict											=	Attention(query_dim=__snake_case	,    heads=__snake_case	,    dim_head=__snake_case	,    out_bias=__snake_case	,    scale_qk=__snake_case	)
								UpperCAmelCase				: Dict											=	TaLayerNorm(__snake_case	,    eps=__snake_case	)
								UpperCAmelCase				: Dict											=	nn.Dropout(__snake_case	)







	def   A			(		self    :  Optional[Any]	,    __snake_case    :  str	,    __snake_case    :  str=None	,    __snake_case    :  Optional[Any]=None	,    )    ->	str:
								UpperCAmelCase				: str											=	self.layer_norm(__snake_case	)
								UpperCAmelCase				: Dict											=	self.attention(
								    __snake_case	,    encoder_hidden_states=__snake_case	,    attention_mask=attention_mask.squeeze(1	)	,    )
								UpperCAmelCase				: Dict											=	hidden_states + self.dropout(__snake_case	)
								return layer_output

class        SCREAMING_SNAKE_CASE(							nn.Module    ):







	"""simple docstring"""






	def __init__(		self    :  Any	,    __snake_case    :  Dict	,    __snake_case    :  Tuple	,    __snake_case    :  Tuple	,    __snake_case    :  Dict	)    ->	int:
								super().__init__()
								UpperCAmelCase				: List[Any]											=	TaDenseGatedActDense(d_model=__snake_case	,    d_ff=__snake_case	,    dropout_rate=__snake_case	)
								UpperCAmelCase				: Tuple											=	TaFiLMLayer(in_features=d_model * 4	,    out_features=__snake_case	)
								UpperCAmelCase				: Optional[Any]											=	TaLayerNorm(__snake_case	,    eps=__snake_case	)
								UpperCAmelCase				: Dict											=	nn.Dropout(__snake_case	)







	def   A			(		self    :  Union[str, Any]	,    __snake_case    :  Optional[int]	,    __snake_case    :  Dict=None	)    ->	Any:
								UpperCAmelCase				: str											=	self.layer_norm(__snake_case	)
								if conditioning_emb is not None:
															UpperCAmelCase				: Optional[Any]											=	self.film(__snake_case	,    __snake_case	)

								UpperCAmelCase				: Optional[Any]											=	self.DenseReluDense(__snake_case	)
								UpperCAmelCase				: List[Any]											=	hidden_states + self.dropout(__snake_case	)
								return hidden_states

class        SCREAMING_SNAKE_CASE(							nn.Module    ):







	"""simple docstring"""






	def __init__(		self    :  Dict	,    __snake_case    :  List[str]	,    __snake_case    :  str	,    __snake_case    :  List[str]	)    ->	int:
								super().__init__()
								UpperCAmelCase				: Dict											=	nn.Linear(__snake_case	,    __snake_case	,    bias=__snake_case	)
								UpperCAmelCase				: Optional[int]											=	nn.Linear(__snake_case	,    __snake_case	,    bias=__snake_case	)
								UpperCAmelCase				: List[str]											=	nn.Linear(__snake_case	,    __snake_case	,    bias=__snake_case	)
								UpperCAmelCase				: Optional[Any]											=	nn.Dropout(__snake_case	)
								UpperCAmelCase				: Optional[int]											=	NewGELUActivation()







	def   A			(		self    :  Union[str, Any]	,    __snake_case    :  Optional[Any]	)    ->	List[Any]:
								UpperCAmelCase				: int											=	self.act(self.wi_a(__snake_case	)	)
								UpperCAmelCase				: List[Any]											=	self.wi_a(__snake_case	)
								UpperCAmelCase				: Tuple											=	hidden_gelu * hidden_linear
								UpperCAmelCase				: int											=	self.dropout(__snake_case	)

								UpperCAmelCase				: Optional[int]											=	self.wo(__snake_case	)
								return hidden_states

class        SCREAMING_SNAKE_CASE(							nn.Module    ):







	"""simple docstring"""






	def __init__(		self    :  Optional[int]	,    __snake_case    :  Optional[Any]	,    __snake_case    :  Union[str, Any]=1E-6	)    ->	Optional[Any]:
								super().__init__()
								UpperCAmelCase				: Union[str, Any]											=	nn.Parameter(torch.ones(__snake_case	)	)
								UpperCAmelCase				: int											=	eps







	def   A			(		self    :  int	,    __snake_case    :  Dict	)    ->	Union[str, Any]:
								# T5 uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean
								# Square Layer Normalization https://arxiv.org/abs/1910.07467 thus variance is calculated
								# w/o mean and there is no bias. Additionally we want to make sure that the accumulation for
								# half-precision inputs is done in fp32

								UpperCAmelCase				: List[Any]											=	hidden_states.to(torch.floataa	).pow(2	).mean(-1	,    keepdim=__snake_case	)
								UpperCAmelCase				: Optional[Any]											=	hidden_states * torch.rsqrt(variance + self.variance_epsilon	)

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

								return self.weight * hidden_states

class        SCREAMING_SNAKE_CASE(							nn.Module    ):







	"""simple docstring"""






	def   A			(		self    :  Any	,    __snake_case    :  torch.Tensor	)    ->	torch.Tensor:
								return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi	) * (input + 0.04_47_15 * torch.pow(__snake_case	,    3.0	))	))

class        SCREAMING_SNAKE_CASE(							nn.Module    ):







	"""simple docstring"""






	def __init__(		self    :  Optional[Any]	,    __snake_case    :  str	,    __snake_case    :  Union[str, Any]	)    ->	int:
								super().__init__()
								UpperCAmelCase				: List[str]											=	nn.Linear(__snake_case	,    out_features * 2	,    bias=__snake_case	)







	def   A			(		self    :  Union[str, Any]	,    __snake_case    :  Union[str, Any]	,    __snake_case    :  int	)    ->	Dict:
								UpperCAmelCase				: List[str]											=	self.scale_bias(__snake_case	)
								UpperCAmelCase							,	UpperCAmelCase				: Any											=	torch.chunk(__snake_case	,    2	,    -1	)
								UpperCAmelCase				: List[Any]											=	x * (1 + scale) + shift
								return x










'''simple docstring'''






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


UpperCamelCase__:     Dict										=    logging.get_logger(__name__)

UpperCamelCase__:     int										=    {
    # See all MEGATRON_BERT models at https://huggingface.co/models?filter=bert
}

class        SCREAMING_SNAKE_CASE(							A__    ):







	"""simple docstring"""






	lowerCamelCase__    				= """megatron-bert"""







	def __init__(		self    :  int	,    __snake_case    :  List[Any]=29056	,    __snake_case    :  Dict=1024	,    __snake_case    :  Tuple=24	,    __snake_case    :  Any=16	,    __snake_case    :  Tuple=4096	,    __snake_case    :  int="gelu"	,    __snake_case    :  Dict=0.1	,    __snake_case    :  str=0.1	,    __snake_case    :  Optional[Any]=512	,    __snake_case    :  Tuple=2	,    __snake_case    :  Tuple=0.02	,    __snake_case    :  int=1E-12	,    __snake_case    :  Union[str, Any]=0	,    __snake_case    :  int="absolute"	,    __snake_case    :  str=True	,    **__snake_case    :  Tuple	,    )    ->	Union[str, Any]:
								super().__init__(pad_token_id=__snake_case	,    **__snake_case	)

								UpperCAmelCase				: Any											=	vocab_size
								UpperCAmelCase				: Tuple											=	hidden_size
								UpperCAmelCase				: List[Any]											=	num_hidden_layers
								UpperCAmelCase				: str											=	num_attention_heads
								UpperCAmelCase				: Optional[int]											=	hidden_act
								UpperCAmelCase				: Optional[int]											=	intermediate_size
								UpperCAmelCase				: Dict											=	hidden_dropout_prob
								UpperCAmelCase				: List[Any]											=	attention_probs_dropout_prob
								UpperCAmelCase				: Dict											=	max_position_embeddings
								UpperCAmelCase				: Any											=	type_vocab_size
								UpperCAmelCase				: str											=	initializer_range
								UpperCAmelCase				: str											=	layer_norm_eps
								UpperCAmelCase				: Optional[Any]											=	position_embedding_type
								UpperCAmelCase				: str											=	use_cache





