






'''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   ,	UpperCamelCase__   ,	UpperCamelCase__=13   ,	UpperCamelCase__=7   ,	UpperCamelCase__=True   ,	UpperCamelCase__=True   ,	UpperCamelCase__=99   ,	UpperCamelCase__=32   ,	UpperCamelCase__=5   ,	UpperCamelCase__=4   ,	UpperCamelCase__=37   ,	UpperCamelCase__="gelu"   ,	UpperCamelCase__=0.1   ,	UpperCamelCase__=0.1   ,	UpperCamelCase__=50   ,	UpperCamelCase__=0.02   ,	UpperCamelCase__=True   ,	UpperCamelCase__=None   ,	):






         """simple docstring"""


         a_						     =			parent
         a_						     =			batch_size
         a_						     =			seq_length
         a_						     =			is_training
         a_						     =			use_input_mask
         a_						     =			vocab_size
         a_						     =			hidden_size
         a_						     =			num_hidden_layers
         a_						     =			num_attention_heads
         a_						     =			intermediate_size
         a_						     =			hidden_act
         a_						     =			hidden_dropout_prob
         a_						     =			attention_probs_dropout_prob
         a_						     =			max_position_embeddings
         a_						     =			initializer_range
         a_						     =			use_labels
         a_						     =			scope







  def    _a       (			self	):






         """simple docstring"""


         a_						     =			ids_tensor([self.batch_size, self.seq_length]   ,	self.vocab_size	)

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

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

         a_						     =			self.get_config()

         return config, input_ids, input_mask, token_labels







  def    _a       (			self	):






         """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=UpperCamelCase__   ,	initializer_range=self.initializer_range   ,	)







  def    _a       (			self	):






         """simple docstring"""


         (
             (
             a_
         )     ,	(
             a_
         )     ,	(
             a_
         )     ,	(
             a_
         )     ,	
         )						     =			self.prepare_config_and_inputs()

         a_						     =			True
         a_						     =			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]	)
         a_						     =			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    _a       (			self   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	**UpperCamelCase__   ,	):






         """simple docstring"""


         a_						     =			BertGenerationEncoder(config=UpperCamelCase__	)
         model.to(UpperCamelCase__	)
         model.eval()
         a_						     =			model(UpperCamelCase__   ,	attention_mask=UpperCamelCase__	)
         a_						     =			model(UpperCamelCase__	)
         self.parent.assertEqual(result.last_hidden_state.shape   ,	(self.batch_size, self.seq_length, self.hidden_size)	)







  def    _a       (			self   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	**UpperCamelCase__   ,	):






         """simple docstring"""


         a_						     =			True
         a_						     =			BertGenerationEncoder(config=UpperCamelCase__	)
         model.to(UpperCamelCase__	)
         model.eval()
         a_						     =			model(
             UpperCamelCase__   ,	attention_mask=UpperCamelCase__   ,	encoder_hidden_states=UpperCamelCase__   ,	encoder_attention_mask=UpperCamelCase__   ,	)
         a_						     =			model(
             UpperCamelCase__   ,	attention_mask=UpperCamelCase__   ,	encoder_hidden_states=UpperCamelCase__   ,	)
         self.parent.assertEqual(result.last_hidden_state.shape   ,	(self.batch_size, self.seq_length, self.hidden_size)	)







  def    _a       (			self   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	**UpperCamelCase__   ,	):






         """simple docstring"""


         a_						     =			True
         a_						     =			True
         a_						     =			BertGenerationDecoder(config=UpperCamelCase__	).to(UpperCamelCase__	).eval()

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

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

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

         a_						     =			model(
             UpperCamelCase__   ,	attention_mask=UpperCamelCase__   ,	encoder_hidden_states=UpperCamelCase__   ,	encoder_attention_mask=UpperCamelCase__   ,	output_hidden_states=UpperCamelCase__   ,	)['hidden_states'][0]
         a_						     =			model(
             UpperCamelCase__   ,	attention_mask=UpperCamelCase__   ,	encoder_hidden_states=UpperCamelCase__   ,	encoder_attention_mask=UpperCamelCase__   ,	past_key_values=UpperCamelCase__   ,	output_hidden_states=UpperCamelCase__   ,	)['hidden_states'][0]

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







  def    _a       (			self   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	*UpperCamelCase__   ,	):






         """simple docstring"""


         a_						     =			BertGenerationDecoder(UpperCamelCase__	)
         model.to(UpperCamelCase__	)
         model.eval()
         a_						     =			model(UpperCamelCase__   ,	attention_mask=UpperCamelCase__   ,	labels=UpperCamelCase__	)
         self.parent.assertEqual(result.logits.shape   ,	(self.batch_size, self.seq_length, self.vocab_size)	)







  def    _a       (			self	):






         """simple docstring"""


         a_     ,	a_     ,	a_     ,	a_						     =			self.prepare_config_and_inputs()
         a_						     =			{'input_ids': input_ids, 'attention_mask': input_mask}
         return config, inputs_dict





@require_torch
class   __SCREAMING_SNAKE_CASE       (__A							,       __A							,       __A							,       unittest.TestCase       ):


  """simple docstring"""


  _a  : int      			=     (BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
  _a  : Any      			=     (BertGenerationDecoder,) if is_torch_available() else ()
  _a  : Dict      			=     (
      {'''feature-extraction''': BertGenerationEncoder, '''text-generation''': BertGenerationDecoder}
      if is_torch_available()
      else {}
  )







  def    _a       (			self	):






         """simple docstring"""


         a_						     =			BertGenerationEncoderTester(self	)
         a_						     =			ConfigTester(self   ,	config_class=UpperCamelCase__   ,	hidden_size=37	)







  def    _a       (			self	):






         """simple docstring"""


         self.config_tester.run_common_tests()







  def    _a       (			self	):






         """simple docstring"""


         a_						     =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*UpperCamelCase__	)







  def    _a       (			self	):






         """simple docstring"""


         a_     ,	a_     ,	a_     ,	a_						     =			self.model_tester.prepare_config_and_inputs()
         a_						     =			'bert'
         self.model_tester.create_and_check_model(UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__	)







  def    _a       (			self	):






         """simple docstring"""


         a_						     =			self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_model_as_decoder(*UpperCamelCase__	)







  def    _a       (			self	):






         """simple docstring"""


         a_						     =			self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_decoder_model_past_large_inputs(*UpperCamelCase__	)







  def    _a       (			self	):






         """simple docstring"""


         (
             (
             a_
         )     ,	(
             a_
         )     ,	(
             a_
         )     ,	(
             a_
         )     ,	(
             a_
         )     ,	(
             a_
         )     ,	
         )						     =			self.model_tester.prepare_config_and_inputs_for_decoder()

         a_						     =			None

         self.model_tester.create_and_check_model_as_decoder(
             UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	)







  def    _a       (			self	):






         """simple docstring"""


         a_						     =			self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_for_causal_lm(*UpperCamelCase__	)







  @slow
  def    _a       (			self	):






         """simple docstring"""


         a_						     =			BertGenerationEncoder.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder'	)
         self.assertIsNotNone(UpperCamelCase__	)





@require_torch
class   __SCREAMING_SNAKE_CASE       (unittest.TestCase       ):


  """simple docstring"""


  @slow
  def    _a       (			self	):






         """simple docstring"""


         a_						     =			BertGenerationEncoder.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder'	)
         a_						     =			torch.tensor([[101, 7_592, 1_010, 2_026, 3_899, 2_003, 10_140, 102]]	)
         with torch.no_grad():
                a_						     =			model(UpperCamelCase__	)[0]
         a_						     =			torch.Size([1, 8, 1_024]	)
         self.assertEqual(output.shape   ,	UpperCamelCase__	)
         a_						     =			torch.tensor(
             [[[0.1_775, 0.0_083, -0.0_321], [1.6_002, 0.1_287, 0.3_912], [2.1_473, 0.5_791, 0.6_066]]]	)
         self.assertTrue(torch.allclose(output[:, :3, :3]   ,	UpperCamelCase__   ,	atol=1e-4	)	)





@require_torch
class   __SCREAMING_SNAKE_CASE       (unittest.TestCase       ):


  """simple docstring"""


  @slow
  def    _a       (			self	):






         """simple docstring"""


         a_						     =			BertGenerationDecoder.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder'	)
         a_						     =			torch.tensor([[101, 7_592, 1_010, 2_026, 3_899, 2_003, 10_140, 102]]	)
         with torch.no_grad():
                a_						     =			model(UpperCamelCase__	)[0]
         a_						     =			torch.Size([1, 8, 50_358]	)
         self.assertEqual(output.shape   ,	UpperCamelCase__	)
         a_						     =			torch.tensor(
             [[[-0.5_788, -2.5_994, -3.7_054], [0.0_438, 4.7_997, 1.8_795], [1.5_862, 6.6_409, 4.4_638]]]	)
         self.assertTrue(torch.allclose(output[:, :3, :3]   ,	UpperCamelCase__   ,	atol=1e-4	)	)












'''simple docstring'''







import functools
import gc
import inspect

import torch

from .imports import is_npu_available, is_xpu_available


def       __UpperCamelCase					(      *lowercase_       :  Union[str, Any]   ):
       """simple docstring"""




       if not isinstance(lowercase_   ,       lowercase_   ):
              a_						     =			list(lowercase_   )
       for i in range(len(lowercase_   )   ):
              a_						     =			None
       gc.collect()
       if is_xpu_available():
              torch.xpu.empty_cache()
       elif is_npu_available():
              torch.npu.empty_cache()
       else:
              torch.cuda.empty_cache()
       return objects


def       __UpperCamelCase					(      lowercase_       :  Exception   ):
       """simple docstring"""




       a_						     =			[
           'CUDA out of memory.',  # CUDA OOM
           'cuDNN error: CUDNN_STATUS_NOT_SUPPORTED.',  # CUDNN SNAFU
           'DefaultCPUAllocator: can\'t allocate memory',  # CPU OOM
       ]
       if isinstance(lowercase_   ,       lowercase_   ) and len(exception.args   ) == 1:
              return any(err in exception.args[0] for err in _statements   )
       return False


def       __UpperCamelCase					(      lowercase_       :  callable = None   ,       lowercase_       :  int = 128   ):
       """simple docstring"""




       if function is None:
              return functools.partial(lowercase_   ,       starting_batch_size=lowercase_   )

       a_						     =			starting_batch_size

       def decorator(*lowercase_       :  str   ,       **lowercase_       :  Dict   ):
              nonlocal batch_size
              gc.collect()
              if is_xpu_available():
                     torch.xpu.empty_cache()
              elif is_npu_available():
                     torch.npu.empty_cache()
              else:
                     torch.cuda.empty_cache()
              a_						     =			list(inspect.signature(lowercase_   ).parameters.keys()   )
              # Guard against user error
              if len(lowercase_   ) < (len(lowercase_   ) + 1):
                     a_						     =			', '.join([F'{arg}={value}' for arg, value in zip(params[1:]   ,       args[1:]   )]   )
                     raise TypeError(
                         F'Batch size was passed into `{function.__name__}` as the first argument when called.'
                         F'Remove this as the decorator already does so: `{function.__name__}({arg_str})`'   )
              while True:
                     if batch_size == 0:
                            raise RuntimeError('No executable batch size found, reached zero.'   )
                     try:
                            return function(lowercase_   ,       *lowercase_   ,       **lowercase_   )
                     except Exception as e:
                            if should_reduce_batch_size(lowercase_   ):
                                   gc.collect()
                                   if is_xpu_available():
                                          torch.xpu.empty_cache()
                                   elif is_npu_available():
                                          torch.npu.empty_cache()
                                   else:
                                          torch.cuda.empty_cache()
                                   batch_size //= 2
                            else:
                                   raise

       return decorator




