


'''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 				UpperCAmelCase			:


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


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[int] ):
    UpperCAmelCase__	:Optional[Any]      =      ids_tensor([self.batch_size, self.seq_length]			,       self.vocab_size )

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

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

    UpperCAmelCase__	:str      =      self.get_config()

    return config, input_ids, input_mask, token_labels


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


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[Any] ):
    (
        (
        UpperCAmelCase__
    )   ,						(
        UpperCAmelCase__
    )   ,						(
        UpperCAmelCase__
    )   ,						(
        UpperCAmelCase__
    )   ,						
    )	:List[Any]      =      self.prepare_config_and_inputs()

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


 def        __SCREAMING_SNAKE_CASE					(      self							:	Tuple			,       __lowerCamelCase							:	Any			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	Union[str, Any]			,       __lowerCamelCase							:	List[Any]			,       __lowerCamelCase							:	Union[str, Any]			,       __lowerCamelCase							:	List[Any]			,       **__lowerCamelCase							:	Any			,       ):
    UpperCAmelCase__	:str      =      True
    UpperCAmelCase__	:str      =      BertGenerationEncoder(config=__lowerCamelCase )
    model.to(__lowerCamelCase )
    model.eval()
    UpperCAmelCase__	:Dict      =      model(
        __lowerCamelCase			,       attention_mask=__lowerCamelCase			,       encoder_hidden_states=__lowerCamelCase			,       encoder_attention_mask=__lowerCamelCase			,       )
    UpperCAmelCase__	:Optional[int]      =      model(
        __lowerCamelCase			,       attention_mask=__lowerCamelCase			,       encoder_hidden_states=__lowerCamelCase			,       )
    self.parent.assertEqual(result.last_hidden_state.shape			,       (self.batch_size, self.seq_length, self.hidden_size) )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Union[str, Any]			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	Tuple			,       __lowerCamelCase							:	Tuple			,       __lowerCamelCase							:	Dict			,       __lowerCamelCase							:	str			,       __lowerCamelCase							:	Any			,       **__lowerCamelCase							:	str			,       ):
    UpperCAmelCase__	:Optional[int]      =      True
    UpperCAmelCase__	:List[str]      =      True
    UpperCAmelCase__	:int      =      BertGenerationDecoder(config=__lowerCamelCase ).to(__lowerCamelCase ).eval()

    # first forward pass
    UpperCAmelCase__	:int      =      model(
        __lowerCamelCase			,       attention_mask=__lowerCamelCase			,       encoder_hidden_states=__lowerCamelCase			,       encoder_attention_mask=__lowerCamelCase			,       use_cache=__lowerCamelCase			,       )
    UpperCAmelCase__	:Union[str, Any]      =      outputs.past_key_values

    # create hypothetical multiple next token and extent to next_input_ids
    UpperCAmelCase__	:Dict      =      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__	:str      =      torch.cat([input_ids, next_tokens]			,       dim=-1 )
    UpperCAmelCase__	:Optional[int]      =      torch.cat([input_mask, next_mask]			,       dim=-1 )

    UpperCAmelCase__	:Optional[int]      =      model(
        __lowerCamelCase			,       attention_mask=__lowerCamelCase			,       encoder_hidden_states=__lowerCamelCase			,       encoder_attention_mask=__lowerCamelCase			,       output_hidden_states=__lowerCamelCase			,       )['''hidden_states'''][0]
    UpperCAmelCase__	:Any      =      model(
        __lowerCamelCase			,       attention_mask=__lowerCamelCase			,       encoder_hidden_states=__lowerCamelCase			,       encoder_attention_mask=__lowerCamelCase			,       past_key_values=__lowerCamelCase			,       output_hidden_states=__lowerCamelCase			,       )['''hidden_states'''][0]

    # select random slice
    UpperCAmelCase__	:Optional[int]      =      ids_tensor((1,)			,       output_from_past.shape[-1] ).item()
    UpperCAmelCase__	:Tuple      =      output_from_no_past[:, -3:, random_slice_idx].detach()
    UpperCAmelCase__	:int      =      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(__lowerCamelCase			,       __lowerCamelCase			,       atol=1e-3 ) )


 def        __SCREAMING_SNAKE_CASE					(      self							:	int			,       __lowerCamelCase							:	Optional[Any]			,       __lowerCamelCase							:	Optional[Any]			,       __lowerCamelCase							:	Optional[int]			,       __lowerCamelCase							:	Tuple			,       *__lowerCamelCase							:	Tuple			,       ):
    UpperCAmelCase__	:Union[str, Any]      =      BertGenerationDecoder(__lowerCamelCase )
    model.to(__lowerCamelCase )
    model.eval()
    UpperCAmelCase__	:Union[str, Any]      =      model(__lowerCamelCase			,       attention_mask=__lowerCamelCase			,       labels=__lowerCamelCase )
    self.parent.assertEqual(result.logits.shape			,       (self.batch_size, self.seq_length, self.vocab_size) )





 def        __SCREAMING_SNAKE_CASE					(      self							:	str ):
    UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__	:Any      =      self.prepare_config_and_inputs()
    UpperCAmelCase__	:List[Any]      =      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
    return config, inputs_dict




@require_torch
class 				UpperCAmelCase			(			_snake_case						, _snake_case						, _snake_case						, unittest.TestCase						):
 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							:	int ):
    UpperCAmelCase__	:Dict      =      BertGenerationEncoderTester(self )
    UpperCAmelCase__	:int      =      ConfigTester(self			,       config_class=__lowerCamelCase			,       hidden_size=3_7 )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Union[str, Any] ):
    self.config_tester.run_common_tests()


 def        __SCREAMING_SNAKE_CASE					(      self							:	Tuple ):
    UpperCAmelCase__	:Optional[Any]      =      self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Union[str, Any] ):
    UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__   ,						UpperCAmelCase__	:Union[str, Any]      =      self.model_tester.prepare_config_and_inputs()
    UpperCAmelCase__	:Optional[Any]      =      '''bert'''
    self.model_tester.create_and_check_model(__lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Any ):
    UpperCAmelCase__	:Optional[int]      =      self.model_tester.prepare_config_and_inputs_for_decoder()
    self.model_tester.create_and_check_model_as_decoder(*__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[int] ):
    UpperCAmelCase__	:str      =      self.model_tester.prepare_config_and_inputs_for_decoder()
    self.model_tester.create_and_check_decoder_model_past_large_inputs(*__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	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__	:int      =      None

    self.model_tester.create_and_check_model_as_decoder(
        __lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase			,       )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[int] ):
    UpperCAmelCase__	:List[str]      =      self.model_tester.prepare_config_and_inputs_for_decoder()
    self.model_tester.create_and_check_for_causal_lm(*__lowerCamelCase )





 @slow
 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[int] ):
    UpperCAmelCase__	:Union[str, Any]      =      BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder''' )
    self.assertIsNotNone(__lowerCamelCase )




@require_torch
class 				UpperCAmelCase			(			unittest.TestCase						):


 @slow
 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str] ):
    UpperCAmelCase__	:List[str]      =      BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder''' )
    UpperCAmelCase__	:Optional[int]      =      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():
       UpperCAmelCase__	:int      =      model(__lowerCamelCase )[0]
    UpperCAmelCase__	:List[str]      =      torch.Size([1, 8, 1_0_2_4] )
    self.assertEqual(output.shape			,       __lowerCamelCase )
    UpperCAmelCase__	:Tuple      =      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]			,       __lowerCamelCase			,       atol=1e-4 ) )




@require_torch
class 				UpperCAmelCase			(			unittest.TestCase						):


 @slow
 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str] ):
    UpperCAmelCase__	:Dict      =      BertGenerationDecoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder''' )
    UpperCAmelCase__	:Any      =      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():
       UpperCAmelCase__	:Optional[int]      =      model(__lowerCamelCase )[0]
    UpperCAmelCase__	:Union[str, Any]      =      torch.Size([1, 8, 5_0_3_5_8] )
    self.assertEqual(output.shape			,       __lowerCamelCase )
    UpperCAmelCase__	:Tuple      =      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]			,       __lowerCamelCase			,       atol=1e-4 ) )







'''simple docstring'''





from .glue import GlueDataset, GlueDataTrainingArguments
from .language_modeling import (
    LineByLineTextDataset,
    LineByLineWithRefDataset,
    LineByLineWithSOPTextDataset,
    TextDataset,
    TextDatasetForNextSentencePrediction,
)
from .squad import SquadDataset, SquadDataTrainingArguments



