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



	"""simple docstring"""




	def __init__(			self    :     List[Any]			,				SCREAMING_SNAKE_CASE_    :     Dict			,				SCREAMING_SNAKE_CASE_    :     Any=1_3			,				SCREAMING_SNAKE_CASE_    :     Union[str, Any]=7			,				SCREAMING_SNAKE_CASE_    :     Tuple=True			,				SCREAMING_SNAKE_CASE_    :     Tuple=True			,				SCREAMING_SNAKE_CASE_    :     str=9_9			,				SCREAMING_SNAKE_CASE_    :     Dict=3_2			,				SCREAMING_SNAKE_CASE_    :     Tuple=5			,				SCREAMING_SNAKE_CASE_    :     List[Any]=4			,				SCREAMING_SNAKE_CASE_    :     Optional[Any]=3_7			,				SCREAMING_SNAKE_CASE_    :     str="gelu"			,				SCREAMING_SNAKE_CASE_    :     List[str]=0.1			,				SCREAMING_SNAKE_CASE_    :     List[Any]=0.1			,				SCREAMING_SNAKE_CASE_    :     str=5_0			,				SCREAMING_SNAKE_CASE_    :     List[str]=0.02			,				SCREAMING_SNAKE_CASE_    :     int=True			,				SCREAMING_SNAKE_CASE_    :     Any=None			,				)			->	Tuple:
					lowercase_	    =   parent
					lowercase_	    =   batch_size
					lowercase_	    =   seq_length
					lowercase_	    =   is_training
					lowercase_	    =   use_input_mask
					lowercase_	    =   vocab_size
					lowercase_	    =   hidden_size
					lowercase_	    =   num_hidden_layers
					lowercase_	    =   num_attention_heads
					lowercase_	    =   intermediate_size
					lowercase_	    =   hidden_act
					lowercase_	    =   hidden_dropout_prob
					lowercase_	    =   attention_probs_dropout_prob
					lowercase_	    =   max_position_embeddings
					lowercase_	    =   initializer_range
					lowercase_	    =   use_labels
					lowercase_	    =   scope
	def 		_lowercase    (			self    :     List[str]			)			->	int:
					lowercase_	    =   ids_tensor([self.batch_size, self.seq_length]			,				self.vocab_size			)

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

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

					lowercase_	    =   self.get_config()

					return config, input_ids, input_mask, token_labels
	def 		_lowercase    (			self    :     List[Any]			)			->	Union[str, Any]:
					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=SCREAMING_SNAKE_CASE_			,				initializer_range=self.initializer_range			,				)
	def 		_lowercase    (			self    :     List[Any]			)			->	Tuple:
					(
					    (
					    lowercase_
					)    ,      (
					    lowercase_
					)    ,      (
					    lowercase_
					)    ,      (
					    lowercase_
					)    ,      
					)	    =   self.prepare_config_and_inputs()

					lowercase_	    =   True
					lowercase_	    =   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
					lowercase_	    =   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 		_lowercase    (			self    :     Tuple			,				SCREAMING_SNAKE_CASE_    :     Union[str, Any]			,				SCREAMING_SNAKE_CASE_    :     List[Any]			,				SCREAMING_SNAKE_CASE_    :     List[str]			,				SCREAMING_SNAKE_CASE_    :     Union[str, Any]			,				**SCREAMING_SNAKE_CASE_    :     str			,				)			->	Any:
					lowercase_	    =   BertGenerationEncoder(config=SCREAMING_SNAKE_CASE_			)
					model.to(SCREAMING_SNAKE_CASE_			)
					model.eval()
					lowercase_	    =   model(SCREAMING_SNAKE_CASE_			,				attention_mask=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   model(SCREAMING_SNAKE_CASE_			)
					self.parent.assertEqual(result.last_hidden_state.shape			,				(self.batch_size, self.seq_length, self.hidden_size)			)
	def 		_lowercase    (			self    :     List[str]			,				SCREAMING_SNAKE_CASE_    :     Union[str, Any]			,				SCREAMING_SNAKE_CASE_    :     List[Any]			,				SCREAMING_SNAKE_CASE_    :     Dict			,				SCREAMING_SNAKE_CASE_    :     Optional[Any]			,				SCREAMING_SNAKE_CASE_    :     List[str]			,				SCREAMING_SNAKE_CASE_    :     Union[str, Any]			,				**SCREAMING_SNAKE_CASE_    :     str			,				)			->	Union[str, Any]:
					lowercase_	    =   True
					lowercase_	    =   BertGenerationEncoder(config=SCREAMING_SNAKE_CASE_			)
					model.to(SCREAMING_SNAKE_CASE_			)
					model.eval()
					lowercase_	    =   model(
					    SCREAMING_SNAKE_CASE_			,				attention_mask=SCREAMING_SNAKE_CASE_			,				encoder_hidden_states=SCREAMING_SNAKE_CASE_			,				encoder_attention_mask=SCREAMING_SNAKE_CASE_			,				)
					lowercase_	    =   model(
					    SCREAMING_SNAKE_CASE_			,				attention_mask=SCREAMING_SNAKE_CASE_			,				encoder_hidden_states=SCREAMING_SNAKE_CASE_			,				)
					self.parent.assertEqual(result.last_hidden_state.shape			,				(self.batch_size, self.seq_length, self.hidden_size)			)
	def 		_lowercase    (			self    :     Optional[int]			,				SCREAMING_SNAKE_CASE_    :     Optional[int]			,				SCREAMING_SNAKE_CASE_    :     Optional[int]			,				SCREAMING_SNAKE_CASE_    :     List[str]			,				SCREAMING_SNAKE_CASE_    :     List[Any]			,				SCREAMING_SNAKE_CASE_    :     Union[str, Any]			,				SCREAMING_SNAKE_CASE_    :     Any			,				**SCREAMING_SNAKE_CASE_    :     List[str]			,				)			->	str:
					lowercase_	    =   True
					lowercase_	    =   True
					lowercase_	    =   BertGenerationDecoder(config=SCREAMING_SNAKE_CASE_			).to(SCREAMING_SNAKE_CASE_			).eval()

					# first forward pass
					lowercase_	    =   model(
					    SCREAMING_SNAKE_CASE_			,				attention_mask=SCREAMING_SNAKE_CASE_			,				encoder_hidden_states=SCREAMING_SNAKE_CASE_			,				encoder_attention_mask=SCREAMING_SNAKE_CASE_			,				use_cache=SCREAMING_SNAKE_CASE_			,				)
					lowercase_	    =   outputs.past_key_values

					# create hypothetical multiple next token and extent to next_input_ids
					lowercase_	    =   ids_tensor((self.batch_size, 3)			,				config.vocab_size			)
					lowercase_	    =   ids_tensor((self.batch_size, 3)			,				vocab_size=2			)

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

					lowercase_	    =   model(
					    SCREAMING_SNAKE_CASE_			,				attention_mask=SCREAMING_SNAKE_CASE_			,				encoder_hidden_states=SCREAMING_SNAKE_CASE_			,				encoder_attention_mask=SCREAMING_SNAKE_CASE_			,				output_hidden_states=SCREAMING_SNAKE_CASE_			,				)['''hidden_states'''][0]
					lowercase_	    =   model(
					    SCREAMING_SNAKE_CASE_			,				attention_mask=SCREAMING_SNAKE_CASE_			,				encoder_hidden_states=SCREAMING_SNAKE_CASE_			,				encoder_attention_mask=SCREAMING_SNAKE_CASE_			,				past_key_values=SCREAMING_SNAKE_CASE_			,				output_hidden_states=SCREAMING_SNAKE_CASE_			,				)['''hidden_states'''][0]

					# select random slice
					lowercase_	    =   ids_tensor((1,)			,				output_from_past.shape[-1]			).item()
					lowercase_	    =   output_from_no_past[:, -3:, random_slice_idx].detach()
					lowercase_	    =   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(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				atol=1e-3			)			)
	def 		_lowercase    (			self    :     Union[str, Any]			,				SCREAMING_SNAKE_CASE_    :     List[Any]			,				SCREAMING_SNAKE_CASE_    :     Tuple			,				SCREAMING_SNAKE_CASE_    :     int			,				SCREAMING_SNAKE_CASE_    :     Dict			,				*SCREAMING_SNAKE_CASE_    :     List[str]			,				)			->	Any:
					lowercase_	    =   BertGenerationDecoder(SCREAMING_SNAKE_CASE_			)
					model.to(SCREAMING_SNAKE_CASE_			)
					model.eval()
					lowercase_	    =   model(SCREAMING_SNAKE_CASE_			,				attention_mask=SCREAMING_SNAKE_CASE_			,				labels=SCREAMING_SNAKE_CASE_			)
					self.parent.assertEqual(result.logits.shape			,				(self.batch_size, self.seq_length, self.vocab_size)			)




	def 		_lowercase    (			self    :     str			)			->	List[str]:
					lowercase_    ,      lowercase_    ,      lowercase_    ,      lowercase_	    =   self.prepare_config_and_inputs()
					lowercase_	    =   {'''input_ids''': input_ids, '''attention_mask''': input_mask}
					return config, inputs_dict

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



	"""simple docstring"""




	a     :List[Any]       	=      (BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
	a     :Optional[int]       	=      (BertGenerationDecoder,) if is_torch_available() else ()
	a     :List[Any]       	=      (
	    {'feature-extraction': BertGenerationEncoder, 'text-generation': BertGenerationDecoder}
	    if is_torch_available()
	    else {}
	)
	def 		_lowercase    (			self    :     int			)			->	str:
					lowercase_	    =   BertGenerationEncoderTester(self			)
					lowercase_	    =   ConfigTester(self			,				config_class=SCREAMING_SNAKE_CASE_			,				hidden_size=3_7			)
	def 		_lowercase    (			self    :     Any			)			->	Optional[int]:
					self.config_tester.run_common_tests()
	def 		_lowercase    (			self    :     Union[str, Any]			)			->	Optional[int]:
					lowercase_	    =   self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_			)
	def 		_lowercase    (			self    :     Tuple			)			->	Dict:
					lowercase_    ,      lowercase_    ,      lowercase_    ,      lowercase_	    =   self.model_tester.prepare_config_and_inputs()
					lowercase_	    =   '''bert'''
					self.model_tester.create_and_check_model(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
	def 		_lowercase    (			self    :     Union[str, Any]			)			->	Union[str, Any]:
					lowercase_	    =   self.model_tester.prepare_config_and_inputs_for_decoder()
					self.model_tester.create_and_check_model_as_decoder(*SCREAMING_SNAKE_CASE_			)
	def 		_lowercase    (			self    :     str			)			->	Tuple:
					lowercase_	    =   self.model_tester.prepare_config_and_inputs_for_decoder()
					self.model_tester.create_and_check_decoder_model_past_large_inputs(*SCREAMING_SNAKE_CASE_			)
	def 		_lowercase    (			self    :     Dict			)			->	Union[str, Any]:
					# This regression test was failing with PyTorch < 1.3
					(
					    (
					    lowercase_
					)    ,      (
					    lowercase_
					)    ,      (
					    lowercase_
					)    ,      (
					    lowercase_
					)    ,      (
					    lowercase_
					)    ,      (
					    lowercase_
					)    ,      
					)	    =   self.model_tester.prepare_config_and_inputs_for_decoder()

					lowercase_	    =   None

					self.model_tester.create_and_check_model_as_decoder(
					    SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				)
	def 		_lowercase    (			self    :     Union[str, Any]			)			->	Optional[Any]:
					lowercase_	    =   self.model_tester.prepare_config_and_inputs_for_decoder()
					self.model_tester.create_and_check_for_causal_lm(*SCREAMING_SNAKE_CASE_			)




	@slow
	def 		_lowercase    (			self    :     List[Any]			)			->	Union[str, Any]:
					lowercase_	    =   BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''			)
					self.assertIsNotNone(SCREAMING_SNAKE_CASE_			)

@require_torch
class 					lowercase__(							unittest.TestCase				):



	"""simple docstring"""




	@slow
	def 		_lowercase    (			self    :     Union[str, Any]			)			->	Union[str, Any]:
					lowercase_	    =   BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''			)
					lowercase_	    =   torch.tensor([[1_0_1, 7_5_9_2, 1_0_1_0, 2_0_2_6, 3_8_9_9, 2_0_0_3, 1_0_1_4_0, 1_0_2]]			)
					with torch.no_grad():
									lowercase_	    =   model(SCREAMING_SNAKE_CASE_			)[0]
					lowercase_	    =   torch.Size([1, 8, 1_0_2_4]			)
					self.assertEqual(output.shape			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   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]			,				SCREAMING_SNAKE_CASE_			,				atol=1e-4			)			)

@require_torch
class 					lowercase__(							unittest.TestCase				):



	"""simple docstring"""




	@slow
	def 		_lowercase    (			self    :     int			)			->	Tuple:
					lowercase_	    =   BertGenerationDecoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''			)
					lowercase_	    =   torch.tensor([[1_0_1, 7_5_9_2, 1_0_1_0, 2_0_2_6, 3_8_9_9, 2_0_0_3, 1_0_1_4_0, 1_0_2]]			)
					with torch.no_grad():
									lowercase_	    =   model(SCREAMING_SNAKE_CASE_			)[0]
					lowercase_	    =   torch.Size([1, 8, 5_0_3_5_8]			)
					self.assertEqual(output.shape			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   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]			,				SCREAMING_SNAKE_CASE_			,				atol=1e-4			)			)






import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__a         =   logging.get_logger(__name__)
class 					lowercase__(							UpperCAmelCase				):



	"""simple docstring"""




	a     :List[str]       	=      'encoder-decoder'
	a     :Any       	=      True
	def __init__(			self    :     List[Any]			,				**SCREAMING_SNAKE_CASE_    :     List[Any]			)			->	Optional[int]:
					super().__init__(**SCREAMING_SNAKE_CASE_			)
					assert (
					    "encoder" in kwargs and "decoder" in kwargs
					), "Config has to be initialized with encoder and decoder config"
					lowercase_	    =   kwargs.pop('''encoder'''			)
					lowercase_	    =   encoder_config.pop('''model_type'''			)
					lowercase_	    =   kwargs.pop('''decoder'''			)
					lowercase_	    =   decoder_config.pop('''model_type'''			)

					from ..auto.configuration_auto import AutoConfig

					lowercase_	    =   AutoConfig.for_model(SCREAMING_SNAKE_CASE_			,				**SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   AutoConfig.for_model(SCREAMING_SNAKE_CASE_			,				**SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   True
	@classmethod
	def 		_lowercase    (			cls    :     Optional[int]			,				SCREAMING_SNAKE_CASE_    :     PretrainedConfig			,				SCREAMING_SNAKE_CASE_    :     PretrainedConfig			,				**SCREAMING_SNAKE_CASE_    :     Optional[int]			)			->	PretrainedConfig:
					logger.info('''Set `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config'''			)
					lowercase_	    =   True
					lowercase_	    =   True

					return cls(encoder=encoder_config.to_dict()			,				decoder=decoder_config.to_dict()			,				**SCREAMING_SNAKE_CASE_			)




	def 		_lowercase    (			self    :     Tuple			)			->	List[Any]:
					lowercase_	    =   copy.deepcopy(self.__dict__			)
					lowercase_	    =   self.encoder.to_dict()
					lowercase_	    =   self.decoder.to_dict()
					lowercase_	    =   self.__class__.model_type
					return output




