
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     SCREAMING_SNAKE_CASE__      :

							"""simple docstring"""






							def __init__(					self       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_=99       ,		SCREAMING_SNAKE_CASE_=13       ,		SCREAMING_SNAKE_CASE_=16       ,		SCREAMING_SNAKE_CASE_=7       ,		SCREAMING_SNAKE_CASE_=True       ,		SCREAMING_SNAKE_CASE_=True       ,		SCREAMING_SNAKE_CASE_=True       ,		SCREAMING_SNAKE_CASE_=False       ,		SCREAMING_SNAKE_CASE_=True       ,		SCREAMING_SNAKE_CASE_=2       ,		SCREAMING_SNAKE_CASE_=32       ,		SCREAMING_SNAKE_CASE_=4       ,		SCREAMING_SNAKE_CASE_=4       ,		SCREAMING_SNAKE_CASE_=30       ,		SCREAMING_SNAKE_CASE_=0       ,		SCREAMING_SNAKE_CASE_=1       ,		SCREAMING_SNAKE_CASE_=2       ,		SCREAMING_SNAKE_CASE_=None       ,		)->				Dict:





													'''simple docstring'''






													__UpperCamelCase       					=     parent
													__UpperCamelCase       					=     batch_size
													__UpperCamelCase       					=     decoder_seq_length
													# For common tests
													__UpperCamelCase       					=     self.decoder_seq_length
													__UpperCamelCase       					=     is_training
													__UpperCamelCase       					=     use_attention_mask
													__UpperCamelCase       					=     use_labels

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

													__UpperCamelCase       					=     None
													__UpperCamelCase       					=     decoder_seq_length
													__UpperCamelCase       					=     2
													__UpperCamelCase       					=     1



							def       A__						(					self     )->				Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase       					=     ids_tensor([self.batch_size, self.decoder_seq_length]       ,		self.vocab_size     )

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

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

													__UpperCamelCase       					=     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       A__						(					self       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		)->				Optional[int]:





													'''simple docstring'''






													__UpperCamelCase       					=     True
													__UpperCamelCase       					=     SpeechaTextaDecoder(config=SCREAMING_SNAKE_CASE_     ).to(SCREAMING_SNAKE_CASE_     ).eval()
													__UpperCamelCase       					=     input_ids[:2]

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

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

													__UpperCamelCase       					=     outputs['''past_key_values''']

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

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

													print(SCREAMING_SNAKE_CASE_     )

													__UpperCamelCase       					=     model(SCREAMING_SNAKE_CASE_     )['''last_hidden_state''']
													__UpperCamelCase       					=     model(SCREAMING_SNAKE_CASE_       ,		past_key_values=SCREAMING_SNAKE_CASE_     )['''last_hidden_state''']

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

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



							def       A__						(					self     )->				Tuple:





													'''simple docstring'''






													__UpperCamelCase       					=     self.prepare_config_and_inputs()
													(
													    (
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,							(
													    __UpperCamelCase							
													)					,														
													)       					=     config_and_inputs

													__UpperCamelCase       					=     {
													    '''input_ids''': input_ids,
													    '''attention_mask''': attention_mask,
													}
													return config, inputs_dict






@require_torch
class     SCREAMING_SNAKE_CASE__      (							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       A__						(					self       ,		)->				Optional[int]:





													'''simple docstring'''






													__UpperCamelCase       					=     SpeechaTextaStandaloneDecoderModelTester(self       ,		is_training=SCREAMING_SNAKE_CASE_     )
													__UpperCamelCase       					=     ConfigTester(self       ,		config_class=SCREAMING_SNAKE_CASE_     )



							def       A__						(					self     )->				Optional[int]:





													'''simple docstring'''






													pass



							def       A__						(					self     )->				Optional[Any]:





													'''simple docstring'''






													pass



							def       A__						(					self     )->				int:





													'''simple docstring'''






													pass



							def       A__						(					self     )->				Optional[int]:





													'''simple docstring'''






													self.config_tester.run_common_tests()



							def       A__						(					self     )->				Any:





													'''simple docstring'''






													__UpperCamelCase       					=     self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_decoder_model_past(*SCREAMING_SNAKE_CASE_     )



							def       A__						(					self     )->				Tuple:





													'''simple docstring'''






													return






import torch

from transformers import PreTrainedModel

from .custom_configuration import CustomConfig, NoSuperInitConfig






class     SCREAMING_SNAKE_CASE__      (							SCREAMING_SNAKE_CASE_      ):

							"""simple docstring"""






							_snake_case       						=					CustomConfig



							def __init__(					self       ,		SCREAMING_SNAKE_CASE_     )->				Tuple:





													'''simple docstring'''






													super().__init__(SCREAMING_SNAKE_CASE_     )
													__UpperCamelCase       					=     torch.nn.Linear(config.hidden_size       ,		config.hidden_size     )



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_     )->				List[str]:





													'''simple docstring'''






													return self.linear(SCREAMING_SNAKE_CASE_     )



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_     )->				Any:





													'''simple docstring'''






													pass






class     SCREAMING_SNAKE_CASE__      (							SCREAMING_SNAKE_CASE_      ):

							"""simple docstring"""






							_snake_case       						=					NoSuperInitConfig



							def __init__(					self       ,		SCREAMING_SNAKE_CASE_     )->				Tuple:





													'''simple docstring'''






													super().__init__(SCREAMING_SNAKE_CASE_     )
													__UpperCamelCase       					=     torch.nn.Linear(config.attribute       ,		config.attribute     )



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_     )->				List[str]:





													'''simple docstring'''






													return self.linear(SCREAMING_SNAKE_CASE_     )



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_     )->				Union[str, Any]:





													'''simple docstring'''






													pass




