






"""simple docstring"""






import unittest

from transformers import BertGenerationConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
  import torch

  from transformers import BertGenerationDecoder, BertGenerationEncoder




class 							lowerCAmelCase__				:






      '''simple docstring'''

      def __init__(						self     :  Tuple						,	lowercase_     :  int						,	lowercase_     :  Optional[int]=13						,	lowercase_     :  List[Any]=7						,	lowercase_     :  Tuple=True						,	lowercase_     :  List[Any]=True						,	lowercase_     :  Union[str, Any]=99						,	lowercase_     :  Any=32						,	lowercase_     :  Optional[Any]=5						,	lowercase_     :  Any=4						,	lowercase_     :  str=37						,	lowercase_     :  Optional[int]="gelu"						,	lowercase_     :  List[Any]=0.1						,	lowercase_     :  List[str]=0.1						,	lowercase_     :  int=50						,	lowercase_     :  List[Any]=0.02						,	lowercase_     :  Any=True						,	lowercase_     :  List[Any]=None						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		parent
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		batch_size
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		seq_length
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		is_training
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		use_input_mask
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		vocab_size
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		hidden_size
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		num_hidden_layers
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		num_attention_heads
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		intermediate_size
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		hidden_act
          SCREAMING_SNAKE_CASE_      :  str    						=		hidden_dropout_prob
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		attention_probs_dropout_prob
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		max_position_embeddings
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		initializer_range
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		use_labels
          SCREAMING_SNAKE_CASE_      :  Dict    						=		scope
      def _SCREAMING_SNAKE_CASE (						self     :  int):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Tuple    						=		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_      :  int    						=		random_attention_mask([self.batch_size, self.seq_length])

          if self.use_labels:
              SCREAMING_SNAKE_CASE_      :  str    						=		ids_tensor([self.batch_size, self.seq_length]						,	self.vocab_size)

          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		self.get_config()

          return config, input_ids, input_mask, token_labels
      def _SCREAMING_SNAKE_CASE (						self     :  int):




          '''simple docstring'''


          return BertGenerationConfig(
              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						,	is_decoder=snake_case__						,	initializer_range=self.initializer_range						,	)
      def _SCREAMING_SNAKE_CASE (						self     :  List[str]):




          '''simple docstring'''


          (
              SCREAMING_SNAKE_CASE_					
          )      :  Tuple    						=		self.prepare_config_and_inputs()

          SCREAMING_SNAKE_CASE_      :  Any    						=		True
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		ids_tensor([self.batch_size, self.seq_length]						,	vocab_size=2)

          return (
              config,
              input_ids,
              input_mask,
              token_labels,
              encoder_hidden_states,
              encoder_attention_mask,
          )
      def _SCREAMING_SNAKE_CASE (						self     :  List[Any]						,	lowercase_     :  str						,	lowercase_     :  str						,	lowercase_     :  Tuple						,	lowercase_     :  int						,	**lowercase_     :  int						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  int    						=		BertGenerationEncoder(config=snake_case__)
          model.to(snake_case__)
          model.eval()
          SCREAMING_SNAKE_CASE_      :  str    						=		model(snake_case__						,	attention_mask=snake_case__)
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		model(snake_case__)
          self.parent.assertEqual(result.last_hidden_state.shape						,	(self.batch_size, self.seq_length, self.hidden_size))
      def _SCREAMING_SNAKE_CASE (						self     :  int						,	lowercase_     :  Tuple						,	lowercase_     :  Optional[Any]						,	lowercase_     :  Optional[int]						,	lowercase_     :  str						,	lowercase_     :  Optional[Any]						,	lowercase_     :  int						,	**lowercase_     :  Optional[Any]						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  str    						=		True
          SCREAMING_SNAKE_CASE_      :  int    						=		BertGenerationEncoder(config=snake_case__)
          model.to(snake_case__)
          model.eval()
          SCREAMING_SNAKE_CASE_      :  Dict    						=		model(
              snake_case__						,	attention_mask=snake_case__						,	encoder_hidden_states=snake_case__						,	encoder_attention_mask=snake_case__						,	)
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		model(
              snake_case__						,	attention_mask=snake_case__						,	encoder_hidden_states=snake_case__						,	)
          self.parent.assertEqual(result.last_hidden_state.shape						,	(self.batch_size, self.seq_length, self.hidden_size))
      def _SCREAMING_SNAKE_CASE (						self     :  int						,	lowercase_     :  Optional[Any]						,	lowercase_     :  Any						,	lowercase_     :  Dict						,	lowercase_     :  Any						,	lowercase_     :  Union[str, Any]						,	lowercase_     :  str						,	**lowercase_     :  str						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		True
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		True
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		BertGenerationDecoder(config=snake_case__).to(snake_case__).eval()

          # first forward pass
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		model(
              snake_case__						,	attention_mask=snake_case__						,	encoder_hidden_states=snake_case__						,	encoder_attention_mask=snake_case__						,	use_cache=snake_case__						,	)
          SCREAMING_SNAKE_CASE_      :  Any    						=		outputs.past_key_values

          # create hypothetical multiple next token and extent to next_input_ids
          SCREAMING_SNAKE_CASE_      :  Any    						=		ids_tensor((self.batch_size, 3)						,	config.vocab_size)
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		ids_tensor((self.batch_size, 3)						,	vocab_size=2)

          # append to next input_ids and
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		torch.cat([input_ids, next_tokens]						,	dim=-1)
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		torch.cat([input_mask, next_mask]						,	dim=-1)

          SCREAMING_SNAKE_CASE_      :  Tuple    						=		model(
              snake_case__						,	attention_mask=snake_case__						,	encoder_hidden_states=snake_case__						,	encoder_attention_mask=snake_case__						,	output_hidden_states=snake_case__						,	)["hidden_states"][0]
          SCREAMING_SNAKE_CASE_      :  Dict    						=		model(
              snake_case__						,	attention_mask=snake_case__						,	encoder_hidden_states=snake_case__						,	encoder_attention_mask=snake_case__						,	past_key_values=snake_case__						,	output_hidden_states=snake_case__						,	)["hidden_states"][0]

          # select random slice
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		ids_tensor((1,)						,	output_from_past.shape[-1]).item()
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		output_from_no_past[:, -3:, random_slice_idx].detach()
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		output_from_past[:, :, random_slice_idx].detach()

          self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1])

          # test that outputs are equal for slice
          self.parent.assertTrue(torch.allclose(snake_case__						,	snake_case__						,	atol=1e-3))
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple						,	lowercase_     :  Tuple						,	lowercase_     :  Dict						,	lowercase_     :  Tuple						,	lowercase_     :  Optional[Any]						,	*lowercase_     :  List[Any]						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  str    						=		BertGenerationDecoder(snake_case__)
          model.to(snake_case__)
          model.eval()
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		model(snake_case__						,	attention_mask=snake_case__						,	labels=snake_case__)
          self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.seq_length, self.vocab_size))



      def _SCREAMING_SNAKE_CASE (						self     :  Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		self.prepare_config_and_inputs()
          SCREAMING_SNAKE_CASE_      :  Dict    						=		{"input_ids": input_ids, "attention_mask": input_mask}
          return config, inputs_dict




@require_torch
class 							lowerCAmelCase__				(			A_				,		A_				,		A_				,		unittest.TestCase ):






      '''simple docstring'''

      __UpperCamelCase	 =  (BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
      __UpperCamelCase	 =  (BertGenerationDecoder,) if is_torch_available() else ()
      __UpperCamelCase	 =  (
          {"feature-extraction": BertGenerationEncoder, "text-generation": BertGenerationDecoder}
          if is_torch_available()
          else {}
      )
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Dict    						=		BertGenerationEncoderTester(self)
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		ConfigTester(self						,	config_class=snake_case__						,	hidden_size=37)
      def _SCREAMING_SNAKE_CASE (						self     :  int):




          '''simple docstring'''


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




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Tuple    						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*snake_case__)
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		self.model_tester.prepare_config_and_inputs()
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		"bert"
          self.model_tester.create_and_check_model(snake_case__						,	snake_case__						,	snake_case__						,	snake_case__)
      def _SCREAMING_SNAKE_CASE (						self     :  Union[str, Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[str]    						=		self.model_tester.prepare_config_and_inputs_for_decoder()
          self.model_tester.create_and_check_model_as_decoder(*snake_case__)
      def _SCREAMING_SNAKE_CASE (						self     :  Any):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Tuple    						=		self.model_tester.prepare_config_and_inputs_for_decoder()
          self.model_tester.create_and_check_decoder_model_past_large_inputs(*snake_case__)
      def _SCREAMING_SNAKE_CASE (						self     :  Dict):




          '''simple docstring'''


          (
              SCREAMING_SNAKE_CASE_					
          )      :  Union[str, Any]    						=		self.model_tester.prepare_config_and_inputs_for_decoder()

          SCREAMING_SNAKE_CASE_      :  Dict    						=		None

          self.model_tester.create_and_check_model_as_decoder(
              snake_case__						,	snake_case__						,	snake_case__						,	snake_case__						,	snake_case__						,	snake_case__						,	)
      def _SCREAMING_SNAKE_CASE (						self     :  int):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		self.model_tester.prepare_config_and_inputs_for_decoder()
          self.model_tester.create_and_check_for_causal_lm(*snake_case__)



      @slow
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder''')
          self.assertIsNotNone(snake_case__)




@require_torch
class 							lowerCAmelCase__				(			unittest.TestCase ):






      '''simple docstring'''

      @slow
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Tuple    						=		BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder''')
          SCREAMING_SNAKE_CASE_      :  Dict    						=		torch.tensor([[101, 7592, 1010, 2026, 3899, 2003, 10140, 102]])
          with torch.no_grad():
              SCREAMING_SNAKE_CASE_      :  List[str]    						=		model(snake_case__)[0]
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		torch.Size([1, 8, 1024])
          self.assertEqual(output.shape						,	snake_case__)
          SCREAMING_SNAKE_CASE_      :  int    						=		torch.tensor(
              [[[0.17_75, 0.00_83, -0.03_21], [1.60_02, 0.12_87, 0.39_12], [2.14_73, 0.57_91, 0.60_66]]])
          self.assertTrue(torch.allclose(output[:, :3, :3]						,	snake_case__						,	atol=1e-4))




@require_torch
class 							lowerCAmelCase__				(			unittest.TestCase ):






      '''simple docstring'''

      @slow
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		BertGenerationDecoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder''')
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		torch.tensor([[101, 7592, 1010, 2026, 3899, 2003, 10140, 102]])
          with torch.no_grad():
              SCREAMING_SNAKE_CASE_      :  List[Any]    						=		model(snake_case__)[0]
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		torch.Size([1, 8, 50358])
          self.assertEqual(output.shape						,	snake_case__)
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		torch.tensor(
              [[[-0.57_88, -2.59_94, -3.70_54], [0.04_38, 4.79_97, 1.87_95], [1.58_62, 6.64_09, 4.46_38]]])
          self.assertTrue(torch.allclose(output[:, :3, :3]						,	snake_case__						,	atol=1e-4))




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


if is_sentencepiece_available():
   from ..ta.tokenization_ta import TaTokenizer
else:
   from ...utils.dummy_sentencepiece_objects import TaTokenizer

__lowerCamelCase							=       TaTokenizer

if is_tokenizers_available():
   from ..ta.tokenization_ta_fast import TaTokenizerFast
else:
   from ...utils.dummy_tokenizers_objects import TaTokenizerFast

__lowerCamelCase							=       TaTokenizerFast

__lowerCamelCase							=       {"""configuration_mt5""": ["""MT5Config""", """MT5OnnxConfig"""]}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __lowerCamelCase							=       [
       """MT5EncoderModel""",
       """MT5ForConditionalGeneration""",
       """MT5ForQuestionAnswering""",
       """MT5Model""",
       """MT5PreTrainedModel""",
       """MT5Stack""",
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __lowerCamelCase							=       ["""TFMT5EncoderModel""", """TFMT5ForConditionalGeneration""", """TFMT5Model"""]

try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __lowerCamelCase							=       ["""FlaxMT5EncoderModel""", """FlaxMT5ForConditionalGeneration""", """FlaxMT5Model"""]


if TYPE_CHECKING:
   from .configuration_mta import MTaConfig, MTaOnnxConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_mta import (
          MTaEncoderModel,
          MTaForConditionalGeneration,
          MTaForQuestionAnswering,
          MTaModel,
          MTaPreTrainedModel,
          MTaStack,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_mta import TFMTaEncoderModel, TFMTaForConditionalGeneration, TFMTaModel

   try:
      if not is_flax_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_flax_mta import FlaxMTaEncoderModel, FlaxMTaForConditionalGeneration, FlaxMTaModel

else:
   import sys

   __lowerCamelCase							=       _LazyModule(
       __name__,
       globals()["""__file__"""],
       _import_structure,
       extra_objects={"""MT5Tokenizer""": MTaTokenizer, """MT5TokenizerFast""": MTaTokenizerFast},
       module_spec=__spec__,
   )





