
"""simple docstring"""




import unittest

from transformers import AlbertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      MODEL_FOR_PRETRAINING_MAPPING,
      AlbertForMaskedLM,
      AlbertForMultipleChoice,
      AlbertForPreTraining,
      AlbertForQuestionAnswering,
      AlbertForSequenceClassification,
      AlbertForTokenClassification,
      AlbertModel,
  )
  from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST







class   _lowercase						:


     '''simple docstring'''





     def __init__(      self ,			__UpperCamelCase ,			__UpperCamelCase=13 ,			__UpperCamelCase=7 ,			__UpperCamelCase=True ,			__UpperCamelCase=True ,			__UpperCamelCase=True ,			__UpperCamelCase=True ,			__UpperCamelCase=99 ,			__UpperCamelCase=16 ,			__UpperCamelCase=36 ,			__UpperCamelCase=6 ,			__UpperCamelCase=6 ,			__UpperCamelCase=6 ,			__UpperCamelCase=37 ,			__UpperCamelCase="gelu" ,			__UpperCamelCase=0.1 ,			__UpperCamelCase=0.1 ,			__UpperCamelCase=5_12 ,			__UpperCamelCase=16 ,			__UpperCamelCase=2 ,			__UpperCamelCase=0.02 ,			__UpperCamelCase=3 ,			__UpperCamelCase=4 ,			__UpperCamelCase=None ,			)->					Dict:
          UpperCAmelCase__							: int 							=						parent
          UpperCAmelCase__							: int 							=						batch_size
          UpperCAmelCase__							: Optional[Any] 							=						seq_length
          UpperCAmelCase__							: List[str] 							=						is_training
          UpperCAmelCase__							: Optional[Any] 							=						use_input_mask
          UpperCAmelCase__							: str 							=						use_token_type_ids
          UpperCAmelCase__							: Union[str, Any] 							=						use_labels
          UpperCAmelCase__							: List[Any] 							=						vocab_size
          UpperCAmelCase__							: Optional[Any] 							=						embedding_size
          UpperCAmelCase__							: int 							=						hidden_size
          UpperCAmelCase__							: Optional[Any] 							=						num_hidden_layers
          UpperCAmelCase__							: Any 							=						num_hidden_groups
          UpperCAmelCase__							: List[Any] 							=						num_attention_heads
          UpperCAmelCase__							: int 							=						intermediate_size
          UpperCAmelCase__							: Optional[int] 							=						hidden_act
          UpperCAmelCase__							: List[Any] 							=						hidden_dropout_prob
          UpperCAmelCase__							: List[Any] 							=						attention_probs_dropout_prob
          UpperCAmelCase__							: Union[str, Any] 							=						max_position_embeddings
          UpperCAmelCase__							: List[Any] 							=						type_vocab_size
          UpperCAmelCase__							: Union[str, Any] 							=						type_sequence_label_size
          UpperCAmelCase__							: List[str] 							=						initializer_range
          UpperCAmelCase__							: Optional[Any] 							=						num_labels
          UpperCAmelCase__							: Union[str, Any] 							=						num_choices
          UpperCAmelCase__							: Optional[int] 							=						scope



     def  lowerCAmelCase__       (      self  )->					str:
          UpperCAmelCase__							: int 							=						ids_tensor([self.batch_size, self.seq_length] ,			self.vocab_size  )

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

          UpperCAmelCase__							: Union[str, Any] 							=						None
          if self.use_token_type_ids:
               UpperCAmelCase__							: int 							=						ids_tensor([self.batch_size, self.seq_length] ,			self.type_vocab_size  )

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

          UpperCAmelCase__							: Optional[Any] 							=						self.get_config()

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



     def  lowerCAmelCase__       (      self  )->					int:
          return AlbertConfig(
              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 ,			num_hidden_groups=self.num_hidden_groups ,			)



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase  )->					Dict:
          UpperCAmelCase__							: Optional[int] 							=						AlbertModel(config=__A  )
          model.to(__A  )
          model.eval()
          UpperCAmelCase__							: Optional[int] 							=						model(__A ,			attention_mask=__A ,			token_type_ids=__A  )
          UpperCAmelCase__							: Tuple 							=						model(__A ,			token_type_ids=__A  )
          UpperCAmelCase__							: Dict 							=						model(__A  )
          self.parent.assertEqual(result.last_hidden_state.shape ,			(self.batch_size, self.seq_length, self.hidden_size)  )
          self.parent.assertEqual(result.pooler_output.shape ,			(self.batch_size, self.hidden_size)  )



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase  )->					str:
          UpperCAmelCase__							: Tuple 							=						AlbertForPreTraining(config=__A  )
          model.to(__A  )
          model.eval()
          UpperCAmelCase__							: Optional[Any] 							=						model(
              __A ,			attention_mask=__A ,			token_type_ids=__A ,			labels=__A ,			sentence_order_label=__A ,			)
          self.parent.assertEqual(result.prediction_logits.shape ,			(self.batch_size, self.seq_length, self.vocab_size)  )
          self.parent.assertEqual(result.sop_logits.shape ,			(self.batch_size, config.num_labels)  )



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase  )->					Union[str, Any]:
          UpperCAmelCase__							: Tuple 							=						AlbertForMaskedLM(config=__A  )
          model.to(__A  )
          model.eval()
          UpperCAmelCase__							: str 							=						model(__A ,			attention_mask=__A ,			token_type_ids=__A ,			labels=__A  )
          self.parent.assertEqual(result.logits.shape ,			(self.batch_size, self.seq_length, self.vocab_size)  )



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase  )->					Dict:
          UpperCAmelCase__							: Tuple 							=						AlbertForQuestionAnswering(config=__A  )
          model.to(__A  )
          model.eval()
          UpperCAmelCase__							: Optional[Any] 							=						model(
              __A ,			attention_mask=__A ,			token_type_ids=__A ,			start_positions=__A ,			end_positions=__A ,			)
          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  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase  )->					Union[str, Any]:
          UpperCAmelCase__							: Tuple 							=						self.num_labels
          UpperCAmelCase__							: Any 							=						AlbertForSequenceClassification(__A  )
          model.to(__A  )
          model.eval()
          UpperCAmelCase__							: int 							=						model(__A ,			attention_mask=__A ,			token_type_ids=__A ,			labels=__A  )
          self.parent.assertEqual(result.logits.shape ,			(self.batch_size, self.num_labels)  )



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase  )->					Union[str, Any]:
          UpperCAmelCase__							: List[Any] 							=						self.num_labels
          UpperCAmelCase__							: List[Any] 							=						AlbertForTokenClassification(config=__A  )
          model.to(__A  )
          model.eval()
          UpperCAmelCase__							: Dict 							=						model(__A ,			attention_mask=__A ,			token_type_ids=__A ,			labels=__A  )
          self.parent.assertEqual(result.logits.shape ,			(self.batch_size, self.seq_length, self.num_labels)  )



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase  )->					Tuple:
          UpperCAmelCase__							: Tuple 							=						self.num_choices
          UpperCAmelCase__							: Dict 							=						AlbertForMultipleChoice(config=__A  )
          model.to(__A  )
          model.eval()
          UpperCAmelCase__							: int 							=						input_ids.unsqueeze(1  ).expand(-1 ,			self.num_choices ,			-1  ).contiguous()
          UpperCAmelCase__							: List[Any] 							=						token_type_ids.unsqueeze(1  ).expand(-1 ,			self.num_choices ,			-1  ).contiguous()
          UpperCAmelCase__							: Union[str, Any] 							=						input_mask.unsqueeze(1  ).expand(-1 ,			self.num_choices ,			-1  ).contiguous()
          UpperCAmelCase__							: Union[str, Any] 							=						model(
              __A ,			attention_mask=__A ,			token_type_ids=__A ,			labels=__A ,			)
          self.parent.assertEqual(result.logits.shape ,			(self.batch_size, self.num_choices)  )



     def  lowerCAmelCase__       (      self  )->					str:
          UpperCAmelCase__							: Tuple 							=						self.prepare_config_and_inputs()
          (
              (
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	
          )							: Union[str, 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_torch
class   _lowercase						(       lowerCAmelCase_	,   lowerCAmelCase_	,   unittest.TestCase				):


     '''simple docstring'''





     _A								=							(
         (
             AlbertModel,
             AlbertForPreTraining,
             AlbertForMaskedLM,
             AlbertForMultipleChoice,
             AlbertForSequenceClassification,
             AlbertForTokenClassification,
             AlbertForQuestionAnswering,
         )
         if is_torch_available()
         else ()
     )
     _A								=							(
         {
             'feature-extraction': AlbertModel,
             'fill-mask': AlbertForMaskedLM,
             'question-answering': AlbertForQuestionAnswering,
             'text-classification': AlbertForSequenceClassification,
             'token-classification': AlbertForTokenClassification,
             'zero-shot': AlbertForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     _A								=							True



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase=False  )->					int:
          UpperCAmelCase__							: Tuple 							=						super()._prepare_for_class(__A ,			__A ,			return_labels=__A  )

          if return_labels:
               if model_class in get_values(__A  ):
                    UpperCAmelCase__							: Dict 							=						torch.zeros(
                        (self.model_tester.batch_size, self.model_tester.seq_length) ,			dtype=torch.long ,			device=__A  )
                    UpperCAmelCase__							: Dict 							=						torch.zeros(
                        self.model_tester.batch_size ,			dtype=torch.long ,			device=__A  )
          return inputs_dict



     def  lowerCAmelCase__       (      self  )->					List[Any]:
          UpperCAmelCase__							: Tuple 							=						AlbertModelTester(self  )
          UpperCAmelCase__							: Optional[int] 							=						ConfigTester(self ,			config_class=__A ,			hidden_size=37  )



     def  lowerCAmelCase__       (      self  )->					int:
          self.config_tester.run_common_tests()



     def  lowerCAmelCase__       (      self  )->					Any:
          UpperCAmelCase__							: Optional[int] 							=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*__A  )



     def  lowerCAmelCase__       (      self  )->					Any:
          UpperCAmelCase__							: Tuple 							=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_pretraining(*__A  )



     def  lowerCAmelCase__       (      self  )->					Tuple:
          UpperCAmelCase__							: Dict 							=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*__A  )



     def  lowerCAmelCase__       (      self  )->					Any:
          UpperCAmelCase__							: Union[str, Any] 							=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_multiple_choice(*__A  )



     def  lowerCAmelCase__       (      self  )->					Dict:
          UpperCAmelCase__							: Union[str, Any] 							=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_question_answering(*__A  )



     def  lowerCAmelCase__       (      self  )->					Tuple:
          UpperCAmelCase__							: Optional[int] 							=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_sequence_classification(*__A  )



     def  lowerCAmelCase__       (      self  )->					str:
          UpperCAmelCase__							: List[str] 							=						self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
               UpperCAmelCase__							: Optional[int] 							=						type
               self.model_tester.create_and_check_model(*__A  )



     @slow
     def  lowerCAmelCase__       (      self  )->					Union[str, Any]:
          for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               UpperCAmelCase__							: str 							=						AlbertModel.from_pretrained(__A  )
               self.assertIsNotNone(__A  )







@require_torch
class   _lowercase						(       unittest.TestCase				):


     '''simple docstring'''





     @slow
     def  lowerCAmelCase__       (      self  )->					Optional[Any]:
          UpperCAmelCase__							: int 							=						AlbertModel.from_pretrained("albert-base-v2"  )
          UpperCAmelCase__							: str 							=						torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]  )
          UpperCAmelCase__							: str 							=						torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]  )
          with torch.no_grad():
               UpperCAmelCase__							: Tuple 							=						model(__A ,			attention_mask=__A  )[0]
          UpperCAmelCase__							: int 							=						torch.Size((1, 11, 7_68)  )
          self.assertEqual(output.shape ,			__A  )
          UpperCAmelCase__							: Optional[Any] 							=						torch.tensor(
              [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]  )

          self.assertTrue(torch.allclose(output[:, 1:4, 1:4] ,			__A ,			atol=1E-4  )  )





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


A__      :	Tuple          =    logging.get_logger(__name__)



def 							a__ (				lowerCAmelCase	:       nn.ModuleList       ,    lowerCAmelCase	:       nn.ModuleList       ,    lowerCAmelCase	:       List[int]		):






     '''simple docstring'''
     UpperCAmelCase__							: Optional[Any] 							=						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()		)


A__      :	List[Any]          =    {
    # 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))},
}
A__      :	List[Any]          =    {
    # 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	:       Dict       ,    lowerCAmelCase	:       Dict		):






     '''simple docstring'''
     try:
          UpperCAmelCase__							: Tuple 							=						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	:       int       ,    lowerCAmelCase	:       Tuple		):






     '''simple docstring'''
     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	:       List[str]=None       ,    lowerCAmelCase	:       List[str]=None       ,    **lowerCAmelCase	:       List[str]       ,    ):






     '''simple docstring'''
     UpperCAmelCase__							: List[str] 							=						"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
          UpperCAmelCase__							: List[str] 							=						AutoModelForSeqaSeqLM.from_pretrained(lowerCAmelCase		).eval()
     else:
          assert isinstance(lowerCAmelCase       ,    lowerCAmelCase		), F"teacher must be a model or string got type {type(lowerCAmelCase		)}"
     UpperCAmelCase__							: int 							=						teacher.config.to_diff_dict()

     try:
          UpperCAmelCase__							,	UpperCAmelCase__							: Optional[Any] 							=						teacher.config.encoder_layers, teacher.config.decoder_layers
          if e is None:
               UpperCAmelCase__							: Tuple 							=						teacher_e
          if d is None:
               UpperCAmelCase__							: str 							=						teacher_d
          init_kwargs.update({"encoder_layers": e, "decoder_layers": d}		)
     except AttributeError:  # T5
          if hasattr(teacher.config       ,    "num_encoder_layers"		):
               UpperCAmelCase__							,	UpperCAmelCase__							: Union[str, Any] 							=						teacher.config.num_encoder_layers, teacher.config.num_decoder_layers
          else:
               UpperCAmelCase__							,	UpperCAmelCase__							: Union[str, Any] 							=						teacher.config.num_layers, teacher.config.num_decoder_layers
          if e is None:
               UpperCAmelCase__							: Optional[Any] 							=						teacher_e
          if d is None:
               UpperCAmelCase__							: Optional[Any] 							=						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
     UpperCAmelCase__							: Tuple 							=						teacher.config_class(**lowerCAmelCase		)
     UpperCAmelCase__							: List[str] 							=						AutoModelForSeqaSeqLM.from_config(lowerCAmelCase		)
     # Start by copying the full teacher state dict this will copy the first N teacher layers to the student.
     UpperCAmelCase__							: Optional[int] 							=						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
          UpperCAmelCase__							,	UpperCAmelCase__							: int 							=						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:
          UpperCAmelCase__							: List[int] 							=						pick_layers_to_copy(lowerCAmelCase       ,    lowerCAmelCase		)
     if d_layers_to_copy is None:
          UpperCAmelCase__							: List[int] 							=						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}"		)
     UpperCAmelCase__							: int 							=						{
         "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)



