





"""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 		SCREAMING_SNAKE_CASE_				:
				"""simple docstring"""




				def __init__(  self	:Any    ,     __lowercase	:List[str]    ,     __lowercase	:Dict=13    ,     __lowercase	:str=7    ,     __lowercase	:List[str]=True    ,     __lowercase	:Optional[Any]=True    ,     __lowercase	:Any=99    ,     __lowercase	:List[str]=32    ,     __lowercase	:List[Any]=5    ,     __lowercase	:int=4    ,     __lowercase	:List[str]=37    ,     __lowercase	:List[str]="gelu"    ,     __lowercase	:Optional[Any]=0.1    ,     __lowercase	:Tuple=0.1    ,     __lowercase	:Tuple=50    ,     __lowercase	:Union[str, Any]=0.02    ,     __lowercase	:Any=True    ,     __lowercase	:Tuple=None    ,     ):
											__lowerCamelCase   :   Optional[Any]       =parent
											__lowerCamelCase   :   Optional[int]       =batch_size
											__lowerCamelCase   :   Dict       =seq_length
											__lowerCamelCase   :   int       =is_training
											__lowerCamelCase   :   Dict       =use_input_mask
											__lowerCamelCase   :   Any       =vocab_size
											__lowerCamelCase   :   str       =hidden_size
											__lowerCamelCase   :   int       =num_hidden_layers
											__lowerCamelCase   :   str       =num_attention_heads
											__lowerCamelCase   :   int       =intermediate_size
											__lowerCamelCase   :   Union[str, Any]       =hidden_act
											__lowerCamelCase   :   List[str]       =hidden_dropout_prob
											__lowerCamelCase   :   Optional[Any]       =attention_probs_dropout_prob
											__lowerCamelCase   :   Optional[Any]       =max_position_embeddings
											__lowerCamelCase   :   int       =initializer_range
											__lowerCamelCase   :   int       =use_labels
											__lowerCamelCase   :   int       =scope




				def        __lowercase	(  self	:Dict      ):
											__lowerCamelCase   :   Any       =ids_tensor([self.batch_size, self.seq_length]    ,     self.vocab_size      )

											__lowerCamelCase   :   List[Any]       =None
											if self.use_input_mask:
																		__lowerCamelCase   :   Optional[int]       =random_attention_mask([self.batch_size, self.seq_length]      )

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

											__lowerCamelCase   :   Tuple       =self.get_config()

											return config, input_ids, input_mask, token_labels




				def        __lowercase	(  self	:Optional[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        __lowercase	(  self	:Optional[int]      ):
											(
											    (
											    __lowerCamelCase
											)     ,    (
											    __lowerCamelCase
											)     ,    (
											    __lowerCamelCase
											)     ,    (
											    __lowerCamelCase
											)     ,    
											)   :   Dict       =self.prepare_config_and_inputs()

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




				def        __lowercase	(  self	:List[Any]    ,     __lowercase	:List[Any]    ,     __lowercase	:Any    ,     __lowercase	:Tuple    ,     __lowercase	:Union[str, Any]    ,     __lowercase	:Union[str, Any]    ,     __lowercase	:Tuple    ,     **__lowercase	:Any    ,     ):
											__lowerCamelCase   :   Optional[Any]       =True
											__lowerCamelCase   :   Any       =BertGenerationEncoder(config=__lowercase      )
											model.to(__lowercase      )
											model.eval()
											__lowerCamelCase   :   Any       =model(
											    __lowercase    ,     attention_mask=__lowercase    ,     encoder_hidden_states=__lowercase    ,     encoder_attention_mask=__lowercase    ,     )
											__lowerCamelCase   :   Dict       =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        __lowercase	(  self	:List[str]    ,     __lowercase	:Optional[Any]    ,     __lowercase	:Optional[int]    ,     __lowercase	:Tuple    ,     __lowercase	:Optional[Any]    ,     __lowercase	:Dict    ,     __lowercase	:Union[str, Any]    ,     **__lowercase	:Union[str, Any]    ,     ):
											__lowerCamelCase   :   Tuple       =True
											__lowerCamelCase   :   Tuple       =True
											__lowerCamelCase   :   List[Any]       =BertGenerationDecoder(config=__lowercase      ).to(__lowercase      ).eval()

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

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

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

											__lowerCamelCase   :   Tuple       =model(
											    __lowercase    ,     attention_mask=__lowercase    ,     encoder_hidden_states=__lowercase    ,     encoder_attention_mask=__lowercase    ,     output_hidden_states=__lowercase    ,     )['''hidden_states'''][0]
											__lowerCamelCase   :   Any       =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
											__lowerCamelCase   :   Any       =ids_tensor((1,)    ,     output_from_past.shape[-1]      ).item()
											__lowerCamelCase   :   int       =output_from_no_past[:, -3:, random_slice_idx].detach()
											__lowerCamelCase   :   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        __lowercase	(  self	:Dict    ,     __lowercase	:Dict    ,     __lowercase	:str    ,     __lowercase	:int    ,     __lowercase	:Optional[int]    ,     *__lowercase	:int    ,     ):
											__lowerCamelCase   :   Optional[int]       =BertGenerationDecoder(__lowercase      )
											model.to(__lowercase      )
											model.eval()
											__lowerCamelCase   :   int       =model(__lowercase    ,     attention_mask=__lowercase    ,     labels=__lowercase      )
											self.parent.assertEqual(result.logits.shape    ,     (self.batch_size, self.seq_length, self.vocab_size)      )




				def        __lowercase	(  self	:Optional[int]      ):
											__lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase   :   Union[str, Any]       =self.prepare_config_and_inputs()
											__lowerCamelCase   :   Dict       ={'''input_ids''': input_ids, '''attention_mask''': input_mask}
											return config, inputs_dict

@require_torch
class 		SCREAMING_SNAKE_CASE_				(     snake_case__							,	snake_case__							,	snake_case__							,	unittest.TestCase ):
				"""simple docstring"""
				__snake_case :							Dict									=		(BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
				__snake_case :							List[str]									=		(BertGenerationDecoder,) if is_torch_available() else ()
				__snake_case :							Optional[Any]									=		(
				    {"""feature-extraction""": BertGenerationEncoder, """text-generation""": BertGenerationDecoder}
				    if is_torch_available()
				    else {}
				)




				def        __lowercase	(  self	:List[Any]      ):
											__lowerCamelCase   :   Dict       =BertGenerationEncoderTester(self      )
											__lowerCamelCase   :   Dict       =ConfigTester(self    ,     config_class=__lowercase    ,     hidden_size=37      )




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




				def        __lowercase	(  self	:List[Any]      ):
											__lowerCamelCase   :   List[str]       =self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__lowercase      )




				def        __lowercase	(  self	:Optional[int]      ):
											__lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase     ,    __lowerCamelCase   :   List[str]       =self.model_tester.prepare_config_and_inputs()
											__lowerCamelCase   :   int       ='''bert'''
											self.model_tester.create_and_check_model(__lowercase    ,     __lowercase    ,     __lowercase    ,     __lowercase      )




				def        __lowercase	(  self	:List[str]      ):
											__lowerCamelCase   :   Tuple       =self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_model_as_decoder(*__lowercase      )




				def        __lowercase	(  self	:Union[str, Any]      ):
											__lowerCamelCase   :   Tuple       =self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_decoder_model_past_large_inputs(*__lowercase      )




				def        __lowercase	(  self	:Tuple      ):
											# This regression test was failing with PyTorch < 1.3
											(
											    (
											    __lowerCamelCase
											)     ,    (
											    __lowerCamelCase
											)     ,    (
											    __lowerCamelCase
											)     ,    (
											    __lowerCamelCase
											)     ,    (
											    __lowerCamelCase
											)     ,    (
											    __lowerCamelCase
											)     ,    
											)   :   str       =self.model_tester.prepare_config_and_inputs_for_decoder()

											__lowerCamelCase   :   List[str]       =None

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




				def        __lowercase	(  self	:Union[str, Any]      ):
											__lowerCamelCase   :   str       =self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_for_causal_lm(*__lowercase      )




				@slow
				def        __lowercase	(  self	:Dict      ):
											__lowerCamelCase   :   Any       =BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''      )
											self.assertIsNotNone(__lowercase      )

@require_torch
class 		SCREAMING_SNAKE_CASE_				(     unittest.TestCase ):
				"""simple docstring"""




				@slow
				def        __lowercase	(  self	:Optional[int]      ):
											__lowerCamelCase   :   int       =BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''      )
											__lowerCamelCase   :   Union[str, Any]       =torch.tensor([[101, 7592, 1010, 2026, 3899, 2003, 1_0140, 102]]      )
											with torch.no_grad():
																		__lowerCamelCase   :   List[str]       =model(__lowercase      )[0]
											__lowerCamelCase   :   List[str]       =torch.Size([1, 8, 1024]      )
											self.assertEqual(output.shape    ,     __lowercase      )
											__lowerCamelCase   :   Optional[Any]       =torch.tensor(
											    [[[0.1775, 0.0083, -0.0321], [1.6002, 0.1287, 0.3912], [2.1473, 0.5791, 0.6066]]]      )
											self.assertTrue(torch.allclose(output[:, :3, :3]    ,     __lowercase    ,     atol=1e-4      )      )

@require_torch
class 		SCREAMING_SNAKE_CASE_				(     unittest.TestCase ):
				"""simple docstring"""




				@slow
				def        __lowercase	(  self	:Dict      ):
											__lowerCamelCase   :   List[str]       =BertGenerationDecoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''      )
											__lowerCamelCase   :   Any       =torch.tensor([[101, 7592, 1010, 2026, 3899, 2003, 1_0140, 102]]      )
											with torch.no_grad():
																		__lowerCamelCase   :   List[str]       =model(__lowercase      )[0]
											__lowerCamelCase   :   str       =torch.Size([1, 8, 5_0358]      )
											self.assertEqual(output.shape    ,     __lowercase      )
											__lowerCamelCase   :   str       =torch.tensor(
											    [[[-0.5788, -2.5994, -3.7054], [0.0438, 4.7997, 1.8795], [1.5862, 6.6409, 4.4638]]]      )
											self.assertTrue(torch.allclose(output[:, :3, :3]    ,     __lowercase    ,     atol=1e-4      )      )











"""simple docstring"""





import os
import sys
import tempfile

import torch

from .state import AcceleratorState
from .utils import PrecisionType, PrepareForLaunch, is_mps_available, patch_environment



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Any     ,    SCREAMING_SNAKE_CASE							:				List[str]=()     ,    SCREAMING_SNAKE_CASE							:				Union[str, Any]=None     ,    SCREAMING_SNAKE_CASE							:				Dict="no"     ,    SCREAMING_SNAKE_CASE							:				int="29500" ):


							'''simple docstring'''







							__lowerCamelCase   :   Optional[int]       =False
							__lowerCamelCase   :   Tuple       =False
							if any(key.startswith('''KAGGLE''' ) for key in os.environ.keys() ):
														__lowerCamelCase   :   Optional[int]       =True
							elif "IPython" in sys.modules:
														__lowerCamelCase   :   int       ='''google.colab''' in str(sys.modules['''IPython'''].get_ipython() )

							try:
														__lowerCamelCase   :   str       =PrecisionType(mixed_precision.lower() )
							except ValueError:
														raise ValueError(
														    F'Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}.' )

							if (in_colab or in_kaggle) and (os.environ.get('''TPU_NAME'''     ,    SCREAMING_SNAKE_CASE ) is not None):
														# TPU launch
														import torch_xla.distributed.xla_multiprocessing as xmp

														if len(AcceleratorState._shared_state ) > 0:
																					raise ValueError(
																					    '''To train on TPU in Colab or Kaggle Kernel, the `Accelerator` should only be initialized inside '''
																					    '''your training function. Restart your notebook and make sure no cells initializes an '''
																					    '''`Accelerator`.''' )
														if num_processes is None:
																					__lowerCamelCase   :   int       =8

														__lowerCamelCase   :   Optional[Any]       =PrepareForLaunch(SCREAMING_SNAKE_CASE     ,    distributed_type='''TPU''' )
														print(F'Launching a training on {num_processes} TPU cores.' )
														xmp.spawn(SCREAMING_SNAKE_CASE     ,    args=SCREAMING_SNAKE_CASE     ,    nprocs=SCREAMING_SNAKE_CASE     ,    start_method='''fork''' )
							elif in_colab:
														# No need for a distributed launch otherwise as it's either CPU or one GPU.
														if torch.cuda.is_available():
																					print('''Launching training on one GPU.''' )
														else:
																					print('''Launching training on one CPU.''' )
														function(*SCREAMING_SNAKE_CASE )
							else:
														if num_processes is None:
																					raise ValueError(
																					    '''You have to specify the number of GPUs you would like to use, add `num_processes=...` to your call.''' )

														if num_processes > 1:
																					# Multi-GPU launch
																					from torch.multiprocessing import start_processes
																					from torch.multiprocessing.spawn import ProcessRaisedException

																					if len(AcceleratorState._shared_state ) > 0:
																												raise ValueError(
																												    '''To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized '''
																												    '''inside your training function. Restart your notebook and make sure no cells initializes an '''
																												    '''`Accelerator`.''' )

																					if torch.cuda.is_initialized():
																												raise ValueError(
																												    '''To launch a multi-GPU training from your notebook, you need to avoid running any instruction '''
																												    '''using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA '''
																												    '''function.''' )

																					# torch.distributed will expect a few environment variable to be here. We set the ones common to each
																					# process here (the other ones will be set be the launcher).
																					with patch_environment(
																					    world_size=SCREAMING_SNAKE_CASE     ,    master_addr='''127.0.01'''     ,    master_port=SCREAMING_SNAKE_CASE     ,    mixed_precision=SCREAMING_SNAKE_CASE ):
																												__lowerCamelCase   :   List[str]       =PrepareForLaunch(SCREAMING_SNAKE_CASE     ,    distributed_type='''MULTI_GPU''' )
																												print(F'Launching training on {num_processes} GPUs.' )
																												try:
																																			start_processes(SCREAMING_SNAKE_CASE     ,    args=SCREAMING_SNAKE_CASE     ,    nprocs=SCREAMING_SNAKE_CASE     ,    start_method='''fork''' )
																												except ProcessRaisedException as e:
																																			if "Cannot re-initialize CUDA in forked subprocess" in e.args[0]:
																																										raise RuntimeError(
																																										    '''CUDA has been initialized before the `notebook_launcher` could create a forked subprocess. '''
																																										    '''This likely stems from an outside import causing issues once the `notebook_launcher()` is called. '''
																																										    '''Please review your imports and test them when running the `notebook_launcher()` to identify '''
																																										    '''which one is problematic.''' ) from e

														else:
																					# No need for a distributed launch otherwise as it's either CPU, GPU or MPS.
																					if is_mps_available():
																												__lowerCamelCase   :   Optional[Any]       ='''1'''
																												print('''Launching training on MPS.''' )
																					elif torch.cuda.is_available():
																												print('''Launching training on one GPU.''' )
																					else:
																												print('''Launching training on CPU.''' )
																					function(*SCREAMING_SNAKE_CASE )



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				str     ,    SCREAMING_SNAKE_CASE							:				Optional[int]=()     ,    SCREAMING_SNAKE_CASE							:				Optional[Any]=2 ):


							'''simple docstring'''







							from torch.multiprocessing import start_processes

							with tempfile.NamedTemporaryFile() as tmp_file:
														# torch.distributed will expect a few environment variable to be here. We set the ones common to each
														# process here (the other ones will be set be the launcher).
														with patch_environment(
														    world_size=SCREAMING_SNAKE_CASE     ,    master_addr='''127.0.01'''     ,    master_port='''29500'''     ,    accelerate_mixed_precision='''no'''     ,    accelerate_debug_rdv_file=tmp_file.name     ,    accelerate_use_cpu='''yes'''     ,    ):
																					__lowerCamelCase   :   Tuple       =PrepareForLaunch(SCREAMING_SNAKE_CASE     ,    debug=SCREAMING_SNAKE_CASE )
																					start_processes(SCREAMING_SNAKE_CASE     ,    args=SCREAMING_SNAKE_CASE     ,    nprocs=SCREAMING_SNAKE_CASE     ,    start_method='''fork''' )




