
import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


lowercase_ 				=			logging.getLogger(__name__)

require_version("""pytorch_lightning>=1.0.4""")

lowercase_ 				=			{
    '''base''': AutoModel,
    '''sequence-classification''': AutoModelForSequenceClassification,
    '''question-answering''': AutoModelForQuestionAnswering,
    '''pretraining''': AutoModelForPreTraining,
    '''token-classification''': AutoModelForTokenClassification,
    '''language-modeling''': AutoModelWithLMHead,
    '''summarization''': AutoModelForSeqaSeqLM,
    '''translation''': AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
lowercase_ 				=			{
    '''linear''': get_linear_schedule_with_warmup,
    '''cosine''': get_cosine_schedule_with_warmup,
    '''cosine_w_restarts''': get_cosine_with_hard_restarts_schedule_with_warmup,
    '''polynomial''': get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
lowercase_ 				=			sorted(arg_to_scheduler.keys())
lowercase_ 				=			'''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''





class     __UpperCamelCase (							pl.LightningModule       ):






     """simple docstring"""







     def __init__(				self							:							Optional[Any]     ,						_A							:							int     ,						_A							:							Any=None     ,						_A							:							List[Any]="base"     ,						_A							:							Optional[Any]=None     ,						_A							:							str=None     ,						_A							:							Union[str, Any]=None     ,						**_A							:							Dict     ,						):







       """simple docstring"""




       super().__init__()
       # TODO: move to self.save_hyperparameters()
       # self.save_hyperparameters()
       # can also expand arguments into trainer signature for easier reading

       self.save_hyperparameters(_snake_case )
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		0
       __SCREAMING_SNAKE_CASE      :	int					=		Path(self.hparams.output_dir )
       __SCREAMING_SNAKE_CASE      :	str					=		self.hparams.cache_dir if self.hparams.cache_dir else None
       if config is None:
         __SCREAMING_SNAKE_CASE      :	List[Any]					=		AutoConfig.from_pretrained(
             self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path     ,						**({'''num_labels''': num_labels} if num_labels is not None else {})     ,						cache_dir=_snake_case     ,						**_snake_case     ,						)
       else:
         __SCREAMING_SNAKE_CASE      :	Tuple					=		config

       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		('''encoder_layerdrop''', '''decoder_layerdrop''', '''dropout''', '''attention_dropout''')
       for p in extra_model_params:
         if getattr(self.hparams     ,						_snake_case     ,						_snake_case ):
           assert hasattr(self.config     ,						_snake_case ), F'''model config doesn\'t have a `{p}` attribute'''
           setattr(self.config     ,						_snake_case     ,						getattr(self.hparams     ,						_snake_case ) )

       if tokenizer is None:
         __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		AutoTokenizer.from_pretrained(
             self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path     ,						cache_dir=_snake_case     ,						)
       else:
         __SCREAMING_SNAKE_CASE      :	Any					=		tokenizer
       __SCREAMING_SNAKE_CASE      :	List[str]					=		MODEL_MODES[mode]
       if model is None:
         __SCREAMING_SNAKE_CASE      :	Optional[int]					=		self.model_type.from_pretrained(
             self.hparams.model_name_or_path     ,						from_tf=bool('''.ckpt''' in self.hparams.model_name_or_path )     ,						config=self.config     ,						cache_dir=_snake_case     ,						)
       else:
         __SCREAMING_SNAKE_CASE      :	int					=		model



     def  UpperCAmelCase__							(				self							:							Optional[int]     ,						*_A							:							str     ,						**_A							:							List[Any] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	str					=		self.model_type.from_pretrained(*_snake_case     ,						**_snake_case )



     def  UpperCAmelCase__							(				self							:							Any ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	str					=		arg_to_scheduler[self.hparams.lr_scheduler]
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		get_schedule_func(
           self.opt     ,						num_warmup_steps=self.hparams.warmup_steps     ,						num_training_steps=self.total_steps() )
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		{'''scheduler''': scheduler, '''interval''': '''step''', '''frequency''': 1}
       return scheduler



     def  UpperCAmelCase__							(				self							:							Optional[Any] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	List[str]					=		self.model
       __SCREAMING_SNAKE_CASE      :	Tuple					=		['''bias''', '''LayerNorm.weight''']
       __SCREAMING_SNAKE_CASE      :	Dict					=		[
           {
               '''params''': [
                   p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay )
               ],  # check this named paramters
               '''weight_decay''': self.hparams.weight_decay,
           },
           {
               '''params''': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay )],
               '''weight_decay''': 0.0,
           },
       ]
       if self.hparams.adafactor:
         __SCREAMING_SNAKE_CASE      :	Any					=		Adafactor(
             _snake_case     ,						lr=self.hparams.learning_rate     ,						scale_parameter=_snake_case     ,						relative_step=_snake_case )

       else:
         __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		AdamW(
             _snake_case     ,						lr=self.hparams.learning_rate     ,						eps=self.hparams.adam_epsilon )
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		optimizer

       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		self.get_lr_scheduler()

       return [optimizer], [scheduler]



     def  UpperCAmelCase__							(				self							:							Optional[int]     ,						_A							:							List[Any]     ,						_A							:							Tuple ):







       """simple docstring"""




       return self.validation_step(_snake_case     ,						_snake_case )



     def  UpperCAmelCase__							(				self							:							List[Any]     ,						_A							:							List[Any] ):







       """simple docstring"""




       return self.validation_end(_snake_case )



     def  UpperCAmelCase__							(				self							:							List[str] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Any					=		max(1     ,						self.hparams.gpus )  # TODO: consider num_tpu_cores
       __SCREAMING_SNAKE_CASE      :	Dict					=		self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
       return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs



     def  UpperCAmelCase__							(				self							:							str     ,						_A							:							List[str] ):







       """simple docstring"""




       if stage == "test":
         __SCREAMING_SNAKE_CASE      :	Dict					=		len(self.test_dataloader().dataset )
       else:
         __SCREAMING_SNAKE_CASE      :	Optional[int]					=		self.get_dataloader('''train'''     ,						self.hparams.train_batch_size     ,						shuffle=_snake_case )
         __SCREAMING_SNAKE_CASE      :	Optional[int]					=		len(self.train_dataloader().dataset )



     def  UpperCAmelCase__							(				self							:							str     ,						_A							:							Any     ,						_A							:							Tuple     ,						_A							:							Optional[Any] = False ):







       """simple docstring"""




       raise NotImplementedError('''You must implement this for your task''' )



     def  UpperCAmelCase__							(				self							:							List[str] ):







       """simple docstring"""




       return self.train_loader



     def  UpperCAmelCase__							(				self							:							Optional[Any] ):







       """simple docstring"""




       return self.get_dataloader('''dev'''     ,						self.hparams.eval_batch_size     ,						shuffle=_snake_case )



     def  UpperCAmelCase__							(				self							:							Optional[int] ):







       """simple docstring"""




       return self.get_dataloader('''test'''     ,						self.hparams.eval_batch_size     ,						shuffle=_snake_case )



     def  UpperCAmelCase__							(				self							:							List[str]     ,						_A							:							Optional[Any] ):







       """simple docstring"""




       return os.path.join(
           self.hparams.data_dir     ,						'''cached_{}_{}_{}'''.format(
               _snake_case     ,						list(filter(_snake_case     ,						self.hparams.model_name_or_path.split('''/''' ) ) ).pop()     ,						str(self.hparams.max_seq_length )     ,						)     ,						)



     @pl.utilities.rank_zero_only
     def  UpperCAmelCase__							(				self							:							List[Any]     ,						_A							:							Union[str, Any] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		self.output_dir.joinpath('''best_tfmr''' )
       __SCREAMING_SNAKE_CASE      :	Any					=		self.step_count
       self.model.save_pretrained(_snake_case )
       self.tokenizer.save_pretrained(_snake_case )




     @staticmethod
     def  UpperCAmelCase__							(				_A							:							List[Any]     ,						_A							:							Optional[Any] ):







       """simple docstring"""




       parser.add_argument(
           '''--model_name_or_path'''     ,						default=_snake_case     ,						type=_snake_case     ,						required=_snake_case     ,						help='''Path to pretrained model or model identifier from huggingface.co/models'''     ,						)
       parser.add_argument(
           '''--config_name'''     ,						default=''''''     ,						type=_snake_case     ,						help='''Pretrained config name or path if not the same as model_name''' )
       parser.add_argument(
           '''--tokenizer_name'''     ,						default=_snake_case     ,						type=_snake_case     ,						help='''Pretrained tokenizer name or path if not the same as model_name'''     ,						)
       parser.add_argument(
           '''--cache_dir'''     ,						default=str(Path(_snake_case ).parent / '''test_run''' / '''cache''' )     ,						type=_snake_case     ,						help='''Where do you want to store the pre-trained models downloaded from huggingface.co'''     ,						)
       parser.add_argument(
           '''--encoder_layerdrop'''     ,						type=_snake_case     ,						help='''Encoder layer dropout probability (Optional). Goes into model.config'''     ,						)
       parser.add_argument(
           '''--decoder_layerdrop'''     ,						type=_snake_case     ,						help='''Decoder layer dropout probability (Optional). Goes into model.config'''     ,						)
       parser.add_argument(
           '''--dropout'''     ,						type=_snake_case     ,						help='''Dropout probability (Optional). Goes into model.config'''     ,						)
       parser.add_argument(
           '''--attention_dropout'''     ,						type=_snake_case     ,						help='''Attention dropout probability (Optional). Goes into model.config'''     ,						)
       parser.add_argument('''--learning_rate'''     ,						default=5e-5     ,						type=_snake_case     ,						help='''The initial learning rate for Adam.''' )
       parser.add_argument(
           '''--lr_scheduler'''     ,						default='''linear'''     ,						choices=_snake_case     ,						metavar=_snake_case     ,						type=_snake_case     ,						help='''Learning rate scheduler'''     ,						)
       parser.add_argument('''--weight_decay'''     ,						default=0.0     ,						type=_snake_case     ,						help='''Weight decay if we apply some.''' )
       parser.add_argument('''--adam_epsilon'''     ,						default=1e-8     ,						type=_snake_case     ,						help='''Epsilon for Adam optimizer.''' )
       parser.add_argument('''--warmup_steps'''     ,						default=0     ,						type=_snake_case     ,						help='''Linear warmup over warmup_steps.''' )
       parser.add_argument('''--num_workers'''     ,						default=4     ,						type=_snake_case     ,						help='''kwarg passed to DataLoader''' )
       parser.add_argument('''--num_train_epochs'''     ,						dest='''max_epochs'''     ,						default=3     ,						type=_snake_case )
       parser.add_argument('''--train_batch_size'''     ,						default=32     ,						type=_snake_case )
       parser.add_argument('''--eval_batch_size'''     ,						default=32     ,						type=_snake_case )
       parser.add_argument('''--adafactor'''     ,						action='''store_true''' )






class     __UpperCamelCase (							pl.Callback       ):






     """simple docstring"""







     def  UpperCAmelCase__							(				self							:							Dict     ,						_A							:							Optional[int]     ,						_A							:							Union[str, Any] ):







       """simple docstring"""




       if (
           trainer.is_global_zero and trainer.global_rank == 0
       ):  # we initialize the retriever only on master worker with RAY. In new pytorch-lightning accelorators are removed.
         pl_module.model.rag.retriever.init_retrieval()  # better to use hook functions.






class     __UpperCamelCase (							pl.Callback       ):






     """simple docstring"""







     def  UpperCAmelCase__							(				self							:							Any     ,						_A							:							Optional[int]     ,						_A							:							Any ):







       """simple docstring"""




       for name, param in pl_module.model.rag.named_parameters():
         if param.grad is None:
           print(_snake_case )






class     __UpperCamelCase (							pl.Callback       ):






     """simple docstring"""







     def  UpperCAmelCase__							(				self							:							int     ,						_A							:							Optional[Any]     ,						_A							:							Optional[Any] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		trainer.lr_schedulers[0]['''scheduler''']
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		{F'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr() )}
       pl_module.logger.log_metrics(_snake_case )



     def  UpperCAmelCase__							(				self							:							str     ,						_A							:							Any     ,						_A							:							Tuple ):







       """simple docstring"""




       rank_zero_info('''***** Validation results *****''' )
       __SCREAMING_SNAKE_CASE      :	Dict					=		trainer.callback_metrics
       # Log results
       for key in sorted(_snake_case ):
         if key not in ["log", "progress_bar"]:
           rank_zero_info('''{} = {}\n'''.format(_snake_case     ,						str(metrics[key] ) ) )




     def  UpperCAmelCase__							(				self							:							str     ,						_A							:							int     ,						_A							:							Any ):







       """simple docstring"""




       rank_zero_info('''***** Test results *****''' )
       __SCREAMING_SNAKE_CASE      :	List[str]					=		trainer.callback_metrics
       # Log and save results to file
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		os.path.join(pl_module.hparams.output_dir     ,						'''test_results.txt''' )
       with open(_snake_case     ,						'''w''' ) as writer:
         for key in sorted(_snake_case ):
           if key not in ["log", "progress_bar"]:
             rank_zero_info('''{} = {}\n'''.format(_snake_case     ,						str(metrics[key] ) ) )
             writer.write('''{} = {}\n'''.format(_snake_case     ,						str(metrics[key] ) ) )






def 							a__					(							snake_case       ,			snake_case							):





  """simple docstring"""
  #  To allow all pl args uncomment the following line
  #  parser = pl.Trainer.add_argparse_args(parser)
  parser.add_argument(
      '''--output_dir'''       ,			default=str(Path(_UpperCAmelCase							).parent / '''test_run''' / '''model_checkpoints'''							)       ,			type=_UpperCAmelCase       ,			help='''The output directory where the model predictions and checkpoints will be written.'''       ,			)
  parser.add_argument(
      '''--fp16'''       ,			action='''store_true'''       ,			help='''Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit'''       ,			)

  parser.add_argument(
      '''--fp16_opt_level'''       ,			type=_UpperCAmelCase       ,			default='''O2'''       ,			help=(
          '''For fp16: Apex AMP optimization level selected in [\'O0\', \'O1\', \'O2\', and \'O3\'].'''
          '''See details at https://nvidia.github.io/apex/amp.html'''
      )       ,			)
  parser.add_argument('''--n_tpu_cores'''       ,			dest='''tpu_cores'''       ,			type=_UpperCAmelCase							)
  parser.add_argument('''--max_grad_norm'''       ,			dest='''gradient_clip_val'''       ,			default=1.0       ,			type=_UpperCAmelCase       ,			help='''Max gradient norm'''							)
  parser.add_argument('''--do_train'''       ,			action='''store_true'''       ,			help='''Whether to run training.'''							)
  parser.add_argument('''--do_predict'''       ,			action='''store_true'''       ,			help='''Whether to run predictions on the test set.'''							)
  parser.add_argument(
      '''--gradient_accumulation_steps'''       ,			dest='''accumulate_grad_batches'''       ,			type=_UpperCAmelCase       ,			default=1       ,			help='''Number of updates steps to accumulate before performing a backward/update pass.'''       ,			)
  parser.add_argument('''--seed'''       ,			type=_UpperCAmelCase       ,			default=42       ,			help='''random seed for initialization'''							)
  parser.add_argument(
      '''--data_dir'''       ,			default=str(Path(_UpperCAmelCase							).parent / '''test_run''' / '''dummy-train-data'''							)       ,			type=_UpperCAmelCase       ,			help='''The input data dir. Should contain the training files for the CoNLL-2003 NER task.'''       ,			)


def 							a__					(							snake_case       ,			snake_case       ,			snake_case=None       ,			snake_case=True       ,			snake_case=[]       ,			snake_case=None       ,			snake_case=None       ,			**snake_case       ,			):





  """simple docstring"""
  pl.seed_everything(args.seed							)

  # init model
  __SCREAMING_SNAKE_CASE      :	str					=		Path(model.hparams.output_dir							)
  odir.mkdir(exist_ok=_UpperCAmelCase							)

  # add custom checkpoints
  if checkpoint_callback is None:
    __SCREAMING_SNAKE_CASE      :	Dict					=		pl.callbacks.ModelCheckpoint(
        filepath=args.output_dir       ,			prefix='''checkpoint'''       ,			monitor='''val_loss'''       ,			mode='''min'''       ,			save_top_k=1							)
  if early_stopping_callback:
    extra_callbacks.append(_UpperCAmelCase							)
  if logging_callback is None:
    __SCREAMING_SNAKE_CASE      :	List[str]					=		LoggingCallback()

  __SCREAMING_SNAKE_CASE      :	Dict					=		{}

  if args.fpaa:
    __SCREAMING_SNAKE_CASE      :	int					=		16

  if args.gpus > 1:
    __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		'''auto'''
    __SCREAMING_SNAKE_CASE      :	int					=		'''ddp'''

  __SCREAMING_SNAKE_CASE      :	Tuple					=		args.accumulate_grad_batches
  __SCREAMING_SNAKE_CASE      :	Optional[int]					=		None
  __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		'''auto'''

  __SCREAMING_SNAKE_CASE      :	Any					=		pl.Trainer.from_argparse_args(
      _UpperCAmelCase       ,			weights_summary=_UpperCAmelCase       ,			callbacks=[logging_callback] + extra_callbacks + [InitCallback()] + [checkpoint_callback]       ,			logger=_UpperCAmelCase       ,			val_check_interval=1       ,			num_sanity_val_steps=2       ,			**_UpperCAmelCase       ,			)

  if args.do_train:
    trainer.fit(_UpperCAmelCase							)

  else:
    print('''RAG modeling tests with new set functions successfuly executed!'''							)
  return trainer







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


__UpperCamelCase     :						List[Any]						       =  logging.get_logger(__name__)

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    nn.ModuleList  ,     _UpperCAmelCase					:    nn.ModuleList  ,     _UpperCAmelCase					:    List[int]			):
 lowerCAmelCase    				= nn.ModuleList([src_layers[i] for i in layers_to_copy]			)
 assert len(_UpperCAmelCase			) == len(_UpperCAmelCase			), F'{len(_UpperCAmelCase			)} != {len(_UpperCAmelCase			)}'
 dest_layers.load_state_dict(layers_to_copy.state_dict()			)


__UpperCamelCase     :						Optional[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))},
}
__UpperCamelCase     :						int						       =  {
    # 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 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Optional[int]  ,     _UpperCAmelCase					:    List[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(_UpperCAmelCase			)			)

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    List[str]  ,     _UpperCAmelCase					:    Dict			):
 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(_UpperCAmelCase			)			)
 elif n_student == 1:
  return [n_teacher - 1]
 else:
  return LAYERS_TO_SUPERVISE[n_teacher][n_student]

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Union[str, PreTrainedModel]  ,     _UpperCAmelCase					:    Union[str, Path] = "student"  ,     _UpperCAmelCase					:    Union[int, None] = None  ,     _UpperCAmelCase					:    Union[int, None] = None  ,     _UpperCAmelCase					:    List[str]=False  ,     _UpperCAmelCase					:    Optional[Any]=None  ,     _UpperCAmelCase					:    Optional[Any]=None  ,     **_UpperCAmelCase					:    str  ,     ):
 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(_UpperCAmelCase  ,     _UpperCAmelCase			):
  AutoTokenizer.from_pretrained(_UpperCAmelCase			).save_pretrained(_UpperCAmelCase			)  # purely for convenience
  lowerCAmelCase    				= AutoModelForSeqaSeqLM.from_pretrained(_UpperCAmelCase			).eval()
 else:
  assert isinstance(_UpperCAmelCase  ,     _UpperCAmelCase			), F'teacher must be a model or string got type {type(_UpperCAmelCase			)}'
 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(_UpperCAmelCase			)

 # Copy weights
 lowerCAmelCase    				= teacher.config_class(**_UpperCAmelCase			)
 lowerCAmelCase    				= AutoModelForSeqaSeqLM.from_config(_UpperCAmelCase			)
 # 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=_UpperCAmelCase			)
 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(_UpperCAmelCase			)			), list(range(_UpperCAmelCase			)			)
  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(_UpperCAmelCase			)
  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(_UpperCAmelCase  ,     _UpperCAmelCase			)
 if d_layers_to_copy is None:
  lowerCAmelCase    				= pick_layers_to_copy(_UpperCAmelCase  ,     _UpperCAmelCase			)

 try:
  if hasattr(
      _UpperCAmelCase  ,     'prophetnet'			):  # For ProphetNet, student.model.encoder.layers is called student.prophetnet.encoder.layers
   copy_layers(teacher.prophetnet.encoder.layers  ,     student.prophetnet.encoder.layers  ,     _UpperCAmelCase			)
   copy_layers(teacher.prophetnet.decoder.layers  ,     student.prophetnet.decoder.layers  ,     _UpperCAmelCase			)
  else:
   copy_layers(teacher.model.encoder.layers  ,     student.model.encoder.layers  ,     _UpperCAmelCase			)
   copy_layers(teacher.model.decoder.layers  ,     student.model.decoder.layers  ,     _UpperCAmelCase			)
 except AttributeError:  # For t5, student.model.encoder.layers is called student.encoder.block
  copy_layers(teacher.encoder.block  ,     student.encoder.block  ,     _UpperCAmelCase			)
  copy_layers(teacher.decoder.block  ,     student.decoder.block  ,     _UpperCAmelCase			)
 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(_UpperCAmelCase			)
 # 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)





