





from __future__ import annotations

import unittest

from transformers import RoFormerConfig, 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 (
					    TFRoFormerForCausalLM,
					    TFRoFormerForMaskedLM,
					    TFRoFormerForMultipleChoice,
					    TFRoFormerForQuestionAnswering,
					    TFRoFormerForSequenceClassification,
					    TFRoFormerForTokenClassification,
					    TFRoFormerModel,
					)
					from transformers.models.roformer.modeling_tf_roformer import (
					    TFRoFormerSelfAttention,
					    TFRoFormerSinusoidalPositionalEmbedding,
					)


class   A      :
			'''simple docstring'''




			def __init__(      self       :   Union[str, Any]							,     __lowerCAmelCase       :   Optional[Any]							,     __lowerCAmelCase       :   int=13							,     __lowerCAmelCase       :   int=7							,     __lowerCAmelCase       :   Optional[int]=True							,     __lowerCAmelCase       :   Optional[int]=True							,     __lowerCAmelCase       :   Optional[int]=True							,     __lowerCAmelCase       :   Optional[int]=True							,     __lowerCAmelCase       :   Union[str, Any]=99							,     __lowerCAmelCase       :   Any=32							,     __lowerCAmelCase       :   Optional[Any]=2							,     __lowerCAmelCase       :   int=4							,     __lowerCAmelCase       :   Union[str, Any]=37							,     __lowerCAmelCase       :   List[str]="gelu"							,     __lowerCAmelCase       :   Dict=0.1							,     __lowerCAmelCase       :   Dict=0.1							,     __lowerCAmelCase       :   Optional[Any]=5_12							,     __lowerCAmelCase       :   Optional[int]=16							,     __lowerCAmelCase       :   Optional[Any]=2							,     __lowerCAmelCase       :   Optional[Any]=0.0_2							,     __lowerCAmelCase       :   Tuple=3							,     __lowerCAmelCase       :   Dict=4							,     __lowerCAmelCase       :   List[str]=None							,     )     ->     Union[str, Any]:







					"""simple docstring"""






					A__  				=       parent
					A__  				=       13
					A__  				=       7
					A__  				=       True
					A__  				=       True
					A__  				=       True
					A__  				=       True
					A__  				=       99
					A__  				=       32
					A__  				=       2
					A__  				=       4
					A__  				=       37
					A__  				=       """gelu"""
					A__  				=       0.1
					A__  				=       0.1
					A__  				=       5_12
					A__  				=       16
					A__  				=       2
					A__  				=       0.0_2
					A__  				=       3
					A__  				=       4
					A__  				=       None






			def        a_	(      self       :   str       )     ->     List[Any]:







					"""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]       )

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

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

					A__  				=       RoFormerConfig(
					    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=__lowerCAmelCase							,     )

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






			def        a_	(      self       :   Optional[int]							,     __lowerCAmelCase       :   Tuple							,     __lowerCAmelCase       :   List[Any]							,     __lowerCAmelCase       :   List[Any]							,     __lowerCAmelCase       :   str							,     __lowerCAmelCase       :   Union[str, Any]							,     __lowerCAmelCase       :   Tuple							,     __lowerCAmelCase       :   int       )     ->     Tuple:







					"""simple docstring"""






					A__  				=       TFRoFormerModel(config=__lowerCAmelCase       )
					A__  				=       {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}

					A__  				=       [input_ids, input_mask]
					A__  				=       model(__lowerCAmelCase       )

					A__  				=       model(__lowerCAmelCase       )

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






			def        a_	(      self       :   Optional[int]							,     __lowerCAmelCase       :   Dict							,     __lowerCAmelCase       :   Optional[Any]							,     __lowerCAmelCase       :   Tuple							,     __lowerCAmelCase       :   Tuple							,     __lowerCAmelCase       :   Tuple							,     __lowerCAmelCase       :   Optional[Any]							,     __lowerCAmelCase       :   Tuple       )     ->     Optional[Any]:







					"""simple docstring"""






					A__  				=       True
					A__  				=       TFRoFormerForCausalLM(config=__lowerCAmelCase       )
					A__  				=       {
					    """input_ids""": input_ids,
					    """attention_mask""": input_mask,
					    """token_type_ids""": token_type_ids,
					}
					A__  				=       model(__lowerCAmelCase       )["""logits"""]
					self.parent.assertListEqual(
					    list(prediction_scores.numpy().shape       )							,     [self.batch_size, self.seq_length, self.vocab_size]       )






			def        a_	(      self       :   int							,     __lowerCAmelCase       :   Optional[int]							,     __lowerCAmelCase       :   Tuple							,     __lowerCAmelCase       :   List[Any]							,     __lowerCAmelCase       :   Optional[Any]							,     __lowerCAmelCase       :   Union[str, Any]							,     __lowerCAmelCase       :   List[str]							,     __lowerCAmelCase       :   int       )     ->     List[Any]:







					"""simple docstring"""






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






			def        a_	(      self       :   Optional[int]							,     __lowerCAmelCase       :   str							,     __lowerCAmelCase       :   List[Any]							,     __lowerCAmelCase       :   int							,     __lowerCAmelCase       :   List[str]							,     __lowerCAmelCase       :   Union[str, Any]							,     __lowerCAmelCase       :   str							,     __lowerCAmelCase       :   List[str]       )     ->     Optional[int]:







					"""simple docstring"""






					A__  				=       self.num_labels
					A__  				=       TFRoFormerForSequenceClassification(config=__lowerCAmelCase       )
					A__  				=       {
					    """input_ids""": input_ids,
					    """attention_mask""": input_mask,
					    """token_type_ids""": token_type_ids,
					}

					A__  				=       model(__lowerCAmelCase       )
					self.parent.assertEqual(result.logits.shape							,     (self.batch_size, self.num_labels)       )






			def        a_	(      self       :   List[Any]							,     __lowerCAmelCase       :   Any							,     __lowerCAmelCase       :   Optional[Any]							,     __lowerCAmelCase       :   Optional[int]							,     __lowerCAmelCase       :   Union[str, Any]							,     __lowerCAmelCase       :   List[str]							,     __lowerCAmelCase       :   str							,     __lowerCAmelCase       :   Optional[Any]       )     ->     List[str]:







					"""simple docstring"""






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






			def        a_	(      self       :   Tuple							,     __lowerCAmelCase       :   str							,     __lowerCAmelCase       :   List[Any]							,     __lowerCAmelCase       :   Dict							,     __lowerCAmelCase       :   str							,     __lowerCAmelCase       :   List[str]							,     __lowerCAmelCase       :   Tuple							,     __lowerCAmelCase       :   str       )     ->     Tuple:







					"""simple docstring"""






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






			def        a_	(      self       :   Tuple							,     __lowerCAmelCase       :   Optional[int]							,     __lowerCAmelCase       :   List[Any]							,     __lowerCAmelCase       :   str							,     __lowerCAmelCase       :   Tuple							,     __lowerCAmelCase       :   List[str]							,     __lowerCAmelCase       :   Optional[Any]							,     __lowerCAmelCase       :   Union[str, Any]       )     ->     List[Any]:







					"""simple docstring"""






					A__  				=       TFRoFormerForQuestionAnswering(config=__lowerCAmelCase       )
					A__  				=       {
					    """input_ids""": input_ids,
					    """attention_mask""": input_mask,
					    """token_type_ids""": token_type_ids,
					}

					A__  				=       model(__lowerCAmelCase       )
					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        a_	(      self       :   Optional[int]       )     ->     Union[str, Any]:







					"""simple docstring"""






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


@require_tf
class   A      (SCREAMING_SNAKE_CASE       ,					SCREAMING_SNAKE_CASE       ,					unittest.TestCase   ):
			'''simple docstring'''




			__lowerCamelCase     :	List[Any]          =   (
			    (
			        TFRoFormerModel,
			        TFRoFormerForCausalLM,
			        TFRoFormerForMaskedLM,
			        TFRoFormerForQuestionAnswering,
			        TFRoFormerForSequenceClassification,
			        TFRoFormerForTokenClassification,
			        TFRoFormerForMultipleChoice,
			    )
			    if is_tf_available()
			    else ()
			)
			__lowerCamelCase     :	Optional[Any]          =   (
			    {
			        '''feature-extraction''': TFRoFormerModel,
			        '''fill-mask''': TFRoFormerForMaskedLM,
			        '''question-answering''': TFRoFormerForQuestionAnswering,
			        '''text-classification''': TFRoFormerForSequenceClassification,
			        '''text-generation''': TFRoFormerForCausalLM,
			        '''token-classification''': TFRoFormerForTokenClassification,
			        '''zero-shot''': TFRoFormerForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)

			__lowerCamelCase     :	Optional[int]          =   False
			__lowerCamelCase     :	List[str]          =   False






			def        a_	(      self       :   str							,     __lowerCAmelCase       :   Optional[int]							,     __lowerCAmelCase       :   Optional[int]							,     __lowerCAmelCase       :   str							,     __lowerCAmelCase       :   Optional[Any]							,     __lowerCAmelCase       :   Dict       )     ->     int:







					"""simple docstring"""






					if pipeline_test_casse_name == "TextGenerationPipelineTests":
							return True

					return False






			def        a_	(      self       :   Dict       )     ->     str:







					"""simple docstring"""






					A__  				=       TFRoFormerModelTester(self       )
					A__  				=       ConfigTester(self							,     config_class=__lowerCAmelCase							,     hidden_size=37       )






			def        a_	(      self       :   str       )     ->     int:







					"""simple docstring"""






					self.config_tester.run_common_tests()






			def        a_	(      self       :   Tuple       )     ->     Dict:







					"""simple docstring"""






					A__  				=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*__lowerCAmelCase       )






			def        a_	(      self       :   Any       )     ->     Dict:







					"""simple docstring"""






					A__  				=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_lm(*__lowerCAmelCase       )






			def        a_	(      self       :   Optional[int]       )     ->     Union[str, Any]:







					"""simple docstring"""






					A__  				=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_lm_head(*__lowerCAmelCase       )






			def        a_	(      self       :   str       )     ->     str:







					"""simple docstring"""






					A__  				=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_multiple_choice(*__lowerCAmelCase       )






			def        a_	(      self       :   Tuple       )     ->     Any:







					"""simple docstring"""






					A__  				=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_question_answering(*__lowerCAmelCase       )






			def        a_	(      self       :   Dict       )     ->     Union[str, Any]:







					"""simple docstring"""






					A__  				=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_sequence_classification(*__lowerCAmelCase       )






			def        a_	(      self       :   Optional[int]       )     ->     Union[str, Any]:







					"""simple docstring"""






					A__  				=       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_token_classification(*__lowerCAmelCase       )







			@slow
			def        a_	(      self       :   Dict       )     ->     str:







					"""simple docstring"""






					A__  				=       TFRoFormerModel.from_pretrained("""junnyu/roformer_chinese_base"""       )
					self.assertIsNotNone(__lowerCAmelCase       )


@require_tf
class   A      (unittest.TestCase   ):
			'''simple docstring'''




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







					"""simple docstring"""






					A__  				=       TFRoFormerForMaskedLM.from_pretrained("""junnyu/roformer_chinese_base"""       )
					A__  				=       tf.constant([[0, 1, 2, 3, 4, 5]]       )
					A__  				=       model(__lowerCAmelCase       )[0]

					# TODO Replace vocab size
					A__  				=       5_00_00

					A__  				=       [1, 6, vocab_size]
					self.assertEqual(output.shape							,     __lowerCAmelCase       )

					print(output[:, :3, :3]       )

					# TODO Replace values below with what was printed above.
					A__  				=       tf.constant(
					    [
					        [
					            [-0.1_2_0_5_3_3_4_1, -1.0_2_6_4_9_0_1, 0.2_9_2_2_1_9_4_6],
					            [-1.5_1_3_3_7_8_3, 0.1_9_7_4_3_3, 0.1_5_1_9_0_6_0_7],
					            [-5.0_1_3_5_4_0_3, -3.9_0_0_2_5_6, -0.8_4_0_3_8_7_6_4],
					        ]
					    ]       )
					tf.debugging.assert_near(output[:, :3, :3]							,     __lowerCAmelCase							,     atol=1e-4       )


@require_tf
class   A      (unittest.TestCase   ):
			'''simple docstring'''




			__lowerCamelCase     :	Any          =   1E-4






			def        a_	(      self       :   Dict       )     ->     Optional[Any]:







					"""simple docstring"""






					A__  				=       tf.constant([[4, 10]]       )
					A__  				=       TFRoFormerSinusoidalPositionalEmbedding(num_positions=6							,     embedding_dim=6       )

					A__  				=       emba(input_ids.shape       )
					A__  				=       tf.constant(
					    [[0.0_0_0_0, 0.0_0_0_0, 0.0_0_0_0, 1.0_0_0_0, 1.0_0_0_0, 1.0_0_0_0], [0.8_4_1_5, 0.0_4_6_4, 0.0_0_2_2, 0.5_4_0_3, 0.9_9_8_9, 1.0_0_0_0]]       )

					tf.debugging.assert_near(__lowerCAmelCase							,     __lowerCAmelCase							,     atol=self.tolerance       )







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







					"""simple docstring"""






					A__  				=       tf.constant(
					    [
					        [0.0_0_0_0, 0.0_0_0_0, 0.0_0_0_0, 0.0_0_0_0, 0.0_0_0_0],
					        [0.8_4_1_5, 0.8_2_1_9, 0.8_0_2_0, 0.7_8_1_9, 0.7_6_1_7],
					        [0.9_0_9_3, 0.9_3_6_4, 0.9_5_8_1, 0.9_7_4_9, 0.9_8_7_0],
					    ]       )
					A__  				=       TFRoFormerSinusoidalPositionalEmbedding(num_positions=5_12							,     embedding_dim=5_12       )
					emba([2, 16, 5_12]       )
					A__  				=       emba.weight[:3, :5]

					tf.debugging.assert_near(__lowerCAmelCase							,     __lowerCAmelCase							,     atol=self.tolerance       )


@require_tf
class   A      (unittest.TestCase   ):
			'''simple docstring'''




			__lowerCamelCase     :	List[str]          =   1E-4






			def        a_	(      self       :   Union[str, Any]       )     ->     Dict:







					"""simple docstring"""






					A__  				=       tf.reshape(tf.range(2 * 12 * 16 * 64							,     dtype=tf.floataa       )							,     shape=(2, 12, 16, 64)       ) / 1_00

					A__  				=       -tf.reshape(tf.range(2 * 12 * 16 * 64							,     dtype=tf.floataa       )							,     shape=(2, 12, 16, 64)       ) / 1_00

					A__  				=       TFRoFormerSinusoidalPositionalEmbedding(num_positions=32							,     embedding_dim=64       )
					A__  				=       embed_positions([2, 16, 7_68]       )[None, None, :, :]

					A__				,      A__  				=       TFRoFormerSelfAttention.apply_rotary_position_embeddings(
					    __lowerCAmelCase							,     __lowerCAmelCase							,     __lowerCAmelCase       )

					A__  				=       tf.constant(
					    [
					        [0.0_0_0_0, 0.0_1_0_0, 0.0_2_0_0, 0.0_3_0_0, 0.0_4_0_0, 0.0_5_0_0, 0.0_6_0_0, 0.0_7_0_0],
					        [-0.2_0_1_2, 0.8_8_9_7, 0.0_2_6_3, 0.9_4_0_1, 0.2_0_7_4, 0.9_4_6_3, 0.3_4_8_1, 0.9_3_4_3],
					        [-1.7_0_5_7, 0.6_2_7_1, -1.2_1_4_5, 1.3_8_9_7, -0.6_3_0_3, 1.7_6_4_7, -0.1_1_7_3, 1.8_9_8_5],
					        [-2.1_7_3_1, -1.6_3_9_7, -2.7_3_5_8, 0.2_8_5_4, -2.1_8_4_0, 1.7_1_8_3, -1.3_0_1_8, 2.4_8_7_1],
					        [0.2_7_1_7, -3.6_1_7_3, -2.9_2_0_6, -2.1_9_8_8, -3.6_6_3_8, 0.3_8_5_8, -2.9_1_5_5, 2.2_9_8_0],
					        [3.9_8_5_9, -2.1_5_8_0, -0.7_9_8_4, -4.4_9_0_4, -4.1_1_8_1, -2.0_2_5_2, -4.4_7_8_2, 1.1_2_5_3],
					    ]       )
					A__  				=       tf.constant(
					    [
					        [0.0_0_0_0, -0.0_1_0_0, -0.0_2_0_0, -0.0_3_0_0, -0.0_4_0_0, -0.0_5_0_0, -0.0_6_0_0, -0.0_7_0_0],
					        [0.2_0_1_2, -0.8_8_9_7, -0.0_2_6_3, -0.9_4_0_1, -0.2_0_7_4, -0.9_4_6_3, -0.3_4_8_1, -0.9_3_4_3],
					        [1.7_0_5_7, -0.6_2_7_1, 1.2_1_4_5, -1.3_8_9_7, 0.6_3_0_3, -1.7_6_4_7, 0.1_1_7_3, -1.8_9_8_5],
					        [2.1_7_3_1, 1.6_3_9_7, 2.7_3_5_8, -0.2_8_5_4, 2.1_8_4_0, -1.7_1_8_3, 1.3_0_1_8, -2.4_8_7_1],
					        [-0.2_7_1_7, 3.6_1_7_3, 2.9_2_0_6, 2.1_9_8_8, 3.6_6_3_8, -0.3_8_5_8, 2.9_1_5_5, -2.2_9_8_0],
					        [-3.9_8_5_9, 2.1_5_8_0, 0.7_9_8_4, 4.4_9_0_4, 4.1_1_8_1, 2.0_2_5_2, 4.4_7_8_2, -1.1_2_5_3],
					    ]       )

					tf.debugging.assert_near(query_layer[0, 0, :6, :8]							,     __lowerCAmelCase							,     atol=self.tolerance       )
					tf.debugging.assert_near(key_layer[0, 0, :6, :8]							,     __lowerCAmelCase							,     atol=self.tolerance       )










import argparse

import torch

from transformers import BertForMaskedLM


if __name__ == "__main__":
					A   : List[str]			=				argparse.ArgumentParser(
					    description=(
					        '''Extraction some layers of the full BertForMaskedLM or RObertaForMaskedLM for Transfer Learned'''
					        ''' Distillation'''
					    )
					)
					parser.add_argument('''--model_type''', default='''bert''', choices=['''bert'''])
					parser.add_argument('''--model_name''', default='''bert-base-uncased''', type=str)
					parser.add_argument('''--dump_checkpoint''', default='''serialization_dir/tf_bert-base-uncased_0247911.pth''', type=str)
					parser.add_argument('''--vocab_transform''', action='''store_true''')
					A   : Tuple			=				parser.parse_args()

					if args.model_type == "bert":
										A   : Dict			=				BertForMaskedLM.from_pretrained(args.model_name)
										A   : List[str]			=				'''bert'''
					else:
										raise ValueError('''args.model_type should be "bert".''')

					A   : Optional[Any]			=				model.state_dict()
					A   : int			=				{}

					for w in ["word_embeddings", "position_embeddings"]:
										A   : str			=				state_dict[F'''{prefix}.embeddings.{w}.weight''']
					for w in ["weight", "bias"]:
										A   : Any			=				state_dict[F'''{prefix}.embeddings.LayerNorm.{w}''']

					A   : Tuple			=				0
					for teacher_idx in [0, 2, 4, 7, 9, 1_1]:
										for w in ["weight", "bias"]:
															A   : Optional[Any]			=				state_dict[
															    F'''{prefix}.encoder.layer.{teacher_idx}.attention.self.query.{w}'''
															]
															A   : Optional[Any]			=				state_dict[
															    F'''{prefix}.encoder.layer.{teacher_idx}.attention.self.key.{w}'''
															]
															A   : Optional[Any]			=				state_dict[
															    F'''{prefix}.encoder.layer.{teacher_idx}.attention.self.value.{w}'''
															]

															A   : int			=				state_dict[
															    F'''{prefix}.encoder.layer.{teacher_idx}.attention.output.dense.{w}'''
															]
															A   : List[Any]			=				state_dict[
															    F'''{prefix}.encoder.layer.{teacher_idx}.attention.output.LayerNorm.{w}'''
															]

															A   : List[str]			=				state_dict[
															    F'''{prefix}.encoder.layer.{teacher_idx}.intermediate.dense.{w}'''
															]
															A   : Union[str, Any]			=				state_dict[
															    F'''{prefix}.encoder.layer.{teacher_idx}.output.dense.{w}'''
															]
															A   : List[str]			=				state_dict[
															    F'''{prefix}.encoder.layer.{teacher_idx}.output.LayerNorm.{w}'''
															]
										std_idx += 1

					A   : int			=				state_dict['''cls.predictions.decoder.weight''']
					A   : str			=				state_dict['''cls.predictions.bias''']
					if args.vocab_transform:
										for w in ["weight", "bias"]:
															A   : List[Any]			=				state_dict[F'''cls.predictions.transform.dense.{w}''']
															A   : List[str]			=				state_dict[F'''cls.predictions.transform.LayerNorm.{w}''']

					print(F'''N layers selected for distillation: {std_idx}''')
					print(F'''Number of params transferred for distillation: {len(compressed_sd.keys())}''')

					print(F'''Save transferred checkpoint to {args.dump_checkpoint}.''')
					torch.save(compressed_sd, args.dump_checkpoint)



