
from __future__ import annotations

import unittest

from transformers import AlbertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
					import tensorflow as tf

					from transformers import TF_MODEL_FOR_PRETRAINING_MAPPING
					from transformers.models.albert.modeling_tf_albert import (
					    TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
					    TFAlbertForMaskedLM,
					    TFAlbertForMultipleChoice,
					    TFAlbertForPreTraining,
					    TFAlbertForQuestionAnswering,
					    TFAlbertForSequenceClassification,
					    TFAlbertForTokenClassification,
					    TFAlbertModel,
					)





class 							UpperCAmelCase   :




	'''simple docstring'''





	def __init__( self      :				Union[str, Any]							,A      :				List[str]							,A      :				int=13							,A      :				List[str]=7							,A      :				Tuple=True							,A      :				Optional[int]=True							,A      :				Any=True							,A      :				Optional[Any]=True							,A      :				Any=99							,A      :				List[Any]=16							,A      :				Tuple=32							,A      :				Any=2							,A      :				Tuple=4							,A      :				Dict=37							,A      :				Optional[int]="gelu"							,A      :				List[Any]=0.1							,A      :				List[Any]=0.1							,A      :				List[str]=5_12							,A      :				str=16							,A      :				List[str]=2							,A      :				int=0.02							,A      :				List[Any]=3							,A      :				Optional[int]=4							,A      :				Optional[Any]=None							,):
					__A         =				parent
					__A         =				13
					__A         =				7
					__A         =				True
					__A         =				True
					__A         =				True
					__A         =				True
					__A         =				99
					__A         =				16
					__A         =				32
					__A         =				2
					__A         =				4
					__A         =				37
					__A         =				"gelu"
					__A         =				0.1
					__A         =				0.1
					__A         =				5_12
					__A         =				16
					__A         =				2
					__A         =				0.02
					__A         =				3
					__A         =				4
					__A         =				None






	def        UpperCamelCase_						( self      :				Optional[Any]				):
					__A         =				ids_tensor([self.batch_size, self.seq_length]							,self.vocab_size				)

					__A         =				None
					if self.use_input_mask:
									__A         =				random_attention_mask([self.batch_size, self.seq_length]				)

					__A         =				None
					if self.use_token_type_ids:
									__A         =				ids_tensor([self.batch_size, self.seq_length]							,self.type_vocab_size				)

					__A         =				None
					__A         =				None
					__A         =				None
					if self.use_labels:
									__A         =				ids_tensor([self.batch_size]							,self.type_sequence_label_size				)
									__A         =				ids_tensor([self.batch_size, self.seq_length]							,self.num_labels				)
									__A         =				ids_tensor([self.batch_size]							,self.num_choices				)

					__A         =				AlbertConfig(
					    vocab_size=self.vocab_size							,hidden_size=self.hidden_size							,embedding_size=self.embedding_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							,)

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






	def        UpperCamelCase_						( self      :				int							,A      :				Union[str, Any]							,A      :				str							,A      :				Optional[int]							,A      :				Any							,A      :				Tuple							,A      :				Union[str, Any]							,A      :				int				):
					__A         =				TFAlbertModel(config=A				)
					# inputs = {'input_ids': input_ids,
					#           'attention_mask': input_mask,
					#           'token_type_ids': token_type_ids}
					# sequence_output, pooled_output = model(**inputs)
					__A         =				{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
					__A         =				model(A				)

					__A         =				[input_ids, input_mask]
					__A         =				model(A				)

					__A         =				model(A				)

					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      :				str							,A      :				List[Any]							,A      :				Any							,A      :				Tuple							,A      :				Optional[int]							,A      :				Union[str, Any]							,A      :				Union[str, Any]							,A      :				Optional[Any]				):
					__A         =				self.num_labels
					__A         =				TFAlbertForPreTraining(config=A				)
					__A         =				{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
					__A         =				model(A				)
					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, self.num_labels)				)






	def        UpperCamelCase_						( self      :				Union[str, Any]							,A      :				Any							,A      :				Tuple							,A      :				List[Any]							,A      :				str							,A      :				Any							,A      :				Any							,A      :				Any				):
					__A         =				TFAlbertForMaskedLM(config=A				)
					__A         =				{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
					__A         =				model(A				)
					self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.seq_length, self.vocab_size)				)






	def        UpperCamelCase_						( self      :				Any							,A      :				Union[str, Any]							,A      :				Tuple							,A      :				Any							,A      :				Union[str, Any]							,A      :				int							,A      :				int							,A      :				str				):
					__A         =				self.num_labels
					__A         =				TFAlbertForSequenceClassification(config=A				)
					__A         =				{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
					__A         =				model(A				)
					self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.num_labels)				)






	def        UpperCamelCase_						( self      :				int							,A      :				List[str]							,A      :				Tuple							,A      :				str							,A      :				int							,A      :				Dict							,A      :				str							,A      :				Tuple				):
					__A         =				TFAlbertForQuestionAnswering(config=A				)
					__A         =				{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
					__A         =				model(A				)
					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]							,A      :				int							,A      :				Dict							,A      :				Any							,A      :				Optional[int]							,A      :				Optional[int]							,A      :				Tuple							,A      :				Tuple				):
					__A         =				self.num_choices
					__A         =				TFAlbertForMultipleChoice(config=A				)
					__A         =				tf.tile(tf.expand_dims(A							,1				)							,(1, self.num_choices, 1)				)
					__A         =				tf.tile(tf.expand_dims(A							,1				)							,(1, self.num_choices, 1)				)
					__A         =				tf.tile(tf.expand_dims(A							,1				)							,(1, self.num_choices, 1)				)
					__A         =				{
					    "input_ids": multiple_choice_inputs_ids,
					    "attention_mask": multiple_choice_input_mask,
					    "token_type_ids": multiple_choice_token_type_ids,
					}
					__A         =				model(A				)
					self.parent.assertListEqual(list(result["logits"].shape				)							,[self.batch_size, self.num_choices]				)






	def        UpperCamelCase_						( self      :				Union[str, Any]							,A      :				List[Any]							,A      :				str							,A      :				Tuple							,A      :				Optional[int]							,A      :				str							,A      :				Tuple							,A      :				Union[str, Any]				):
					__A         =				self.num_labels
					__A         =				TFAlbertForTokenClassification(config=A				)
					__A         =				{
					    "input_ids": input_ids,
					    "attention_mask": input_mask,
					    "token_type_ids": token_type_ids,
					}
					__A         =				model(A				)
					self.parent.assertListEqual(list(result["logits"].shape				)							,[self.batch_size, self.seq_length, self.num_labels]				)







	def        UpperCamelCase_						( self      :				int				):
					__A         =				self.prepare_config_and_inputs()
					(
					    (
					    __A
					)    ,							(
					    __A
					)    ,							(
					    __A
					)    ,							(
					    __A
					)    ,							(
					    __A
					)    ,							(
					    __A
					)    ,							(
					    __A
					)    ,							
					)         =				config_and_inputs
					__A         =				{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
					return config, inputs_dict





@require_tf
class 							UpperCAmelCase   (						__SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE    ,    unittest.TestCase						):




	'''simple docstring'''





	snake_case_      					=  (
	    (
	        TFAlbertModel,
	        TFAlbertForPreTraining,
	        TFAlbertForMaskedLM,
	        TFAlbertForSequenceClassification,
	        TFAlbertForQuestionAnswering,
	        TFAlbertForTokenClassification,
	        TFAlbertForMultipleChoice,
	    )
	    if is_tf_available()
	    else ()
	)
	snake_case_      					=  (
	    {
	        "feature-extraction": TFAlbertModel,
	        "fill-mask": TFAlbertForMaskedLM,
	        "question-answering": TFAlbertForQuestionAnswering,
	        "text-classification": TFAlbertForSequenceClassification,
	        "token-classification": TFAlbertForTokenClassification,
	        "zero-shot": TFAlbertForSequenceClassification,
	    }
	    if is_tf_available()
	    else {}
	)
	snake_case_      					=  False
	snake_case_      					=  False






	def        UpperCamelCase_						( self      :				Optional[int]							,A      :				Any							,A      :				Dict							,A      :				List[str]=False				):
					__A         =				super()._prepare_for_class(A							,A							,return_labels=A				)

					if return_labels:
									if model_class in get_values(A				):
													__A         =				tf.zeros(self.model_tester.batch_size							,dtype=tf.intaa				)

					return inputs_dict






	def        UpperCamelCase_						( self      :				int				):
					__A         =				TFAlbertModelTester(self				)
					__A         =				ConfigTester(self							,config_class=A							,hidden_size=37				)






	def        UpperCamelCase_						( self      :				Dict				):
					self.config_tester.run_common_tests()






	def        UpperCamelCase_						( self      :				str				):
					__A         =				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_albert_model(*A				)






	def        UpperCamelCase_						( self      :				Optional[Any]				):
					__A         =				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_albert_for_pretraining(*A				)






	def        UpperCamelCase_						( self      :				Any				):
					__A         =				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_albert_for_masked_lm(*A				)






	def        UpperCamelCase_						( self      :				Dict				):
					__A         =				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_albert_for_multiple_choice(*A				)






	def        UpperCamelCase_						( self      :				List[str]				):
					__A         =				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_albert_for_sequence_classification(*A				)






	def        UpperCamelCase_						( self      :				int				):
					__A         =				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_albert_for_question_answering(*A				)







	@slow
	def        UpperCamelCase_						( self      :				Dict				):
					for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									__A         =				TFAlbertModel.from_pretrained(A				)
									self.assertIsNotNone(A				)





@require_tf
class 							UpperCAmelCase   (						unittest.TestCase						):




	'''simple docstring'''





	@slow
	def        UpperCamelCase_						( self      :				Tuple				):
					__A         =				TFAlbertForPreTraining.from_pretrained("albert-base-v2"				)
					__A         =				tf.constant([[0, 1, 2, 3, 4, 5]]				)
					__A         =				model(A				)[0]

					__A         =				[1, 6, 3_00_00]
					self.assertEqual(output.shape							,A				)

					__A         =				tf.constant(
					    [
					        [
					            [4.59_56_68, 0.74_46_27_54, -1.81_81_47],
					            [4.5_95_43_47, 0.7_45_41_84, -1.8_18_82_58],
					            [4.5_95_49_05, 0.7_44_82_35, -1.8_18_23_16],
					        ]
					    ]				)
					tf.debugging.assert_near(output[:, :3, :3]							,A							,atol=1E-4				)


import unittest

from transformers import SpeechaTextaConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers.models.speech_to_text_a.modeling_speech_to_text_a import (
					    SpeechaTextaDecoder,
					    SpeechaTextaForCausalLM,
					)





@require_torch
class 							UpperCAmelCase   :




	'''simple docstring'''





	def __init__( self      :				Dict							,A      :				Tuple							,A      :				Union[str, Any]=99							,A      :				Dict=13							,A      :				int=16							,A      :				int=7							,A      :				Optional[int]=True							,A      :				str=True							,A      :				Optional[Any]=True							,A      :				Optional[int]=False							,A      :				Dict=True							,A      :				Any=2							,A      :				List[Any]=32							,A      :				Optional[Any]=4							,A      :				Tuple=4							,A      :				List[Any]=30							,A      :				Optional[Any]=0							,A      :				str=1							,A      :				List[Any]=2							,A      :				List[Any]=None							,):
					__A         =				parent
					__A         =				batch_size
					__A         =				decoder_seq_length
					# For common tests
					__A         =				self.decoder_seq_length
					__A         =				is_training
					__A         =				use_attention_mask
					__A         =				use_labels

					__A         =				vocab_size
					__A         =				d_model
					__A         =				d_model
					__A         =				decoder_layers
					__A         =				decoder_layers
					__A         =				decoder_ffn_dim
					__A         =				decoder_attention_heads
					__A         =				decoder_attention_heads
					__A         =				eos_token_id
					__A         =				bos_token_id
					__A         =				pad_token_id
					__A         =				decoder_start_token_id
					__A         =				use_cache
					__A         =				max_position_embeddings

					__A         =				None
					__A         =				decoder_seq_length
					__A         =				2
					__A         =				1






	def        UpperCamelCase_						( self      :				Tuple				):
					__A         =				ids_tensor([self.batch_size, self.decoder_seq_length]							,self.vocab_size				)

					__A         =				None
					if self.use_attention_mask:
									__A         =				ids_tensor([self.batch_size, self.decoder_seq_length]							,vocab_size=2				)

					__A         =				None
					if self.use_labels:
									__A         =				ids_tensor([self.batch_size, self.decoder_seq_length]							,self.vocab_size				)

					__A         =				SpeechaTextaConfig(
					    vocab_size=self.vocab_size							,d_model=self.d_model							,decoder_layers=self.decoder_layers							,decoder_ffn_dim=self.decoder_ffn_dim							,decoder_attention_heads=self.decoder_attention_heads							,eos_token_id=self.eos_token_id							,bos_token_id=self.bos_token_id							,use_cache=self.use_cache							,pad_token_id=self.pad_token_id							,decoder_start_token_id=self.decoder_start_token_id							,max_position_embeddings=self.max_position_embeddings							,)

					return (
					    config,
					    input_ids,
					    attention_mask,
					    lm_labels,
					)






	def        UpperCamelCase_						( self      :				Tuple							,A      :				List[str]							,A      :				Tuple							,A      :				Optional[int]							,A      :				List[str]							,):
					__A         =				True
					__A         =				SpeechaTextaDecoder(config=A				).to(A				).eval()
					__A         =				input_ids[:2]

					input_ids[input_ids == 0] += 1
					# first forward pass
					__A         =				model(A							,use_cache=A				)
					__A         =				model(A				)
					__A         =				model(A							,use_cache=A				)

					self.parent.assertTrue(len(A				) == len(A				)				)
					self.parent.assertTrue(len(A				) == len(A				) + 1				)

					__A         =				outputs["past_key_values"]

					# create hypothetical next token and extent to next_input_ids
					__A         =				ids_tensor((2, 1)							,config.vocab_size - 1				) + 1

					# append to next input_ids and
					__A         =				torch.cat([input_ids, next_tokens]							,dim=-1				)

					print(A				)

					__A         =				model(A				)["last_hidden_state"]
					__A         =				model(A							,past_key_values=A				)["last_hidden_state"]

					# select random slice
					__A         =				ids_tensor((1,)							,output_from_past.shape[-1]				).item()
					__A         =				output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
					__A         =				output_from_past[:, 0, random_slice_idx].detach()

					# test that outputs are equal for slice
					assert torch.allclose(A							,A							,atol=1E-3				)







	def        UpperCamelCase_						( self      :				Any				):
					__A         =				self.prepare_config_and_inputs()
					(
					    (
					    __A
					)    ,							(
					    __A
					)    ,							(
					    __A
					)    ,							(
					    __A
					)    ,							
					)         =				config_and_inputs

					__A         =				{
					    "input_ids": input_ids,
					    "attention_mask": attention_mask,
					}
					return config, inputs_dict





@require_torch
class 							UpperCAmelCase   (						__SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE    ,    unittest.TestCase						):




	'''simple docstring'''





	snake_case_      					=  (SpeechaTextaDecoder, SpeechaTextaForCausalLM) if is_torch_available() else ()
	snake_case_      					=  (SpeechaTextaForCausalLM,) if is_torch_available() else ()
	snake_case_      					=  {"text-generation": SpeechaTextaForCausalLM} if is_torch_available() else {}
	snake_case_      					=  True
	snake_case_      					=  False






	def        UpperCamelCase_						( self      :				Any							,):
					__A         =				SpeechaTextaStandaloneDecoderModelTester(self							,is_training=A				)
					__A         =				ConfigTester(self							,config_class=A				)






	def        UpperCamelCase_						( self      :				Optional[Any]				):
					pass






	def        UpperCamelCase_						( self      :				int				):
					pass






	def        UpperCamelCase_						( self      :				Tuple				):
					pass






	def        UpperCamelCase_						( self      :				str				):
					self.config_tester.run_common_tests()






	def        UpperCamelCase_						( self      :				Dict				):
					__A         =				self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_decoder_model_past(*A				)







	def        UpperCamelCase_						( self      :				str				):
					return
