


import argparse
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.local_sgd import LocalSGD


########################################################################
# This is a fully working simple example to use Accelerate
# with LocalSGD, which is a method to synchronize model
# parameters every K batches. It is different, but complementary
# to gradient accumulation.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


UpperCAmelCase__   	=							16
UpperCAmelCase__   	=							32
def 				UpperCAmelCase_					(							__snake_case     ,						__snake_case = 16						)     ->			str:







 """simple docstring"""


 _lowercase							      =AutoTokenizer.from_pretrained('''bert-base-cased'''						)
 _lowercase							      =load_dataset('''glue'''     ,						'''mrpc'''						)

 def tokenize_function(__snake_case						):
  # max_length=None => use the model max length (it's actually the default)
  _lowercase							      =tokenizer(examples['''sentence1''']     ,						examples['''sentence2''']     ,						truncation=__snake_case     ,						max_length=__snake_case						)
  return outputs

 # Apply the method we just defined to all the examples in all the splits of the dataset
 # starting with the main process first:
 with accelerator.main_process_first():
  _lowercase							      =datasets.map(
      __snake_case     ,						batched=__snake_case     ,						remove_columns=['''idx''', '''sentence1''', '''sentence2''']     ,						)

 # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
 # transformers library
 _lowercase							      =tokenized_datasets.rename_column('''label'''     ,						'''labels'''						)

 def collate_fn(__snake_case						):
  # On TPU it's best to pad everything to the same length or training will be very slow.
  _lowercase							      =128 if accelerator.distributed_type == DistributedType.TPU else None
  # When using mixed precision we want round multiples of 8/16
  if accelerator.mixed_precision == "fp8":
   _lowercase							      =16
  elif accelerator.mixed_precision != "no":
   _lowercase							      =8
  else:
   _lowercase							      =None

  return tokenizer.pad(
      __snake_case     ,						padding='''longest'''     ,						max_length=__snake_case     ,						pad_to_multiple_of=__snake_case     ,						return_tensors='''pt'''     ,						)

 # Instantiate dataloaders.
 _lowercase							      =DataLoader(
     tokenized_datasets['''train''']     ,						shuffle=__snake_case     ,						collate_fn=__snake_case     ,						batch_size=__snake_case						)
 _lowercase							      =DataLoader(
     tokenized_datasets['''validation''']     ,						shuffle=__snake_case     ,						collate_fn=__snake_case     ,						batch_size=__snake_case						)

 return train_dataloader, eval_dataloader


# For testing only
if os.environ.get('''TESTING_MOCKED_DATALOADERS''', None) == "1":
  from accelerate.test_utils.training import mocked_dataloaders

  UpperCAmelCase__   	=							mocked_dataloaders  # noqa: F811
def 				UpperCAmelCase_					(							__snake_case     ,						__snake_case						)     ->			List[Any]:







 """simple docstring"""


 if os.environ.get('''TESTING_MOCKED_DATALOADERS'''     ,						__snake_case						) == "1":
  _lowercase							      =2
 # New Code #
 _lowercase							      =int(args.gradient_accumulation_steps						)
 _lowercase							      =int(args.local_sgd_steps						)
 # Initialize accelerator
 _lowercase							      =Accelerator(
     cpu=args.cpu     ,						mixed_precision=args.mixed_precision     ,						gradient_accumulation_steps=__snake_case						)
 if accelerator.distributed_type not in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
  raise NotImplementedError('''LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)'''						)
 # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
 _lowercase							      =config['''lr''']
 _lowercase							      =int(config['''num_epochs''']						)
 _lowercase							      =int(config['''seed''']						)
 _lowercase							      =int(config['''batch_size''']						)

 _lowercase							      =evaluate.load('''glue'''     ,						'''mrpc'''						)

 set_seed(__snake_case						)
 _lowercase      ,    _lowercase							      =get_dataloaders(__snake_case     ,						__snake_case						)
 # Instantiate the model (we build the model here so that the seed also control new weights initialization)
 _lowercase							      =AutoModelForSequenceClassification.from_pretrained('''bert-base-cased'''     ,						return_dict=__snake_case						)

 # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
 # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
 # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
 _lowercase							      =model.to(accelerator.device						)

 # Instantiate optimizer
 _lowercase							      =AdamW(params=model.parameters()     ,						lr=__snake_case						)

 # Instantiate scheduler
 _lowercase							      =get_linear_schedule_with_warmup(
     optimizer=__snake_case     ,						num_warmup_steps=100     ,						num_training_steps=(len(__snake_case						) * num_epochs)     ,						)

 # Prepare everything
 # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
 # prepare method.
 _lowercase      ,    _lowercase      ,    _lowercase      ,    _lowercase      ,    _lowercase							      =accelerator.prepare(
     __snake_case     ,						__snake_case     ,						__snake_case     ,						__snake_case     ,						__snake_case						)

 # Now we train the model
 for epoch in range(__snake_case						):
  model.train()
  with LocalSGD(
      accelerator=__snake_case     ,						model=__snake_case     ,						local_sgd_steps=__snake_case     ,						enabled=local_sgd_steps is not None						) as local_sgd:
   for step, batch in enumerate(__snake_case						):
    # We could avoid this line since we set the accelerator with `device_placement=True`.
    batch.to(accelerator.device						)
    # New code #
    # We use the new `accumulate` context manager to perform gradient accumulation
    # We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
    with accelerator.accumulate(__snake_case						):
     _lowercase							      =model(**__snake_case						)
     _lowercase							      =output.loss
     accelerator.backward(__snake_case						)
     optimizer.step()
     lr_scheduler.step()
     optimizer.zero_grad()
     # LocalSGD-specific line
     local_sgd.step()

  model.eval()
  for step, batch in enumerate(__snake_case						):
   # We could avoid this line since we set the accelerator with `device_placement=True`.
   batch.to(accelerator.device						)
   with torch.no_grad():
    _lowercase							      =model(**__snake_case						)
   _lowercase							      =outputs.logits.argmax(dim=-1						)
   _lowercase      ,    _lowercase							      =accelerator.gather_for_metrics((predictions, batch['''labels'''])						)
   metric.add_batch(
       predictions=__snake_case     ,						references=__snake_case     ,						)

  _lowercase							      =metric.compute()
  # Use accelerator.print to print only on the main process.
  accelerator.print(F"epoch {epoch}:"     ,						__snake_case						)
def 				UpperCAmelCase_					(							)     ->			Optional[int]:







 """simple docstring"""


 _lowercase							      =argparse.ArgumentParser(description='''Simple example of training script.'''						)
 parser.add_argument(
     '''--mixed_precision'''     ,						type=__snake_case     ,						default=__snake_case     ,						choices=['''no''', '''fp16''', '''bf16''', '''fp8''']     ,						help='''Whether to use mixed precision. Choose'''
     '''between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'''
     '''and an Nvidia Ampere GPU.'''     ,						)
 # New Code #
 parser.add_argument(
     '''--gradient_accumulation_steps'''     ,						type=__snake_case     ,						default=1     ,						help='''The number of minibatches to be ran before gradients are accumulated.'''     ,						)
 parser.add_argument(
     '''--local_sgd_steps'''     ,						type=__snake_case     ,						default=8     ,						help='''Number of local SGD steps or None to disable local SGD'''						)
 parser.add_argument('''--cpu'''     ,						action='''store_true'''     ,						help='''If passed, will train on the CPU.'''						)
 _lowercase							      =parser.parse_args()
 _lowercase							      ={'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
 training_function(__snake_case     ,						__snake_case						)


if __name__ == "__main__":
  main()








from __future__ import annotations

import os
import tempfile
import unittest

from transformers import ConvBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers import (
      TFConvBertForMaskedLM,
      TFConvBertForMultipleChoice,
      TFConvBertForQuestionAnswering,
      TFConvBertForSequenceClassification,
      TFConvBertForTokenClassification,
      TFConvBertModel,
  )






class     lowercase__       :

   '''simple docstring'''


   def __init__(    self,		__magic_name__,		__magic_name__=13,		__magic_name__=7,		__magic_name__=True,		__magic_name__=True,		__magic_name__=True,		__magic_name__=True,		__magic_name__=99,		__magic_name__=32,		__magic_name__=2,		__magic_name__=4,		__magic_name__=37,		__magic_name__="gelu",		__magic_name__=0.1,		__magic_name__=0.1,		__magic_name__=512,		__magic_name__=16,		__magic_name__=2,		__magic_name__=0.02,		__magic_name__=3,		__magic_name__=4,		__magic_name__=None,		) ->      List[str]:






       """simple docstring"""
       UpperCamelCase__							:				List[Any]          =      parent
       UpperCamelCase__							:				int          =      13
       UpperCamelCase__							:				Optional[Any]          =      7
       UpperCamelCase__							:				Dict          =      True
       UpperCamelCase__							:				Dict          =      True
       UpperCamelCase__							:				str          =      True
       UpperCamelCase__							:				List[str]          =      True
       UpperCamelCase__							:				Tuple          =      99
       UpperCamelCase__							:				Any          =      384
       UpperCamelCase__							:				Tuple          =      2
       UpperCamelCase__							:				int          =      4
       UpperCamelCase__							:				Optional[int]          =      37
       UpperCamelCase__							:				str          =      '''gelu'''
       UpperCamelCase__							:				Any          =      0.1
       UpperCamelCase__							:				Dict          =      0.1
       UpperCamelCase__							:				Tuple          =      512
       UpperCamelCase__							:				int          =      16
       UpperCamelCase__							:				int          =      2
       UpperCamelCase__							:				Any          =      0.02
       UpperCamelCase__							:				Optional[Any]          =      3
       UpperCamelCase__							:				Any          =      4
       UpperCamelCase__							:				Any          =      128
       UpperCamelCase__							:				List[str]          =      2
       UpperCamelCase__							:				List[str]          =      9
       UpperCamelCase__							:				Optional[int]          =      1
       UpperCamelCase__							:				List[Any]          =      None




   def       UpperCamelCase__  (    self							) ->      Dict:






       """simple docstring"""
       UpperCamelCase__							:				Tuple          =      ids_tensor([self.batch_size, self.seq_length],		self.vocab_size							)

       UpperCamelCase__							:				Optional[int]          =      None
       if self.use_input_mask:
           UpperCamelCase__							:				Dict          =      random_attention_mask([self.batch_size, self.seq_length]							)

       UpperCamelCase__							:				Dict          =      None
       if self.use_token_type_ids:
           UpperCamelCase__							:				List[str]          =      ids_tensor([self.batch_size, self.seq_length],		self.type_vocab_size							)

       UpperCamelCase__							:				str          =      None
       UpperCamelCase__							:				Union[str, Any]          =      None
       UpperCamelCase__							:				Dict          =      None
       if self.use_labels:
           UpperCamelCase__							:				str          =      ids_tensor([self.batch_size],		self.type_sequence_label_size							)
           UpperCamelCase__							:				List[str]          =      ids_tensor([self.batch_size, self.seq_length],		self.num_labels							)
           UpperCamelCase__							:				Tuple          =      ids_tensor([self.batch_size],		self.num_choices							)

       UpperCamelCase__							:				Any          =      ConvBertConfig(
           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,		type_vocab_size=self.type_vocab_size,		initializer_range=self.initializer_range,		return_dict=__magic_name__,		)

       return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




   def       UpperCamelCase__  (    self,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__							) ->      Tuple:






       """simple docstring"""
       UpperCamelCase__							:				Tuple          =      TFConvBertModel(config=__magic_name__							)
       UpperCamelCase__							:				str          =      {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}

       UpperCamelCase__							:				Any          =      [input_ids, input_mask]
       UpperCamelCase__							:				int          =      model(__magic_name__							)

       UpperCamelCase__							:				int          =      model(__magic_name__							)

       self.parent.assertEqual(result.last_hidden_state.shape,		(self.batch_size, self.seq_length, self.hidden_size)							)




   def       UpperCamelCase__  (    self,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__							) ->      int:






       """simple docstring"""
       UpperCamelCase__							:				Any          =      TFConvBertForMaskedLM(config=__magic_name__							)
       UpperCamelCase__							:				List[str]          =      {
           '''input_ids''': input_ids,
           '''attention_mask''': input_mask,
           '''token_type_ids''': token_type_ids,
       }
       UpperCamelCase__							:				Tuple          =      model(__magic_name__							)
       self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.seq_length, self.vocab_size)							)




   def       UpperCamelCase__  (    self,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__							) ->      Dict:






       """simple docstring"""
       UpperCamelCase__							:				str          =      self.num_labels
       UpperCamelCase__							:				Optional[int]          =      TFConvBertForSequenceClassification(config=__magic_name__							)
       UpperCamelCase__							:				str          =      {
           '''input_ids''': input_ids,
           '''attention_mask''': input_mask,
           '''token_type_ids''': token_type_ids,
       }

       UpperCamelCase__							:				Any          =      model(__magic_name__							)
       self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.num_labels)							)




   def       UpperCamelCase__  (    self,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__							) ->      Tuple:






       """simple docstring"""
       UpperCamelCase__							:				int          =      self.num_choices
       UpperCamelCase__							:				str          =      TFConvBertForMultipleChoice(config=__magic_name__							)
       UpperCamelCase__							:				Optional[int]          =      tf.tile(tf.expand_dims(__magic_name__,		1							),		(1, self.num_choices, 1)							)
       UpperCamelCase__							:				Any          =      tf.tile(tf.expand_dims(__magic_name__,		1							),		(1, self.num_choices, 1)							)
       UpperCamelCase__							:				Tuple          =      tf.tile(tf.expand_dims(__magic_name__,		1							),		(1, self.num_choices, 1)							)
       UpperCamelCase__							:				int          =      {
           '''input_ids''': multiple_choice_inputs_ids,
           '''attention_mask''': multiple_choice_input_mask,
           '''token_type_ids''': multiple_choice_token_type_ids,
       }
       UpperCamelCase__							:				Tuple          =      model(__magic_name__							)
       self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.num_choices)							)




   def       UpperCamelCase__  (    self,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__							) ->      Optional[int]:






       """simple docstring"""
       UpperCamelCase__							:				Optional[Any]          =      self.num_labels
       UpperCamelCase__							:				Any          =      TFConvBertForTokenClassification(config=__magic_name__							)
       UpperCamelCase__							:				str          =      {
           '''input_ids''': input_ids,
           '''attention_mask''': input_mask,
           '''token_type_ids''': token_type_ids,
       }
       UpperCamelCase__							:				Dict          =      model(__magic_name__							)
       self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.seq_length, self.num_labels)							)




   def       UpperCamelCase__  (    self,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__							) ->      str:






       """simple docstring"""
       UpperCamelCase__							:				Union[str, Any]          =      TFConvBertForQuestionAnswering(config=__magic_name__							)
       UpperCamelCase__							:				Optional[Any]          =      {
           '''input_ids''': input_ids,
           '''attention_mask''': input_mask,
           '''token_type_ids''': token_type_ids,
       }

       UpperCamelCase__							:				List[str]          =      model(__magic_name__							)
       self.parent.assertEqual(result.start_logits.shape,		(self.batch_size, self.seq_length)							)
       self.parent.assertEqual(result.end_logits.shape,		(self.batch_size, self.seq_length)							)




   def       UpperCamelCase__  (    self							) ->      str:






       """simple docstring"""
       UpperCamelCase__							:				Any          =      self.prepare_config_and_inputs()
       (
           (
           UpperCamelCase__
       )   ,(
           UpperCamelCase__
       )   ,(
           UpperCamelCase__
       )   ,(
           UpperCamelCase__
       )   ,(
           UpperCamelCase__
       )   ,(
           UpperCamelCase__
       )   ,(
           UpperCamelCase__
       )   ,
       )							:				Any          =      config_and_inputs
       UpperCamelCase__							:				List[str]          =      {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
       return config, inputs_dict






@require_tf
class     lowercase__       (	__lowerCamelCase			,				__lowerCamelCase			,				unittest.TestCase				):

   '''simple docstring'''


   a						:							Union[str, Any]   		=		(
       (
           TFConvBertModel,
           TFConvBertForMaskedLM,
           TFConvBertForQuestionAnswering,
           TFConvBertForSequenceClassification,
           TFConvBertForTokenClassification,
           TFConvBertForMultipleChoice,
       )
       if is_tf_available()
       else ()
   )
   a						:							int   		=		(
       {
           "feature-extraction": TFConvBertModel,
           "fill-mask": TFConvBertForMaskedLM,
           "question-answering": TFConvBertForQuestionAnswering,
           "text-classification": TFConvBertForSequenceClassification,
           "token-classification": TFConvBertForTokenClassification,
           "zero-shot": TFConvBertForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   a						:							Optional[Any]   		=		False
   a						:							List[str]   		=		False
   a						:							List[Any]   		=		False




   def       UpperCamelCase__  (    self							) ->      Optional[Any]:






       """simple docstring"""
       UpperCamelCase__							:				List[str]          =      TFConvBertModelTester(self							)
       UpperCamelCase__							:				Optional[int]          =      ConfigTester(self,		config_class=__magic_name__,		hidden_size=37							)




   def       UpperCamelCase__  (    self							) ->      int:






       """simple docstring"""
       self.config_tester.run_common_tests()




   def       UpperCamelCase__  (    self							) ->      List[Any]:






       """simple docstring"""
       UpperCamelCase__							:				Any          =      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*__magic_name__							)




   def       UpperCamelCase__  (    self							) ->      Tuple:






       """simple docstring"""
       UpperCamelCase__							:				str          =      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*__magic_name__							)




   def       UpperCamelCase__  (    self							) ->      Dict:






       """simple docstring"""
       UpperCamelCase__							:				List[Any]          =      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_multiple_choice(*__magic_name__							)




   def       UpperCamelCase__  (    self							) ->      Optional[Any]:






       """simple docstring"""
       UpperCamelCase__							:				str          =      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*__magic_name__							)




   def       UpperCamelCase__  (    self							) ->      Tuple:






       """simple docstring"""
       UpperCamelCase__							:				Optional[int]          =      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*__magic_name__							)




   def       UpperCamelCase__  (    self							) ->      Dict:






       """simple docstring"""
       UpperCamelCase__							:				str          =      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*__magic_name__							)




   @slow
   def       UpperCamelCase__  (    self							) ->      Optional[int]:






       """simple docstring"""
       UpperCamelCase__   ,UpperCamelCase__							:				Tuple          =      self.model_tester.prepare_config_and_inputs_for_common()
       UpperCamelCase__							:				Union[str, Any]          =      True
       UpperCamelCase__							:				Optional[int]          =      True

       if hasattr(__magic_name__,		'''use_cache'''							):
           UpperCamelCase__							:				Union[str, Any]          =      True

       UpperCamelCase__							:				Tuple          =      getattr(self.model_tester,		'''encoder_seq_length''',		self.model_tester.seq_length							)
       UpperCamelCase__							:				Any          =      getattr(self.model_tester,		'''key_length''',		__magic_name__							)

       for model_class in self.all_model_classes:
           UpperCamelCase__							:				Optional[int]          =      self._prepare_for_class(__magic_name__,		__magic_name__							)
           UpperCamelCase__							:				Any          =      model_class(__magic_name__							)
           UpperCamelCase__							:				List[str]          =      len(model(__magic_name__							)							)

           with tempfile.TemporaryDirectory() as tmpdirname:
               model.save_pretrained(__magic_name__,		saved_model=__magic_name__							)
               UpperCamelCase__							:				List[Any]          =      os.path.join(__magic_name__,		'''saved_model''',		'''1'''							)
               UpperCamelCase__							:				Dict          =      tf.keras.models.load_model(__magic_name__							)
               UpperCamelCase__							:				Optional[int]          =      model(__magic_name__							)

               if self.is_encoder_decoder:
                   UpperCamelCase__							:				Union[str, Any]          =      outputs['''encoder_hidden_states''']
                   UpperCamelCase__							:				Any          =      outputs['''encoder_attentions''']
               else:
                   UpperCamelCase__							:				Optional[Any]          =      outputs['''hidden_states''']
                   UpperCamelCase__							:				Dict          =      outputs['''attentions''']

               self.assertEqual(len(__magic_name__							),		__magic_name__							)

               UpperCamelCase__							:				Optional[Any]          =      getattr(
                   self.model_tester,		'''expected_num_hidden_layers''',		self.model_tester.num_hidden_layers + 1							)

               self.assertEqual(len(__magic_name__							),		__magic_name__							)
               self.assertListEqual(
                   list(output_hidden_states[0].shape[-2:]							),		[self.model_tester.seq_length, self.model_tester.hidden_size],		)

               self.assertEqual(len(__magic_name__							),		self.model_tester.num_hidden_layers							)
               self.assertListEqual(
                   list(output_attentions[0].shape[-3:]							),		[self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length],		)




   @slow
   def       UpperCamelCase__  (    self							) ->      str:






       """simple docstring"""
       UpperCamelCase__							:				str          =      TFConvBertModel.from_pretrained('''YituTech/conv-bert-base'''							)
       self.assertIsNotNone(__magic_name__							)




   def       UpperCamelCase__  (    self							) ->      Dict:






       """simple docstring"""
       UpperCamelCase__   ,UpperCamelCase__							:				int          =      self.model_tester.prepare_config_and_inputs_for_common()
       UpperCamelCase__							:				str          =      True
       UpperCamelCase__							:				Union[str, Any]          =      getattr(self.model_tester,		'''decoder_seq_length''',		self.model_tester.seq_length							)
       UpperCamelCase__							:				str          =      getattr(self.model_tester,		'''encoder_seq_length''',		self.model_tester.seq_length							)
       UpperCamelCase__							:				List[str]          =      getattr(self.model_tester,		'''key_length''',		__magic_name__							)
       UpperCamelCase__							:				int          =      getattr(self.model_tester,		'''key_length''',		__magic_name__							)

       def check_decoder_attentions_output(__magic_name__							):
           UpperCamelCase__							:				List[str]          =      len(__magic_name__							)
           self.assertEqual(out_len % 2,		0							)
           UpperCamelCase__							:				Optional[Any]          =      outputs.decoder_attentions
           self.assertEqual(len(__magic_name__							),		self.model_tester.num_hidden_layers							)
           self.assertListEqual(
               list(decoder_attentions[0].shape[-3:]							),		[self.model_tester.num_attention_heads / 2, decoder_seq_length, decoder_key_length],		)

       def check_encoder_attentions_output(__magic_name__							):
           UpperCamelCase__							:				int          =      [
               t.numpy() for t in (outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions)
           ]
           self.assertEqual(len(__magic_name__							),		self.model_tester.num_hidden_layers							)
           self.assertListEqual(
               list(attentions[0].shape[-3:]							),		[self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length],		)

       for model_class in self.all_model_classes:
           UpperCamelCase__							:				Dict          =      True
           UpperCamelCase__							:				str          =      False
           UpperCamelCase__							:				Optional[Any]          =      model_class(__magic_name__							)
           UpperCamelCase__							:				Dict          =      model(self._prepare_for_class(__magic_name__,		__magic_name__							)							)
           UpperCamelCase__							:				int          =      len(__magic_name__							)
           self.assertEqual(config.output_hidden_states,		__magic_name__							)
           check_encoder_attentions_output(__magic_name__							)

           if self.is_encoder_decoder:
               UpperCamelCase__							:				str          =      model_class(__magic_name__							)
               UpperCamelCase__							:				List[str]          =      model(self._prepare_for_class(__magic_name__,		__magic_name__							)							)
               self.assertEqual(config.output_hidden_states,		__magic_name__							)
               check_decoder_attentions_output(__magic_name__							)

           # Check that output attentions can also be changed via the config
           del inputs_dict["output_attentions"]
           UpperCamelCase__							:				List[Any]          =      True
           UpperCamelCase__							:				Dict          =      model_class(__magic_name__							)
           UpperCamelCase__							:				Optional[Any]          =      model(self._prepare_for_class(__magic_name__,		__magic_name__							)							)
           self.assertEqual(config.output_hidden_states,		__magic_name__							)
           check_encoder_attentions_output(__magic_name__							)

           # Check attention is always last and order is fine
           UpperCamelCase__							:				List[Any]          =      True
           UpperCamelCase__							:				List[Any]          =      True
           UpperCamelCase__							:				Dict          =      model_class(__magic_name__							)
           UpperCamelCase__							:				List[Any]          =      model(self._prepare_for_class(__magic_name__,		__magic_name__							)							)

           self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1),		len(__magic_name__							)							)
           self.assertEqual(model.config.output_hidden_states,		__magic_name__							)
           check_encoder_attentions_output(__magic_name__							)






@require_tf
class     lowercase__       (	unittest.TestCase				):

   '''simple docstring'''


   @slow
   def       UpperCamelCase__  (    self							) ->      int:






       """simple docstring"""
       UpperCamelCase__							:				Union[str, Any]          =      TFConvBertModel.from_pretrained('''YituTech/conv-bert-base'''							)
       UpperCamelCase__							:				Tuple          =      tf.constant([[0, 1, 2, 3, 4, 5]]							)
       UpperCamelCase__							:				int          =      model(__magic_name__							)[0]

       UpperCamelCase__							:				Dict          =      [1, 6, 768]
       self.assertEqual(output.shape,		__magic_name__							)

       UpperCamelCase__							:				Union[str, Any]          =      tf.constant(
           [
               [
                   [-0.0347_5493, -0.468_6034, -0.3063_8832],
                   [0.2263_7248, -0.2698_8646, -0.742_3424],
                   [0.1032_4868, -0.4501_3508, -0.5828_0784],
               ]
           ]							)
       tf.debugging.assert_near(output[:, :3, :3],		__magic_name__,		atol=1E-4							)

