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 , __lowercase , __lowercase=13 , __lowercase=7 , __lowercase=True , __lowercase=True , __lowercase=99 , __lowercase=32 , __lowercase=5 , __lowercase=4 , __lowercase=37 , __lowercase="gelu" , __lowercase=0.1 , __lowercase=0.1 , __lowercase=50 , __lowercase=0.02 , __lowercase=True , __lowercase=None , )       ->	Dict:
        __UpperCamelCase     :Union[str, Any]			  =		parent
        __UpperCamelCase     :List[Any]			  =		batch_size
        __UpperCamelCase     :Tuple			  =		seq_length
        __UpperCamelCase     :List[Any]			  =		is_training
        __UpperCamelCase     :str			  =		use_input_mask
        __UpperCamelCase     :str			  =		vocab_size
        __UpperCamelCase     :str			  =		hidden_size
        __UpperCamelCase     :Tuple			  =		num_hidden_layers
        __UpperCamelCase     :int			  =		num_attention_heads
        __UpperCamelCase     :Union[str, Any]			  =		intermediate_size
        __UpperCamelCase     :Tuple			  =		hidden_act
        __UpperCamelCase     :List[Any]			  =		hidden_dropout_prob
        __UpperCamelCase     :Any			  =		attention_probs_dropout_prob
        __UpperCamelCase     :Any			  =		max_position_embeddings
        __UpperCamelCase     :int			  =		initializer_range
        __UpperCamelCase     :Tuple			  =		use_labels
        __UpperCamelCase     :Any			  =		scope







    def 			UpperCamelCase__   (	self)       ->	Any:
        __UpperCamelCase     :List[Any]			  =		ids_tensor([self.batch_size, self.seq_length] , self.vocab_size)

        __UpperCamelCase     :Any			  =		None
        if self.use_input_mask:
            __UpperCamelCase     :Union[str, Any]			  =		random_attention_mask([self.batch_size, self.seq_length])

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

        __UpperCamelCase     :int			  =		self.get_config()

        return config, input_ids, input_mask, token_labels







    def 			UpperCamelCase__   (	self)       ->	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=__lowercase , initializer_range=self.initializer_range , )







    def 			UpperCamelCase__   (	self)       ->	int:
        (
            (
            __UpperCamelCase
        )  ,       (
            __UpperCamelCase
        )  ,       (
            __UpperCamelCase
        )  ,       (
            __UpperCamelCase
        )  ,       
        )     :Dict			  =		self.prepare_config_and_inputs()

        __UpperCamelCase     :Dict			  =		True
        __UpperCamelCase     :str			  =		floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
        __UpperCamelCase     :str			  =		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 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , **__lowercase , )       ->	Dict:
        __UpperCamelCase     :Tuple			  =		BertGenerationEncoder(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase     :Tuple			  =		model(__lowercase , attention_mask=__lowercase)
        __UpperCamelCase     :Dict			  =		model(__lowercase)
        self.parent.assertEqual(result.last_hidden_state.shape , (self.batch_size, self.seq_length, self.hidden_size))







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , **__lowercase , )       ->	int:
        __UpperCamelCase     :str			  =		True
        __UpperCamelCase     :Any			  =		BertGenerationEncoder(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase     :Union[str, Any]			  =		model(
            __lowercase , attention_mask=__lowercase , encoder_hidden_states=__lowercase , encoder_attention_mask=__lowercase , )
        __UpperCamelCase     :Optional[Any]			  =		model(
            __lowercase , attention_mask=__lowercase , encoder_hidden_states=__lowercase , )
        self.parent.assertEqual(result.last_hidden_state.shape , (self.batch_size, self.seq_length, self.hidden_size))







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , **__lowercase , )       ->	Optional[int]:
        __UpperCamelCase     :Any			  =		True
        __UpperCamelCase     :Union[str, Any]			  =		True
        __UpperCamelCase     :Optional[int]			  =		BertGenerationDecoder(config=__lowercase).to(__lowercase).eval()

        # first forward pass
        __UpperCamelCase     :int			  =		model(
            __lowercase , attention_mask=__lowercase , encoder_hidden_states=__lowercase , encoder_attention_mask=__lowercase , use_cache=__lowercase , )
        __UpperCamelCase     :List[Any]			  =		outputs.past_key_values

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

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

        __UpperCamelCase     :Optional[Any]			  =		model(
            __lowercase , attention_mask=__lowercase , encoder_hidden_states=__lowercase , encoder_attention_mask=__lowercase , output_hidden_states=__lowercase , )['''hidden_states'''][0]
        __UpperCamelCase     :str			  =		model(
            __lowercase , attention_mask=__lowercase , encoder_hidden_states=__lowercase , encoder_attention_mask=__lowercase , past_key_values=__lowercase , output_hidden_states=__lowercase , )['''hidden_states'''][0]

        # select random slice
        __UpperCamelCase     :Union[str, Any]			  =		ids_tensor((1,) , output_from_past.shape[-1]).item()
        __UpperCamelCase     :Optional[int]			  =		output_from_no_past[:, -3:, random_slice_idx].detach()
        __UpperCamelCase     :Tuple			  =		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(__lowercase , __lowercase , atol=1E-3))







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , *__lowercase , )       ->	str:
        __UpperCamelCase     :str			  =		BertGenerationDecoder(__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase     :List[str]			  =		model(__lowercase , attention_mask=__lowercase , labels=__lowercase)
        self.parent.assertEqual(result.logits.shape , (self.batch_size, self.seq_length, self.vocab_size))







    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        __UpperCamelCase  ,       __UpperCamelCase  ,       __UpperCamelCase  ,       __UpperCamelCase     :Union[str, Any]			  =		self.prepare_config_and_inputs()
        __UpperCamelCase     :Union[str, Any]			  =		{'''input_ids''': input_ids, '''attention_mask''': input_mask}
        return config, inputs_dict


@require_torch
class 	lowerCamelCase_     (   UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_						,	unittest.TestCase ):






    '''simple docstring'''






    a__	:						List[str]						=      (BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
    a__	:						Union[str, Any]						=      (BertGenerationDecoder,) if is_torch_available() else ()
    a__	:						Optional[Any]						=      (
        {"""feature-extraction""": BertGenerationEncoder, """text-generation""": BertGenerationDecoder}
        if is_torch_available()
        else {}
    )







    def 			UpperCamelCase__   (	self)       ->	str:
        __UpperCamelCase     :List[Any]			  =		BertGenerationEncoderTester(self)
        __UpperCamelCase     :Optional[Any]			  =		ConfigTester(self , config_class=__lowercase , hidden_size=37)







    def 			UpperCamelCase__   (	self)       ->	int:
        self.config_tester.run_common_tests()







    def 			UpperCamelCase__   (	self)       ->	str:
        __UpperCamelCase     :List[Any]			  =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowercase)







    def 			UpperCamelCase__   (	self)       ->	List[str]:
        __UpperCamelCase  ,       __UpperCamelCase  ,       __UpperCamelCase  ,       __UpperCamelCase     :Any			  =		self.model_tester.prepare_config_and_inputs()
        __UpperCamelCase     :Any			  =		'''bert'''
        self.model_tester.create_and_check_model(__lowercase , __lowercase , __lowercase , __lowercase)







    def 			UpperCamelCase__   (	self)       ->	Optional[int]:
        __UpperCamelCase     :str			  =		self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_model_as_decoder(*__lowercase)







    def 			UpperCamelCase__   (	self)       ->	Optional[Any]:
        __UpperCamelCase     :int			  =		self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_decoder_model_past_large_inputs(*__lowercase)







    def 			UpperCamelCase__   (	self)       ->	List[Any]:
        # This regression test was failing with PyTorch < 1.3
        (
            (
            __UpperCamelCase
        )  ,       (
            __UpperCamelCase
        )  ,       (
            __UpperCamelCase
        )  ,       (
            __UpperCamelCase
        )  ,       (
            __UpperCamelCase
        )  ,       (
            __UpperCamelCase
        )  ,       
        )     :int			  =		self.model_tester.prepare_config_and_inputs_for_decoder()

        __UpperCamelCase     :List[Any]			  =		None

        self.model_tester.create_and_check_model_as_decoder(
            __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , )







    def 			UpperCamelCase__   (	self)       ->	Dict:
        __UpperCamelCase     :Dict			  =		self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_for_causal_lm(*__lowercase)







    @slow
    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        __UpperCamelCase     :Dict			  =		BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder''')
        self.assertIsNotNone(__lowercase)


@require_torch
class 	lowerCamelCase_     (   unittest.TestCase ):






    '''simple docstring'''






    @slow
    def 			UpperCamelCase__   (	self)       ->	Any:
        __UpperCamelCase     :Any			  =		BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder''')
        __UpperCamelCase     :Optional[Any]			  =		torch.tensor([[101, 7_592, 1_010, 2_026, 3_899, 2_003, 10_140, 102]])
        with torch.no_grad():
            __UpperCamelCase     :Any			  =		model(__lowercase)[0]
        __UpperCamelCase     :Optional[Any]			  =		torch.Size([1, 8, 1_024])
        self.assertEqual(output.shape , __lowercase)
        __UpperCamelCase     :str			  =		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] , __lowercase , atol=1E-4))


@require_torch
class 	lowerCamelCase_     (   unittest.TestCase ):






    '''simple docstring'''






    @slow
    def 			UpperCamelCase__   (	self)       ->	List[Any]:
        __UpperCamelCase     :List[str]			  =		BertGenerationDecoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder''')
        __UpperCamelCase     :Tuple			  =		torch.tensor([[101, 7_592, 1_010, 2_026, 3_899, 2_003, 10_140, 102]])
        with torch.no_grad():
            __UpperCamelCase     :int			  =		model(__lowercase)[0]
        __UpperCamelCase     :Optional[int]			  =		torch.Size([1, 8, 50_358])
        self.assertEqual(output.shape , __lowercase)
        __UpperCamelCase     :Optional[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] , __lowercase , atol=1E-4))






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


__lowercase							 =					logging.get_logger(__name__)

__lowercase							 =					{'''ctrl''': '''https://huggingface.co/ctrl/resolve/main/config.json'''}


class 	lowerCamelCase_     (   UpperCAmelCase_ ):






    '''simple docstring'''






    a__	:						str						=      """ctrl"""
    a__	:						Dict						=      ["""past_key_values"""]
    a__	:						Tuple						=      {
        """max_position_embeddings""": """n_positions""",
        """hidden_size""": """n_embd""",
        """num_attention_heads""": """n_head""",
        """num_hidden_layers""": """n_layer""",
    }







    def __init__(	self , __lowercase=246_534 , __lowercase=256 , __lowercase=1_280 , __lowercase=8_192 , __lowercase=48 , __lowercase=16 , __lowercase=0.1 , __lowercase=0.1 , __lowercase=1E-6 , __lowercase=0.02 , __lowercase=True , **__lowercase , )       ->	List[Any]:
        __UpperCamelCase     :List[str]			  =		vocab_size
        __UpperCamelCase     :Optional[Any]			  =		n_positions
        __UpperCamelCase     :Dict			  =		n_embd
        __UpperCamelCase     :Dict			  =		n_layer
        __UpperCamelCase     :List[Any]			  =		n_head
        __UpperCamelCase     :int			  =		dff
        __UpperCamelCase     :Union[str, Any]			  =		resid_pdrop
        __UpperCamelCase     :Optional[int]			  =		embd_pdrop
        __UpperCamelCase     :List[Any]			  =		layer_norm_epsilon
        __UpperCamelCase     :Dict			  =		initializer_range

        __UpperCamelCase     :Any			  =		use_cache

        super().__init__(**__lowercase)





