



'''simple docstring'''




import unittest

from transformers import AlbertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import (
						    MODEL_FOR_PRETRAINING_MAPPING,
						    AlbertForMaskedLM,
						    AlbertForMultipleChoice,
						    AlbertForPreTraining,
						    AlbertForQuestionAnswering,
						    AlbertForSequenceClassification,
						    AlbertForTokenClassification,
						    AlbertModel,
						)
						from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST






class 		_UpperCamelCase						:


			'''simple docstring'''

			def __init__(							self  :				List[str]							,  lowerCAmelCase__  :				int							,  lowerCAmelCase__  :				List[Any]=1_3							,  lowerCAmelCase__  :				Tuple=7							,  lowerCAmelCase__  :				Tuple=True							,  lowerCAmelCase__  :				Dict=True							,  lowerCAmelCase__  :				Optional[int]=True							,  lowerCAmelCase__  :				Any=True							,  lowerCAmelCase__  :				Any=9_9							,  lowerCAmelCase__  :				Union[str, Any]=1_6							,  lowerCAmelCase__  :				List[Any]=3_6							,  lowerCAmelCase__  :				Union[str, Any]=6							,  lowerCAmelCase__  :				Any=6							,  lowerCAmelCase__  :				Any=6							,  lowerCAmelCase__  :				Union[str, Any]=3_7							,  lowerCAmelCase__  :				List[str]="gelu"							,  lowerCAmelCase__  :				List[Any]=0.1							,  lowerCAmelCase__  :				List[Any]=0.1							,  lowerCAmelCase__  :				Any=5_1_2							,  lowerCAmelCase__  :				Dict=1_6							,  lowerCAmelCase__  :				int=2							,  lowerCAmelCase__  :				Tuple=0.02							,  lowerCAmelCase__  :				Union[str, Any]=3							,  lowerCAmelCase__  :				int=4							,  lowerCAmelCase__  :				List[Any]=None							,  ):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Tuple					 =	parent
					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	batch_size
					__SCREAMING_SNAKE_CASE :	str					 =	seq_length
					__SCREAMING_SNAKE_CASE :	Dict					 =	is_training
					__SCREAMING_SNAKE_CASE :	Dict					 =	use_input_mask
					__SCREAMING_SNAKE_CASE :	str					 =	use_token_type_ids
					__SCREAMING_SNAKE_CASE :	Dict					 =	use_labels
					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	vocab_size
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	embedding_size
					__SCREAMING_SNAKE_CASE :	Dict					 =	hidden_size
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	num_hidden_layers
					__SCREAMING_SNAKE_CASE :	str					 =	num_hidden_groups
					__SCREAMING_SNAKE_CASE :	str					 =	num_attention_heads
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	intermediate_size
					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	hidden_act
					__SCREAMING_SNAKE_CASE :	List[str]					 =	hidden_dropout_prob
					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	attention_probs_dropout_prob
					__SCREAMING_SNAKE_CASE :	str					 =	max_position_embeddings
					__SCREAMING_SNAKE_CASE :	List[str]					 =	type_vocab_size
					__SCREAMING_SNAKE_CASE :	int					 =	type_sequence_label_size
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	initializer_range
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	num_labels
					__SCREAMING_SNAKE_CASE :	str					 =	num_choices
					__SCREAMING_SNAKE_CASE :	str					 =	scope

			def UpperCamelCase__      (							self  :				List[str]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	int					 =	ids_tensor([self.batch_size, self.seq_length]							,  self.vocab_size		)

					__SCREAMING_SNAKE_CASE :	List[Any]					 =	None
					if self.use_input_mask:
							__SCREAMING_SNAKE_CASE :	Tuple					 =	random_attention_mask([self.batch_size, self.seq_length]		)

					__SCREAMING_SNAKE_CASE :	List[str]					 =	None
					if self.use_token_type_ids:
							__SCREAMING_SNAKE_CASE :	Tuple					 =	ids_tensor([self.batch_size, self.seq_length]							,  self.type_vocab_size		)

					__SCREAMING_SNAKE_CASE :	List[str]					 =	None
					__SCREAMING_SNAKE_CASE :	Any					 =	None
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	None
					if self.use_labels:
							__SCREAMING_SNAKE_CASE :	int					 =	ids_tensor([self.batch_size]							,  self.type_sequence_label_size		)
							__SCREAMING_SNAKE_CASE :	Tuple					 =	ids_tensor([self.batch_size, self.seq_length]							,  self.num_labels		)
							__SCREAMING_SNAKE_CASE :	str					 =	ids_tensor([self.batch_size]							,  self.num_choices		)

					__SCREAMING_SNAKE_CASE :	Any					 =	self.get_config()

					return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

			def UpperCamelCase__      (							self  :				Optional[int]		):





					"""simple docstring"""






					return AlbertConfig(
					    vocab_size=self.vocab_size							,  hidden_size=self.hidden_size							,  num_hidden_layers=self.num_hidden_layers							,  num_attention_heads=self.num_attention_heads							,  intermediate_size=self.intermediate_size							,  hidden_act=self.hidden_act							,  hidden_dropout_prob=self.hidden_dropout_prob							,  attention_probs_dropout_prob=self.attention_probs_dropout_prob							,  max_position_embeddings=self.max_position_embeddings							,  type_vocab_size=self.type_vocab_size							,  initializer_range=self.initializer_range							,  num_hidden_groups=self.num_hidden_groups							,  )

			def UpperCamelCase__      (							self  :				str							,  lowerCAmelCase__  :				int							,  lowerCAmelCase__  :				Union[str, Any]							,  lowerCAmelCase__  :				Optional[int]							,  lowerCAmelCase__  :				Tuple							,  lowerCAmelCase__  :				str							,  lowerCAmelCase__  :				Dict							,  lowerCAmelCase__  :				Any		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	int					 =	AlbertModel(config=lowerCAmelCase__		)
					model.to(lowerCAmelCase__		)
					model.eval()
					__SCREAMING_SNAKE_CASE :	Any					 =	model(lowerCAmelCase__							,  attention_mask=lowerCAmelCase__							,  token_type_ids=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Dict					 =	model(lowerCAmelCase__							,  token_type_ids=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	model(lowerCAmelCase__		)
					self.parent.assertEqual(result.last_hidden_state.shape							,  (self.batch_size, self.seq_length, self.hidden_size)		)
					self.parent.assertEqual(result.pooler_output.shape							,  (self.batch_size, self.hidden_size)		)

			def UpperCamelCase__      (							self  :				List[Any]							,  lowerCAmelCase__  :				Optional[Any]							,  lowerCAmelCase__  :				List[str]							,  lowerCAmelCase__  :				int							,  lowerCAmelCase__  :				int							,  lowerCAmelCase__  :				Optional[Any]							,  lowerCAmelCase__  :				Dict							,  lowerCAmelCase__  :				List[str]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	str					 =	AlbertForPreTraining(config=lowerCAmelCase__		)
					model.to(lowerCAmelCase__		)
					model.eval()
					__SCREAMING_SNAKE_CASE :	Dict					 =	model(
					    lowerCAmelCase__							,  attention_mask=lowerCAmelCase__							,  token_type_ids=lowerCAmelCase__							,  labels=lowerCAmelCase__							,  sentence_order_label=lowerCAmelCase__							,  )
					self.parent.assertEqual(result.prediction_logits.shape							,  (self.batch_size, self.seq_length, self.vocab_size)		)
					self.parent.assertEqual(result.sop_logits.shape							,  (self.batch_size, config.num_labels)		)

			def UpperCamelCase__      (							self  :				Optional[Any]							,  lowerCAmelCase__  :				List[str]							,  lowerCAmelCase__  :				Union[str, Any]							,  lowerCAmelCase__  :				Any							,  lowerCAmelCase__  :				List[Any]							,  lowerCAmelCase__  :				str							,  lowerCAmelCase__  :				Optional[int]							,  lowerCAmelCase__  :				Tuple		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	List[Any]					 =	AlbertForMaskedLM(config=lowerCAmelCase__		)
					model.to(lowerCAmelCase__		)
					model.eval()
					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	model(lowerCAmelCase__							,  attention_mask=lowerCAmelCase__							,  token_type_ids=lowerCAmelCase__							,  labels=lowerCAmelCase__		)
					self.parent.assertEqual(result.logits.shape							,  (self.batch_size, self.seq_length, self.vocab_size)		)

			def UpperCamelCase__      (							self  :				str							,  lowerCAmelCase__  :				List[str]							,  lowerCAmelCase__  :				List[str]							,  lowerCAmelCase__  :				Dict							,  lowerCAmelCase__  :				Optional[Any]							,  lowerCAmelCase__  :				Dict							,  lowerCAmelCase__  :				Union[str, Any]							,  lowerCAmelCase__  :				int		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	str					 =	AlbertForQuestionAnswering(config=lowerCAmelCase__		)
					model.to(lowerCAmelCase__		)
					model.eval()
					__SCREAMING_SNAKE_CASE :	Tuple					 =	model(
					    lowerCAmelCase__							,  attention_mask=lowerCAmelCase__							,  token_type_ids=lowerCAmelCase__							,  start_positions=lowerCAmelCase__							,  end_positions=lowerCAmelCase__							,  )
					self.parent.assertEqual(result.start_logits.shape							,  (self.batch_size, self.seq_length)		)
					self.parent.assertEqual(result.end_logits.shape							,  (self.batch_size, self.seq_length)		)

			def UpperCamelCase__      (							self  :				Union[str, Any]							,  lowerCAmelCase__  :				Dict							,  lowerCAmelCase__  :				Optional[int]							,  lowerCAmelCase__  :				List[str]							,  lowerCAmelCase__  :				Optional[Any]							,  lowerCAmelCase__  :				Union[str, Any]							,  lowerCAmelCase__  :				str							,  lowerCAmelCase__  :				Union[str, Any]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Dict					 =	self.num_labels
					__SCREAMING_SNAKE_CASE :	Tuple					 =	AlbertForSequenceClassification(lowerCAmelCase__		)
					model.to(lowerCAmelCase__		)
					model.eval()
					__SCREAMING_SNAKE_CASE :	Tuple					 =	model(lowerCAmelCase__							,  attention_mask=lowerCAmelCase__							,  token_type_ids=lowerCAmelCase__							,  labels=lowerCAmelCase__		)
					self.parent.assertEqual(result.logits.shape							,  (self.batch_size, self.num_labels)		)

			def UpperCamelCase__      (							self  :				Any							,  lowerCAmelCase__  :				int							,  lowerCAmelCase__  :				List[str]							,  lowerCAmelCase__  :				Optional[Any]							,  lowerCAmelCase__  :				Tuple							,  lowerCAmelCase__  :				Dict							,  lowerCAmelCase__  :				Optional[Any]							,  lowerCAmelCase__  :				List[str]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	str					 =	self.num_labels
					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	AlbertForTokenClassification(config=lowerCAmelCase__		)
					model.to(lowerCAmelCase__		)
					model.eval()
					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	model(lowerCAmelCase__							,  attention_mask=lowerCAmelCase__							,  token_type_ids=lowerCAmelCase__							,  labels=lowerCAmelCase__		)
					self.parent.assertEqual(result.logits.shape							,  (self.batch_size, self.seq_length, self.num_labels)		)

			def UpperCamelCase__      (							self  :				Union[str, Any]							,  lowerCAmelCase__  :				Optional[Any]							,  lowerCAmelCase__  :				int							,  lowerCAmelCase__  :				Dict							,  lowerCAmelCase__  :				Dict							,  lowerCAmelCase__  :				int							,  lowerCAmelCase__  :				List[Any]							,  lowerCAmelCase__  :				str		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	self.num_choices
					__SCREAMING_SNAKE_CASE :	List[str]					 =	AlbertForMultipleChoice(config=lowerCAmelCase__		)
					model.to(lowerCAmelCase__		)
					model.eval()
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	input_ids.unsqueeze(1		).expand(-1							,  self.num_choices							,  -1		).contiguous()
					__SCREAMING_SNAKE_CASE :	Any					 =	token_type_ids.unsqueeze(1		).expand(-1							,  self.num_choices							,  -1		).contiguous()
					__SCREAMING_SNAKE_CASE :	Any					 =	input_mask.unsqueeze(1		).expand(-1							,  self.num_choices							,  -1		).contiguous()
					__SCREAMING_SNAKE_CASE :	Any					 =	model(
					    lowerCAmelCase__							,  attention_mask=lowerCAmelCase__							,  token_type_ids=lowerCAmelCase__							,  labels=lowerCAmelCase__							,  )
					self.parent.assertEqual(result.logits.shape							,  (self.batch_size, self.num_choices)		)



			def UpperCamelCase__      (							self  :				str		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	int					 =	self.prepare_config_and_inputs()
					(
					    (
					    __SCREAMING_SNAKE_CASE
					) ,		(
					    __SCREAMING_SNAKE_CASE
					) ,		(
					    __SCREAMING_SNAKE_CASE
					) ,		(
					    __SCREAMING_SNAKE_CASE
					) ,		(
					    __SCREAMING_SNAKE_CASE
					) ,		(
					    __SCREAMING_SNAKE_CASE
					) ,		(
					    __SCREAMING_SNAKE_CASE
					) ,		
					) :	List[str]					 =	config_and_inputs
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
					return config, inputs_dict






@require_torch
class 		_UpperCamelCase						(							lowerCamelCase__					,      lowerCamelCase__					,      unittest.TestCase					):


			'''simple docstring'''

			_A   :	List[str]    					=	(
			    (
			        AlbertModel,
			        AlbertForPreTraining,
			        AlbertForMaskedLM,
			        AlbertForMultipleChoice,
			        AlbertForSequenceClassification,
			        AlbertForTokenClassification,
			        AlbertForQuestionAnswering,
			    )
			    if is_torch_available()
			    else ()
			)
			_A   :	List[str]    					=	(
			    {
			        '''feature-extraction''': AlbertModel,
			        '''fill-mask''': AlbertForMaskedLM,
			        '''question-answering''': AlbertForQuestionAnswering,
			        '''text-classification''': AlbertForSequenceClassification,
			        '''token-classification''': AlbertForTokenClassification,
			        '''zero-shot''': AlbertForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			_A   :	Union[str, Any]    					=	True

			def UpperCamelCase__      (							self  :				Any							,  lowerCAmelCase__  :				List[str]							,  lowerCAmelCase__  :				int							,  lowerCAmelCase__  :				Union[str, Any]=False		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	str					 =	super()._prepare_for_class(lowerCAmelCase__							,  lowerCAmelCase__							,  return_labels=lowerCAmelCase__		)

					if return_labels:
							if model_class in get_values(lowerCAmelCase__		):
									__SCREAMING_SNAKE_CASE :	Tuple					 =	torch.zeros(
									    (self.model_tester.batch_size, self.model_tester.seq_length)							,  dtype=torch.long							,  device=lowerCAmelCase__		)
									__SCREAMING_SNAKE_CASE :	Dict					 =	torch.zeros(
									    self.model_tester.batch_size							,  dtype=torch.long							,  device=lowerCAmelCase__		)
					return inputs_dict

			def UpperCamelCase__      (							self  :				Optional[int]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	str					 =	AlbertModelTester(self		)
					__SCREAMING_SNAKE_CASE :	Tuple					 =	ConfigTester(self							,  config_class=lowerCAmelCase__							,  hidden_size=3_7		)

			def UpperCamelCase__      (							self  :				int		):





					"""simple docstring"""






					self.config_tester.run_common_tests()

			def UpperCamelCase__      (							self  :				Any		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Any					 =	self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*lowerCAmelCase__		)

			def UpperCamelCase__      (							self  :				Optional[Any]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	str					 =	self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_pretraining(*lowerCAmelCase__		)

			def UpperCamelCase__      (							self  :				int		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase__		)

			def UpperCamelCase__      (							self  :				Any		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Dict					 =	self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_multiple_choice(*lowerCAmelCase__		)

			def UpperCamelCase__      (							self  :				Union[str, Any]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	List[str]					 =	self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_question_answering(*lowerCAmelCase__		)

			def UpperCamelCase__      (							self  :				Optional[int]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_sequence_classification(*lowerCAmelCase__		)

			def UpperCamelCase__      (							self  :				Optional[int]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	self.model_tester.prepare_config_and_inputs()
					for type in ["absolute", "relative_key", "relative_key_query"]:
							__SCREAMING_SNAKE_CASE :	Dict					 =	type
							self.model_tester.create_and_check_model(*lowerCAmelCase__		)



			@slow
			def UpperCamelCase__      (							self  :				Tuple		):





					"""simple docstring"""






					for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							__SCREAMING_SNAKE_CASE :	str					 =	AlbertModel.from_pretrained(lowerCAmelCase__		)
							self.assertIsNotNone(lowerCAmelCase__		)






@require_torch
class 		_UpperCamelCase						(							unittest.TestCase					):


			'''simple docstring'''

			@slow
			def UpperCamelCase__      (							self  :				Dict		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Tuple					 =	AlbertModel.from_pretrained("""albert-base-v2"""		)
					__SCREAMING_SNAKE_CASE :	List[str]					 =	torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]		)
					__SCREAMING_SNAKE_CASE :	Dict					 =	torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
					with torch.no_grad():
							__SCREAMING_SNAKE_CASE :	Any					 =	model(lowerCAmelCase__							,  attention_mask=lowerCAmelCase__		)[0]
					__SCREAMING_SNAKE_CASE :	str					 =	torch.Size((1, 1_1, 7_6_8)		)
					self.assertEqual(output.shape							,  lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	str					 =	torch.tensor(
					    [[[-0.65_13, 1.50_35, -0.27_66], [-0.65_15, 1.50_46, -0.27_80], [-0.65_12, 1.50_49, -0.27_84]]]		)

					self.assertTrue(torch.allclose(output[:, 1:4, 1:4]							,  lowerCAmelCase__							,  atol=1E-4		)		)




'''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 		_UpperCamelCase						(							lowerCamelCase__					,      lowerCamelCase__					):


			'''simple docstring'''

			@register_to_config
			def __init__(							self  :				Any							,  lowerCAmelCase__  :				int = 1_2_8							,  lowerCAmelCase__  :				int = 2_5_6							,  lowerCAmelCase__  :				float = 20_00.0							,  lowerCAmelCase__  :				int = 7_6_8							,  lowerCAmelCase__  :				int = 1_2							,  lowerCAmelCase__  :				int = 1_2							,  lowerCAmelCase__  :				int = 6_4							,  lowerCAmelCase__  :				int = 2_0_4_8							,  lowerCAmelCase__  :				float = 0.1							,  ):





					"""simple docstring"""






					super().__init__()

					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	nn.Sequential(
					    nn.Linear(lowerCAmelCase__							,  d_model * 4							,  bias=lowerCAmelCase__		)							,  nn.SiLU()							,  nn.Linear(d_model * 4							,  d_model * 4							,  bias=lowerCAmelCase__		)							,  nn.SiLU()							,  )

					__SCREAMING_SNAKE_CASE :	Dict					 =	nn.Embedding(lowerCAmelCase__							,  lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	False

					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	nn.Linear(lowerCAmelCase__							,  lowerCAmelCase__							,  bias=lowerCAmelCase__		)

					__SCREAMING_SNAKE_CASE :	str					 =	nn.Dropout(p=lowerCAmelCase__		)

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

					__SCREAMING_SNAKE_CASE :	List[str]					 =	TaLayerNorm(lowerCAmelCase__		)

					__SCREAMING_SNAKE_CASE :	int					 =	nn.Dropout(p=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	nn.Linear(lowerCAmelCase__							,  lowerCAmelCase__							,  bias=lowerCAmelCase__		)

			def UpperCamelCase__      (							self  :				Optional[int]							,  lowerCAmelCase__  :				Optional[Any]							,  lowerCAmelCase__  :				List[Any]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	List[Any]					 =	torch.mul(query_input.unsqueeze(-1		)							,  key_input.unsqueeze(-2		)		)
					return mask.unsqueeze(-3		)



			def UpperCamelCase__      (							self  :				Dict							,  lowerCAmelCase__  :				Optional[int]							,  lowerCAmelCase__  :				Union[str, Any]							,  lowerCAmelCase__  :				List[str]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	List[str]					 =	decoder_input_tokens.shape
					assert decoder_noise_time.shape == (batch,)

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

					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	self.conditioning_emb(lowerCAmelCase__		).unsqueeze(1		)

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

					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	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.
					__SCREAMING_SNAKE_CASE :	Tuple					 =	torch.broadcast_to(
					    torch.arange(lowerCAmelCase__							,  device=decoder_input_tokens.device		)							,  (batch, seq_length)							,  )

					__SCREAMING_SNAKE_CASE :	str					 =	self.position_encoding(lowerCAmelCase__		)

					__SCREAMING_SNAKE_CASE :	Tuple					 =	self.continuous_inputs_projection(lowerCAmelCase__		)
					inputs += position_encodings
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	self.dropout(lowerCAmelCase__		)

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

					# Translate encoding masks to encoder-decoder masks.
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	[(x, self.encoder_decoder_mask(lowerCAmelCase__							,  lowerCAmelCase__		)) for x, y in encodings_and_masks]

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

					for lyr in self.decoders:
							__SCREAMING_SNAKE_CASE :	Dict					 =	lyr(
							    lowerCAmelCase__							,  conditioning_emb=lowerCAmelCase__							,  encoder_hidden_states=lowerCAmelCase__							,  encoder_attention_mask=lowerCAmelCase__							,  )[0]

					__SCREAMING_SNAKE_CASE :	List[str]					 =	self.decoder_norm(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	self.post_dropout(lowerCAmelCase__		)

					__SCREAMING_SNAKE_CASE :	str					 =	self.spec_out(lowerCAmelCase__		)
					return spec_out






class 		_UpperCamelCase						(							nn.Module					):


			'''simple docstring'''

			def __init__(							self  :				Dict							,  lowerCAmelCase__  :				int							,  lowerCAmelCase__  :				Any							,  lowerCAmelCase__  :				Optional[Any]							,  lowerCAmelCase__  :				Optional[int]							,  lowerCAmelCase__  :				str							,  lowerCAmelCase__  :				Tuple=1E-6		):





					"""simple docstring"""






					super().__init__()
					__SCREAMING_SNAKE_CASE :	Dict					 =	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 UpperCamelCase__      (							self  :				int							,  lowerCAmelCase__  :				Dict							,  lowerCAmelCase__  :				Optional[Any]=None							,  lowerCAmelCase__  :				List[Any]=None							,  lowerCAmelCase__  :				Any=None							,  lowerCAmelCase__  :				Dict=None							,  lowerCAmelCase__  :				Optional[int]=None							,  ):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	str					 =	self.layer[0](
					    lowerCAmelCase__							,  conditioning_emb=lowerCAmelCase__							,  attention_mask=lowerCAmelCase__							,  )

					if encoder_hidden_states is not None:
							__SCREAMING_SNAKE_CASE :	List[str]					 =	torch.where(encoder_attention_mask > 0							,  0							,  -1E10		).to(
							    encoder_hidden_states.dtype		)

							__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	self.layer[1](
							    lowerCAmelCase__							,  key_value_states=lowerCAmelCase__							,  attention_mask=lowerCAmelCase__							,  )

					# Apply Film Conditional Feed Forward layer
					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	self.layer[-1](lowerCAmelCase__							,  lowerCAmelCase__		)

					return (hidden_states,)






class 		_UpperCamelCase						(							nn.Module					):


			'''simple docstring'''

			def __init__(							self  :				List[str]							,  lowerCAmelCase__  :				List[str]							,  lowerCAmelCase__  :				str							,  lowerCAmelCase__  :				int							,  lowerCAmelCase__  :				Union[str, Any]		):





					"""simple docstring"""






					super().__init__()
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	TaLayerNorm(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Any					 =	TaFiLMLayer(in_features=d_model * 4							,  out_features=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	List[str]					 =	Attention(query_dim=lowerCAmelCase__							,  heads=lowerCAmelCase__							,  dim_head=lowerCAmelCase__							,  out_bias=lowerCAmelCase__							,  scale_qk=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	nn.Dropout(lowerCAmelCase__		)



			def UpperCamelCase__      (							self  :				int							,  lowerCAmelCase__  :				Union[str, Any]							,  lowerCAmelCase__  :				List[str]=None							,  lowerCAmelCase__  :				int=None							,  ):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Tuple					 =	self.layer_norm(lowerCAmelCase__		)

					if conditioning_emb is not None:
							__SCREAMING_SNAKE_CASE :	str					 =	self.FiLMLayer(lowerCAmelCase__							,  lowerCAmelCase__		)

					# Self-attention block
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	self.attention(lowerCAmelCase__		)

					__SCREAMING_SNAKE_CASE :	str					 =	hidden_states + self.dropout(lowerCAmelCase__		)

					return hidden_states






class 		_UpperCamelCase						(							nn.Module					):


			'''simple docstring'''

			def __init__(							self  :				Union[str, Any]							,  lowerCAmelCase__  :				int							,  lowerCAmelCase__  :				Dict							,  lowerCAmelCase__  :				Tuple							,  lowerCAmelCase__  :				Dict							,  lowerCAmelCase__  :				str		):





					"""simple docstring"""






					super().__init__()
					__SCREAMING_SNAKE_CASE :	Dict					 =	Attention(query_dim=lowerCAmelCase__							,  heads=lowerCAmelCase__							,  dim_head=lowerCAmelCase__							,  out_bias=lowerCAmelCase__							,  scale_qk=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Any					 =	TaLayerNorm(lowerCAmelCase__							,  eps=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	nn.Dropout(lowerCAmelCase__		)



			def UpperCamelCase__      (							self  :				List[Any]							,  lowerCAmelCase__  :				Tuple							,  lowerCAmelCase__  :				List[str]=None							,  lowerCAmelCase__  :				Any=None							,  ):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Dict					 =	self.layer_norm(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	List[str]					 =	self.attention(
					    lowerCAmelCase__							,  encoder_hidden_states=lowerCAmelCase__							,  attention_mask=attention_mask.squeeze(1		)							,  )
					__SCREAMING_SNAKE_CASE :	Dict					 =	hidden_states + self.dropout(lowerCAmelCase__		)
					return layer_output






class 		_UpperCamelCase						(							nn.Module					):


			'''simple docstring'''

			def __init__(							self  :				List[str]							,  lowerCAmelCase__  :				Dict							,  lowerCAmelCase__  :				Any							,  lowerCAmelCase__  :				Tuple							,  lowerCAmelCase__  :				int		):





					"""simple docstring"""






					super().__init__()
					__SCREAMING_SNAKE_CASE :	Any					 =	TaDenseGatedActDense(d_model=lowerCAmelCase__							,  d_ff=lowerCAmelCase__							,  dropout_rate=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	TaFiLMLayer(in_features=d_model * 4							,  out_features=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Tuple					 =	TaLayerNorm(lowerCAmelCase__							,  eps=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Tuple					 =	nn.Dropout(lowerCAmelCase__		)



			def UpperCamelCase__      (							self  :				int							,  lowerCAmelCase__  :				List[Any]							,  lowerCAmelCase__  :				Tuple=None		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	str					 =	self.layer_norm(lowerCAmelCase__		)
					if conditioning_emb is not None:
							__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	self.film(lowerCAmelCase__							,  lowerCAmelCase__		)

					__SCREAMING_SNAKE_CASE :	int					 =	self.DenseReluDense(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Any					 =	hidden_states + self.dropout(lowerCAmelCase__		)
					return hidden_states






class 		_UpperCamelCase						(							nn.Module					):


			'''simple docstring'''

			def __init__(							self  :				Tuple							,  lowerCAmelCase__  :				List[str]							,  lowerCAmelCase__  :				Optional[int]							,  lowerCAmelCase__  :				Dict		):





					"""simple docstring"""






					super().__init__()
					__SCREAMING_SNAKE_CASE :	Tuple					 =	nn.Linear(lowerCAmelCase__							,  lowerCAmelCase__							,  bias=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Any					 =	nn.Linear(lowerCAmelCase__							,  lowerCAmelCase__							,  bias=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Dict					 =	nn.Linear(lowerCAmelCase__							,  lowerCAmelCase__							,  bias=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	nn.Dropout(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Dict					 =	NewGELUActivation()



			def UpperCamelCase__      (							self  :				List[str]							,  lowerCAmelCase__  :				Optional[Any]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	int					 =	self.act(self.wi_a(lowerCAmelCase__		)		)
					__SCREAMING_SNAKE_CASE :	Any					 =	self.wi_a(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	hidden_gelu * hidden_linear
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	self.dropout(lowerCAmelCase__		)

					__SCREAMING_SNAKE_CASE :	List[str]					 =	self.wo(lowerCAmelCase__		)
					return hidden_states






class 		_UpperCamelCase						(							nn.Module					):


			'''simple docstring'''

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





					"""simple docstring"""






					super().__init__()
					__SCREAMING_SNAKE_CASE :	Dict					 =	nn.Parameter(torch.ones(lowerCAmelCase__		)		)
					__SCREAMING_SNAKE_CASE :	str					 =	eps



			def UpperCamelCase__      (							self  :				Dict							,  lowerCAmelCase__  :				Optional[Any]		):





					"""simple docstring"""






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

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

					return self.weight * hidden_states






class 		_UpperCamelCase						(							nn.Module					):


			'''simple docstring'''

			def UpperCamelCase__      (							self  :				Union[str, Any]							,  lowerCAmelCase__  :				torch.Tensor		):





					"""simple docstring"""






					return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi		) * (input + 0.04_47_15 * torch.pow(lowerCAmelCase__							,  3.0		))		))






class 		_UpperCamelCase						(							nn.Module					):


			'''simple docstring'''

			def __init__(							self  :				Any							,  lowerCAmelCase__  :				Dict							,  lowerCAmelCase__  :				Union[str, Any]		):





					"""simple docstring"""






					super().__init__()
					__SCREAMING_SNAKE_CASE :	Tuple					 =	nn.Linear(lowerCAmelCase__							,  out_features * 2							,  bias=lowerCAmelCase__		)



			def UpperCamelCase__      (							self  :				List[Any]							,  lowerCAmelCase__  :				List[str]							,  lowerCAmelCase__  :				Tuple		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Dict					 =	self.scale_bias(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	Tuple					 =	torch.chunk(lowerCAmelCase__							,  2							,  -1		)
					__SCREAMING_SNAKE_CASE :	Dict					 =	x * (1 + scale) + shift
					return x