


'''simple docstring'''

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





      '''simple docstring'''






      def __init__(self   :   Any       ,						UpperCamelCase   :   List[Any]       ,						UpperCamelCase   :   int=99       ,						UpperCamelCase   :   str=13       ,						UpperCamelCase   :   Union[str, Any]=16       ,						UpperCamelCase   :   Optional[Any]=7       ,						UpperCamelCase   :   int=True       ,						UpperCamelCase   :   Optional[Any]=True       ,						UpperCamelCase   :   Optional[int]=True       ,						UpperCamelCase   :   List[Any]=False       ,						UpperCamelCase   :   Optional[int]=True       ,						UpperCamelCase   :   Optional[Any]=2       ,						UpperCamelCase   :   Optional[Any]=32       ,						UpperCamelCase   :   List[Any]=4       ,						UpperCamelCase   :   Union[str, Any]=4       ,						UpperCamelCase   :   str=30       ,						UpperCamelCase   :   str=0       ,						UpperCamelCase   :   List[str]=1       ,						UpperCamelCase   :   List[Any]=2       ,						UpperCamelCase   :   List[Any]=None       ,						):

             '''simple docstring'''







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

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

             lowercase__				      =			None
             lowercase__				      =			decoder_seq_length
             lowercase__				      =			2
             lowercase__				      =			1



      def 				UpperCamelCase__		(self   :   Any			):

             '''simple docstring'''







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

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

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

             lowercase__				      =			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       ,						UpperCamelCase   :   Optional[int]       ,						UpperCamelCase   :   List[Any]       ,						UpperCamelCase   :   int       ,						UpperCamelCase   :   Dict       ,						):

             '''simple docstring'''







             lowercase__				      =			True
             lowercase__				      =			SpeechaTextaDecoder(config=lowercase_			).to(lowercase_			).eval()
             lowercase__				      =			input_ids[:2]

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

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

             lowercase__				      =			outputs['past_key_values']

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

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

             print(lowercase_			)

             lowercase__				      =			model(lowercase_			)['last_hidden_state']
             lowercase__				      =			model(lowercase_       ,						past_key_values=lowercase_			)['last_hidden_state']

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

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






      def 				UpperCamelCase__		(self   :   Any			):

             '''simple docstring'''







             lowercase__				      =			self.prepare_config_and_inputs()
             (
                 lowercase__					
             )				      =			config_and_inputs

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






@require_torch
class   __lowerCAmelCase						(UpperCamelCase_  ,					UpperCamelCase_  ,					UpperCamelCase_  ,					unittest.TestCase					):





      '''simple docstring'''






      lowerCAmelCase__ :					Dict   =							(SpeechaTextaDecoder, SpeechaTextaForCausalLM) if is_torch_available() else ()
      lowerCAmelCase__ :					Optional[Any]   =							(SpeechaTextaForCausalLM,) if is_torch_available() else ()
      lowerCAmelCase__ :					List[str]   =							{'text-generation': SpeechaTextaForCausalLM} if is_torch_available() else {}
      lowerCAmelCase__ :					List[str]   =							True
      lowerCAmelCase__ :					Optional[Any]   =							False



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

             '''simple docstring'''







             lowercase__				      =			SpeechaTextaStandaloneDecoderModelTester(self       ,						is_training=lowercase_			)
             lowercase__				      =			ConfigTester(self       ,						config_class=lowercase_			)



      def 				UpperCamelCase__		(self   :   Optional[int]			):

             '''simple docstring'''







             pass



      def 				UpperCamelCase__		(self   :   str			):

             '''simple docstring'''







             pass



      def 				UpperCamelCase__		(self   :   str			):

             '''simple docstring'''







             pass



      def 				UpperCamelCase__		(self   :   str			):

             '''simple docstring'''







             self.config_tester.run_common_tests()



      def 				UpperCamelCase__		(self   :   List[str]			):

             '''simple docstring'''







             lowercase__				      =			self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_decoder_model_past(*lowercase_			)






      def 				UpperCamelCase__		(self   :   Optional[Any]			):

             '''simple docstring'''







             return








'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available, is_torch_available


lowerCamelCase							:       Tuple  						=     {'configuration_bert_generation': ['BertGenerationConfig']}

try:
 if not is_sentencepiece_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase							:       Union[str, Any]  						=     ['BertGenerationTokenizer']

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase							:       Dict  						=     [
     'BertGenerationDecoder',
     'BertGenerationEncoder',
     'BertGenerationPreTrainedModel',
     'load_tf_weights_in_bert_generation',
 ]


if TYPE_CHECKING:
 from .configuration_bert_generation import BertGenerationConfig

 try:
  if not is_sentencepiece_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_bert_generation import BertGenerationTokenizer

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_bert_generation import (
      BertGenerationDecoder,
      BertGenerationEncoder,
      BertGenerationPreTrainedModel,
      load_tf_weights_in_bert_generation,
  )

else:
 import sys

 lowerCamelCase							:       int  						=     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




