
'''simple docstring'''



import warnings
from pathlib import Path
from typing import List, Tuple, Union

import fire
from torch import nn

from transformers import AutoModelForSeqaSeqLM, AutoTokenizer, PreTrainedModel
from transformers.utils import logging


_lowercase    =		logging.get_logger(__name__)


def A				(__lowerCamelCase				:nn.ModuleList			,     __lowerCamelCase				:nn.ModuleList			,     __lowerCamelCase				:List[int]		):
						_lowerCAmelCase  						=			nn.ModuleList([src_layers[i] for i in layers_to_copy]		)
						assert len(__lowerCamelCase		) == len(__lowerCamelCase		), f'{len(__lowerCamelCase		)} != {len(__lowerCamelCase		)}'
						dest_layers.load_state_dict(layers_to_copy.state_dict()		)


_lowercase    =		{
    # maps  num layers in teacher -> num_layers in student -> which teacher layers to copy.
    # 12: bart, 16: pegasus, 6: marian/Helsinki-NLP
    12: {
        1: [0],  # This says that if the teacher has 12 layers and the student has 1, copy layer 0 of the teacher
        2: [0, 6],
        3: [0, 6, 11],
        4: [0, 4, 8, 11],
        6: [0, 2, 4, 7, 9, 11],
        9: [0, 1, 2, 4, 5, 7, 9, 10, 11],
        12: list(range(12)),
    },
    16: {  # maps  num layers in student -> which teacher layers to copy
        1: [0],
        2: [0, 15],
        3: [0, 8, 15],
        4: [0, 5, 10, 15],
        6: [0, 3, 6, 9, 12, 15],
        8: [0, 2, 4, 6, 8, 10, 12, 15],
        9: [0, 1, 3, 5, 7, 9, 11, 13, 15],
        12: [0, 1, 2, 3, 4, 5, 6, 7, 9, 11, 13, 15],
        16: list(range(16)),
    },
    6: {1: [0], 2: [0, 5], 3: [0, 2, 5], 4: [0, 1, 3, 5], 6: list(range(6))},
}
_lowercase    =		{
    # maps  num layers in student -> which teacher layers to copy.
    6: {1: [5], 2: [3, 5], 3: [1, 4, 5], 4: [1, 2, 4, 5]},
    12: {1: [11], 2: [5, 11], 3: [3, 7, 11], 6: [1, 3, 5, 8, 10, 11]},
    16: {1: [15], 4: [4, 9, 12, 15], 8: [1, 3, 5, 7, 9, 11, 13, 15]},
}


def A				(__lowerCamelCase				:Union[str, Any]			,     __lowerCamelCase				:Optional[Any]		):
						try:
												_lowerCAmelCase  						=			LAYERS_TO_COPY[n_teacher][n_student]
												return val
						except KeyError:
												if n_student != n_teacher:
																		warnings.warn(
																		    f'no hardcoded layers to copy for teacher {n_teacher} -> student {n_student}, defaulting to first'
																		    f' {n_student}'		)
												return list(range(__lowerCamelCase		)		)


def A				(__lowerCamelCase				:Dict			,     __lowerCamelCase				:int		):
						if n_student > n_teacher:
												raise ValueError(f'Cannot perform intermediate supervision for student {n_student} > teacher {n_teacher}'		)
						elif n_teacher == n_student:
												return list(range(__lowerCamelCase		)		)
						elif n_student == 1:
												return [n_teacher - 1]
						else:
												return LAYERS_TO_SUPERVISE[n_teacher][n_student]


def A				(__lowerCamelCase				:Union[str, PreTrainedModel]			,     __lowerCamelCase				:Union[str, Path] = "student"			,     __lowerCamelCase				:Union[int, None] = None			,     __lowerCamelCase				:Union[int, None] = None			,     __lowerCamelCase				:List[str]=False			,     __lowerCamelCase				:Optional[int]=None			,     __lowerCamelCase				:List[str]=None			,     **__lowerCamelCase				:Any			,     ):
						_lowerCAmelCase  						=			"""encoder_layers and decoder_layers cannot be both None-- you would just have an identical teacher."""
						assert (e is not None) or (d is not None), _msg
						if isinstance(__lowerCamelCase			,     __lowerCamelCase		):
												AutoTokenizer.from_pretrained(__lowerCamelCase		).save_pretrained(__lowerCamelCase		)  # purely for convenience
												_lowerCAmelCase  						=			AutoModelForSeqaSeqLM.from_pretrained(__lowerCamelCase		).eval()
						else:
												assert isinstance(__lowerCamelCase			,     __lowerCamelCase		), f'teacher must be a model or string got type {type(__lowerCamelCase		)}'
						_lowerCAmelCase  						=			teacher.config.to_diff_dict()

						try:
												_lowerCAmelCase						,   _lowerCAmelCase  						=			teacher.config.encoder_layers, teacher.config.decoder_layers
												if e is None:
																		_lowerCAmelCase  						=			teacher_e
												if d is None:
																		_lowerCAmelCase  						=			teacher_d
												init_kwargs.update({"""encoder_layers""": e, """decoder_layers""": d}		)
						except AttributeError:  # T5
												if hasattr(teacher.config			,     """num_encoder_layers"""		):
																		_lowerCAmelCase						,   _lowerCAmelCase  						=			teacher.config.num_encoder_layers, teacher.config.num_decoder_layers
												else:
																		_lowerCAmelCase						,   _lowerCAmelCase  						=			teacher.config.num_layers, teacher.config.num_decoder_layers
												if e is None:
																		_lowerCAmelCase  						=			teacher_e
												if d is None:
																		_lowerCAmelCase  						=			teacher_d
												if hasattr(teacher.config			,     """num_encoder_layers"""		):
																		init_kwargs.update({"""num_encoder_layers""": e, """num_decoder_layers""": d}		)
												else:
																		init_kwargs.update({"""num_layers""": e, """num_decoder_layers""": d}		)

    # Kwargs to instantiate student: teacher kwargs with updated layer numbers + **extra_config_kwargs
						init_kwargs.update(__lowerCamelCase		)

						# Copy weights
						_lowerCAmelCase  						=			teacher.config_class(**__lowerCamelCase		)
						_lowerCAmelCase  						=			AutoModelForSeqaSeqLM.from_config(__lowerCamelCase		)
						# Start by copying the full teacher state dict this will copy the first N teacher layers to the student.
						_lowerCAmelCase  						=			student.load_state_dict(teacher.state_dict()			,     strict=__lowerCamelCase		)
						assert info.missing_keys == [], info.missing_keys  # every student key should have a teacher keys.

						if copy_first_teacher_layers:  # Our copying is done. We just log and save
												_lowerCAmelCase						,   _lowerCAmelCase  						=			list(range(__lowerCamelCase		)		), list(range(__lowerCamelCase		)		)
												logger.info(
												    f'Copied encoder layers {e_layers_to_copy} and decoder layers {d_layers_to_copy}. Saving them to'
												    f' {save_path}'		)
												student.save_pretrained(__lowerCamelCase		)
												return student, e_layers_to_copy, d_layers_to_copy

						# Decide which layers of the teacher to copy. Not exactly alternating -- we try to keep first and last layer.
						if e_layers_to_copy is None:
												_lowerCAmelCase  						=			pick_layers_to_copy(__lowerCamelCase			,     __lowerCamelCase		)
						if d_layers_to_copy is None:
												_lowerCAmelCase  						=			pick_layers_to_copy(__lowerCamelCase			,     __lowerCamelCase		)

						try:
												if hasattr(
												    __lowerCamelCase			,     """prophetnet"""		):  # For ProphetNet, student.model.encoder.layers is called student.prophetnet.encoder.layers
																		copy_layers(teacher.prophetnet.encoder.layers			,     student.prophetnet.encoder.layers			,     __lowerCamelCase		)
																		copy_layers(teacher.prophetnet.decoder.layers			,     student.prophetnet.decoder.layers			,     __lowerCamelCase		)
												else:
																		copy_layers(teacher.model.encoder.layers			,     student.model.encoder.layers			,     __lowerCamelCase		)
																		copy_layers(teacher.model.decoder.layers			,     student.model.decoder.layers			,     __lowerCamelCase		)
						except AttributeError:  # For t5, student.model.encoder.layers is called student.encoder.block
												copy_layers(teacher.encoder.block			,     student.encoder.block			,     __lowerCamelCase		)
												copy_layers(teacher.decoder.block			,     student.decoder.block			,     __lowerCamelCase		)
						logger.info(
						    f'Copied encoder layers {e_layers_to_copy} and decoder layers {d_layers_to_copy}. Saving them to {save_path}'		)
						_lowerCAmelCase  						=			{
						    """teacher_type""": teacher.config.model_type,
						    """copied_encoder_layers""": e_layers_to_copy,
						    """copied_decoder_layers""": d_layers_to_copy,
						}
						student.save_pretrained(__lowerCamelCase		)
						# Save information about copying for easier reproducibility

						return student, e_layers_to_copy, d_layers_to_copy


if __name__ == "__main__":
						fire.Fire(create_student_by_copying_alternating_layers)


'''simple docstring'''



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   UpperCAmelCase_					:






						'''simple docstring'''






						def __init__( self   ,       _lowercase   ,       _lowercase=13   ,       _lowercase=7   ,       _lowercase=True   ,       _lowercase=True   ,       _lowercase=True   ,       _lowercase=True   ,       _lowercase=99   ,       _lowercase=32   ,       _lowercase=2   ,       _lowercase=4   ,       _lowercase=37   ,       _lowercase="gelu"   ,       _lowercase=0.1   ,       _lowercase=0.1   ,       _lowercase=512   ,       _lowercase=16   ,       _lowercase=2   ,       _lowercase=0.02   ,       _lowercase=3   ,       _lowercase=4   ,       _lowercase=None   ,       ):

												"""simple docstring"""



												_lowerCAmelCase  						=			parent
												_lowerCAmelCase  						=			13
												_lowerCAmelCase  						=			7
												_lowerCAmelCase  						=			True
												_lowerCAmelCase  						=			True
												_lowerCAmelCase  						=			True
												_lowerCAmelCase  						=			True
												_lowerCAmelCase  						=			99
												_lowerCAmelCase  						=			384
												_lowerCAmelCase  						=			2
												_lowerCAmelCase  						=			4
												_lowerCAmelCase  						=			37
												_lowerCAmelCase  						=			"""gelu"""
												_lowerCAmelCase  						=			0.1
												_lowerCAmelCase  						=			0.1
												_lowerCAmelCase  						=			512
												_lowerCAmelCase  						=			16
												_lowerCAmelCase  						=			2
												_lowerCAmelCase  						=			0.02
												_lowerCAmelCase  						=			3
												_lowerCAmelCase  						=			4
												_lowerCAmelCase  						=			128
												_lowerCAmelCase  						=			2
												_lowerCAmelCase  						=			9
												_lowerCAmelCase  						=			1
												_lowerCAmelCase  						=			None




						def  _lowercase     ( self				):

												"""simple docstring"""



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

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

												_lowerCAmelCase  						=			None
												if self.use_token_type_ids:
																		_lowerCAmelCase  						=			ids_tensor([self.batch_size, self.seq_length]   ,       self.type_vocab_size				)

												_lowerCAmelCase  						=			None
												_lowerCAmelCase  						=			None
												_lowerCAmelCase  						=			None
												if self.use_labels:
																		_lowerCAmelCase  						=			ids_tensor([self.batch_size]   ,       self.type_sequence_label_size				)
																		_lowerCAmelCase  						=			ids_tensor([self.batch_size, self.seq_length]   ,       self.num_labels				)
																		_lowerCAmelCase  						=			ids_tensor([self.batch_size]   ,       self.num_choices				)

												_lowerCAmelCase  						=			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=_lowercase   ,       )

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




						def  _lowercase     ( self   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			TFConvBertModel(config=_lowercase				)
												_lowerCAmelCase  						=			{"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}

												_lowerCAmelCase  						=			[input_ids, input_mask]
												_lowerCAmelCase  						=			model(_lowercase				)

												_lowerCAmelCase  						=			model(_lowercase				)

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




						def  _lowercase     ( self   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			TFConvBertForMaskedLM(config=_lowercase				)
												_lowerCAmelCase  						=			{
												    """input_ids""": input_ids,
												    """attention_mask""": input_mask,
												    """token_type_ids""": token_type_ids,
												}
												_lowerCAmelCase  						=			model(_lowercase				)
												self.parent.assertEqual(result.logits.shape   ,       (self.batch_size, self.seq_length, self.vocab_size)				)




						def  _lowercase     ( self   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			self.num_labels
												_lowerCAmelCase  						=			TFConvBertForSequenceClassification(config=_lowercase				)
												_lowerCAmelCase  						=			{
												    """input_ids""": input_ids,
												    """attention_mask""": input_mask,
												    """token_type_ids""": token_type_ids,
												}

												_lowerCAmelCase  						=			model(_lowercase				)
												self.parent.assertEqual(result.logits.shape   ,       (self.batch_size, self.num_labels)				)




						def  _lowercase     ( self   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			self.num_choices
												_lowerCAmelCase  						=			TFConvBertForMultipleChoice(config=_lowercase				)
												_lowerCAmelCase  						=			tf.tile(tf.expand_dims(_lowercase   ,       1				)   ,       (1, self.num_choices, 1)				)
												_lowerCAmelCase  						=			tf.tile(tf.expand_dims(_lowercase   ,       1				)   ,       (1, self.num_choices, 1)				)
												_lowerCAmelCase  						=			tf.tile(tf.expand_dims(_lowercase   ,       1				)   ,       (1, self.num_choices, 1)				)
												_lowerCAmelCase  						=			{
												    """input_ids""": multiple_choice_inputs_ids,
												    """attention_mask""": multiple_choice_input_mask,
												    """token_type_ids""": multiple_choice_token_type_ids,
												}
												_lowerCAmelCase  						=			model(_lowercase				)
												self.parent.assertEqual(result.logits.shape   ,       (self.batch_size, self.num_choices)				)




						def  _lowercase     ( self   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			self.num_labels
												_lowerCAmelCase  						=			TFConvBertForTokenClassification(config=_lowercase				)
												_lowerCAmelCase  						=			{
												    """input_ids""": input_ids,
												    """attention_mask""": input_mask,
												    """token_type_ids""": token_type_ids,
												}
												_lowerCAmelCase  						=			model(_lowercase				)
												self.parent.assertEqual(result.logits.shape   ,       (self.batch_size, self.seq_length, self.num_labels)				)




						def  _lowercase     ( self   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			TFConvBertForQuestionAnswering(config=_lowercase				)
												_lowerCAmelCase  						=			{
												    """input_ids""": input_ids,
												    """attention_mask""": input_mask,
												    """token_type_ids""": token_type_ids,
												}

												_lowerCAmelCase  						=			model(_lowercase				)
												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  _lowercase     ( self				):

												"""simple docstring"""



												_lowerCAmelCase  						=			self.prepare_config_and_inputs()
												(
												    (
												    _lowerCAmelCase
												)						,   (
												    _lowerCAmelCase
												)						,   (
												    _lowerCAmelCase
												)						,   (
												    _lowerCAmelCase
												)						,   (
												    _lowerCAmelCase
												)						,   (
												    _lowerCAmelCase
												)						,   (
												    _lowerCAmelCase
												)						,   
												)  						=			config_and_inputs
												_lowerCAmelCase  						=			{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
												return config, inputs_dict






@require_tf
class   UpperCAmelCase_					(		_SCREAMING_SNAKE_CASE	,							_SCREAMING_SNAKE_CASE	,							unittest.TestCase			):






						'''simple docstring'''






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




						def  _lowercase     ( self				):

												"""simple docstring"""



												_lowerCAmelCase  						=			TFConvBertModelTester(self				)
												_lowerCAmelCase  						=			ConfigTester(self   ,       config_class=_lowercase   ,       hidden_size=37				)




						def  _lowercase     ( self				):

												"""simple docstring"""



												self.config_tester.run_common_tests()




						def  _lowercase     ( self				):

												"""simple docstring"""



												_lowerCAmelCase  						=			self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*_lowercase				)




						def  _lowercase     ( self				):

												"""simple docstring"""



												_lowerCAmelCase  						=			self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_masked_lm(*_lowercase				)




						def  _lowercase     ( self				):

												"""simple docstring"""



												_lowerCAmelCase  						=			self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_multiple_choice(*_lowercase				)




						def  _lowercase     ( self				):

												"""simple docstring"""



												_lowerCAmelCase  						=			self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_question_answering(*_lowercase				)




						def  _lowercase     ( self				):

												"""simple docstring"""



												_lowerCAmelCase  						=			self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_sequence_classification(*_lowercase				)




						def  _lowercase     ( self				):

												"""simple docstring"""



												_lowerCAmelCase  						=			self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_token_classification(*_lowercase				)




						@slow
						def  _lowercase     ( self				):

												"""simple docstring"""



												_lowerCAmelCase						,   _lowerCAmelCase  						=			self.model_tester.prepare_config_and_inputs_for_common()
												_lowerCAmelCase  						=			True
												_lowerCAmelCase  						=			True

												if hasattr(_lowercase   ,       """use_cache"""				):
																		_lowerCAmelCase  						=			True

												_lowerCAmelCase  						=			getattr(self.model_tester   ,       """encoder_seq_length"""   ,       self.model_tester.seq_length				)
												_lowerCAmelCase  						=			getattr(self.model_tester   ,       """key_length"""   ,       _lowercase				)

												for model_class in self.all_model_classes:
																		_lowerCAmelCase  						=			self._prepare_for_class(_lowercase   ,       _lowercase				)
																		_lowerCAmelCase  						=			model_class(_lowercase				)
																		_lowerCAmelCase  						=			len(model(_lowercase				)				)

																		with tempfile.TemporaryDirectory() as tmpdirname:
																								model.save_pretrained(_lowercase   ,       saved_model=_lowercase				)
																								_lowerCAmelCase  						=			os.path.join(_lowercase   ,       """saved_model"""   ,       """1"""				)
																								_lowerCAmelCase  						=			tf.keras.models.load_model(_lowercase				)
																								_lowerCAmelCase  						=			model(_lowercase				)

																								if self.is_encoder_decoder:
																														_lowerCAmelCase  						=			outputs["""encoder_hidden_states"""]
																														_lowerCAmelCase  						=			outputs["""encoder_attentions"""]
																								else:
																														_lowerCAmelCase  						=			outputs["""hidden_states"""]
																														_lowerCAmelCase  						=			outputs["""attentions"""]

																								self.assertEqual(len(_lowercase				)   ,       _lowercase				)

																								_lowerCAmelCase  						=			getattr(
																								    self.model_tester   ,       """expected_num_hidden_layers"""   ,       self.model_tester.num_hidden_layers + 1				)

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

																								self.assertEqual(len(_lowercase				)   ,       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  _lowercase     ( self				):

												"""simple docstring"""



												_lowerCAmelCase  						=			TFConvBertModel.from_pretrained("""YituTech/conv-bert-base"""				)
												self.assertIsNotNone(_lowercase				)




						def  _lowercase     ( self				):

												"""simple docstring"""



												_lowerCAmelCase						,   _lowerCAmelCase  						=			self.model_tester.prepare_config_and_inputs_for_common()
												_lowerCAmelCase  						=			True
												_lowerCAmelCase  						=			getattr(self.model_tester   ,       """decoder_seq_length"""   ,       self.model_tester.seq_length				)
												_lowerCAmelCase  						=			getattr(self.model_tester   ,       """encoder_seq_length"""   ,       self.model_tester.seq_length				)
												_lowerCAmelCase  						=			getattr(self.model_tester   ,       """key_length"""   ,       _lowercase				)
												_lowerCAmelCase  						=			getattr(self.model_tester   ,       """key_length"""   ,       _lowercase				)

												def check_decoder_attentions_output(_lowercase				):
																		_lowerCAmelCase  						=			len(_lowercase				)
																		self.assertEqual(out_len % 2   ,       0				)
																		_lowerCAmelCase  						=			outputs.decoder_attentions
																		self.assertEqual(len(_lowercase				)   ,       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(_lowercase				):
																		_lowerCAmelCase  						=			[
																		    t.numpy() for t in (outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions)
																		]
																		self.assertEqual(len(_lowercase				)   ,       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:
																		_lowerCAmelCase  						=			True
																		_lowerCAmelCase  						=			False
																		_lowerCAmelCase  						=			model_class(_lowercase				)
																		_lowerCAmelCase  						=			model(self._prepare_for_class(_lowercase   ,       _lowercase				)				)
																		_lowerCAmelCase  						=			len(_lowercase				)
																		self.assertEqual(config.output_hidden_states   ,       _lowercase				)
																		check_encoder_attentions_output(_lowercase				)

																		if self.is_encoder_decoder:
																								_lowerCAmelCase  						=			model_class(_lowercase				)
																								_lowerCAmelCase  						=			model(self._prepare_for_class(_lowercase   ,       _lowercase				)				)
																								self.assertEqual(config.output_hidden_states   ,       _lowercase				)
																								check_decoder_attentions_output(_lowercase				)

																		# Check that output attentions can also be changed via the config
																		del inputs_dict["output_attentions"]
																		_lowerCAmelCase  						=			True
																		_lowerCAmelCase  						=			model_class(_lowercase				)
																		_lowerCAmelCase  						=			model(self._prepare_for_class(_lowercase   ,       _lowercase				)				)
																		self.assertEqual(config.output_hidden_states   ,       _lowercase				)
																		check_encoder_attentions_output(_lowercase				)

																		# Check attention is always last and order is fine
																		_lowerCAmelCase  						=			True
																		_lowerCAmelCase  						=			True
																		_lowerCAmelCase  						=			model_class(_lowercase				)
																		_lowerCAmelCase  						=			model(self._prepare_for_class(_lowercase   ,       _lowercase				)				)

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






@require_tf
class   UpperCAmelCase_					(		unittest.TestCase			):






						'''simple docstring'''






						@slow
						def  _lowercase     ( self				):

												"""simple docstring"""



												_lowerCAmelCase  						=			TFConvBertModel.from_pretrained("""YituTech/conv-bert-base"""				)
												_lowerCAmelCase  						=			tf.constant([[0, 1, 2, 3, 4, 5]]				)
												_lowerCAmelCase  						=			model(_lowercase				)[0]

												_lowerCAmelCase  						=			[1, 6, 768]
												self.assertEqual(output.shape   ,       _lowercase				)

												_lowerCAmelCase  						=			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]   ,       _lowercase   ,       atol=1e-4				)
