
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


lowerCamelCase      =logging.get_logger(__name__)


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ ,		UpperCamelCase__ ,		UpperCamelCase__		):
      UpperCamelCase__    :   List[str]					=				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()		)


lowerCamelCase      ={
    # maps  num layers in teacher -> num_layers in student -> which teacher layers to copy.
    # 12: bart, 16: pegasus, 6: marian/Helsinki-NLP
    1_2: {
        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, 1_1],
        4: [0, 4, 8, 1_1],
        6: [0, 2, 4, 7, 9, 1_1],
        9: [0, 1, 2, 4, 5, 7, 9, 1_0, 1_1],
        1_2: list(range(1_2)),
    },
    1_6: {  # maps  num layers in student -> which teacher layers to copy
        1: [0],
        2: [0, 1_5],
        3: [0, 8, 1_5],
        4: [0, 5, 1_0, 1_5],
        6: [0, 3, 6, 9, 1_2, 1_5],
        8: [0, 2, 4, 6, 8, 1_0, 1_2, 1_5],
        9: [0, 1, 3, 5, 7, 9, 1_1, 1_3, 1_5],
        1_2: [0, 1, 2, 3, 4, 5, 6, 7, 9, 1_1, 1_3, 1_5],
        1_6: list(range(1_6)),
    },
    6: {1: [0], 2: [0, 5], 3: [0, 2, 5], 4: [0, 1, 3, 5], 6: list(range(6))},
}
lowerCamelCase      ={
    # 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]},
    1_2: {1: [1_1], 2: [5, 1_1], 3: [3, 7, 1_1], 6: [1, 3, 5, 8, 1_0, 1_1]},
    1_6: {1: [1_5], 4: [4, 9, 1_2, 1_5], 8: [1, 3, 5, 7, 9, 1_1, 1_3, 1_5]},
}


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ ,		UpperCamelCase__		):
      try:
            UpperCamelCase__    :   Dict					=				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__ ,		UpperCamelCase__		):
      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__ ,		UpperCamelCase__ = "student" ,		UpperCamelCase__ = None ,		UpperCamelCase__ = None ,		UpperCamelCase__=False ,		UpperCamelCase__=None ,		UpperCamelCase__=None ,		**UpperCamelCase__ ,		):
      UpperCamelCase__    :   Optional[int]					=				'''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
            UpperCamelCase__    :   List[str]					=				AutoModelForSeqaSeqLM.from_pretrained(UpperCamelCase__		).eval()
      else:
            assert isinstance(UpperCamelCase__ ,		UpperCamelCase__		), f'''teacher must be a model or string got type {type(UpperCamelCase__		)}'''
      UpperCamelCase__    :   Union[str, Any]					=				teacher.config.to_diff_dict()

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




import argparse

from torch import nn

# transformers_old should correspond to branch `save_old_prophetnet_model_structure` here
# original prophetnet_checkpoints are saved under `patrickvonplaten/..._old` respectively
from transformers_old.modeling_prophetnet import (
    ProphetNetForConditionalGeneration as ProphetNetForConditionalGenerationOld,
)
from transformers_old.modeling_xlm_prophetnet import (
    XLMProphetNetForConditionalGeneration as XLMProphetNetForConditionalGenerationOld,
)

from transformers import ProphetNetForConditionalGeneration, XLMProphetNetForConditionalGeneration, logging


lowerCamelCase      =logging.get_logger(__name__)
logging.set_verbosity_info()


def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ ,		UpperCamelCase__		):
      if "xprophetnet" in prophetnet_checkpoint_path:
            UpperCamelCase__    :   int					=				XLMProphetNetForConditionalGenerationOld.from_pretrained(UpperCamelCase__		)
            UpperCamelCase__				,UpperCamelCase__    :   Union[str, Any]					=				XLMProphetNetForConditionalGeneration.from_pretrained(
                UpperCamelCase__ ,		output_loading_info=UpperCamelCase__		)
      else:
            UpperCamelCase__    :   Tuple					=				ProphetNetForConditionalGenerationOld.from_pretrained(UpperCamelCase__		)
            UpperCamelCase__				,UpperCamelCase__    :   Union[str, Any]					=				ProphetNetForConditionalGeneration.from_pretrained(
                UpperCamelCase__ ,		output_loading_info=UpperCamelCase__		)

      UpperCamelCase__    :   Optional[Any]					=				['''key_proj''', '''value_proj''', '''query_proj''']

      UpperCamelCase__    :   Tuple					=				{
          '''self_attn''': '''ngram_self_attn''',
          '''cross_attn''': '''encoder_attn''',
          '''cross_attn_layer_norm''': '''encoder_attn_layer_norm''',
          '''feed_forward_layer_norm''': '''final_layer_norm''',
          '''feed_forward''': '''''',
          '''intermediate''': '''fc1''',
          '''output''': '''fc2''',
          '''key_proj''': '''k_proj''',
          '''query_proj''': '''q_proj''',
          '''value_proj''': '''v_proj''',
          '''word_embeddings''': '''embed_tokens''',
          '''embeddings_layer_norm''': '''emb_layer_norm''',
          '''relative_pos_embeddings''': '''relative_linear''',
          '''ngram_embeddings''': '''ngram_input_embed''',
          '''position_embeddings''': '''embed_positions''',
      }

      for key in loading_info["missing_keys"]:
            UpperCamelCase__    :   Tuple					=				key.split('''.'''		)

            if attributes[0] == "lm_head":
                  UpperCamelCase__    :   Union[str, Any]					=				prophet
                  UpperCamelCase__    :   int					=				prophet_old
            else:
                  UpperCamelCase__    :   int					=				prophet.prophetnet
                  UpperCamelCase__    :   str					=				prophet_old.model

            UpperCamelCase__    :   Tuple					=				False
            for attribute in attributes:
                  if attribute in mapping:
                        UpperCamelCase__    :   List[str]					=				mapping[attribute]
                        if not hasattr(UpperCamelCase__ ,		UpperCamelCase__		) and len(UpperCamelCase__		) > 0:
                              UpperCamelCase__    :   Tuple					=				attribute
                  elif hasattr(UpperCamelCase__ ,		UpperCamelCase__		):
                        UpperCamelCase__    :   Tuple					=				attribute

                  if attribute == "weight":
                        assert old_model.weight.shape == model.weight.shape, "Shapes have to match!"
                        UpperCamelCase__    :   Dict					=				old_model.weight
                        logger.info(f'''{attribute} is initialized.'''		)
                        UpperCamelCase__    :   Optional[Any]					=				True
                        break
                  elif attribute == "bias":
                        assert old_model.bias.shape == model.bias.shape, "Shapes have to match!"
                        UpperCamelCase__    :   Any					=				old_model.bias
                        logger.info(f'''{attribute} is initialized'''		)
                        UpperCamelCase__    :   Any					=				True
                        break
                  elif attribute in special_keys and hasattr(UpperCamelCase__ ,		'''in_proj_weight'''		):
                        UpperCamelCase__    :   Optional[int]					=				old_model.in_proj_weight.shape[0] // 3
                        UpperCamelCase__    :   Optional[Any]					=				getattr(UpperCamelCase__ ,		UpperCamelCase__		)
                        param.weight.shape == old_model.in_proj_weight[:embed_dim, :].shape, "Shapes have to match"
                        param.bias.shape == old_model.in_proj_bias[:embed_dim].shape, "Shapes have to match"
                        if attribute == "query_proj":
                              UpperCamelCase__    :   Optional[Any]					=				nn.Parameter(old_model.in_proj_weight[:embed_dim, :]		)
                              UpperCamelCase__    :   str					=				nn.Parameter(old_model.in_proj_bias[:embed_dim]		)

                        elif attribute == "key_proj":
                              UpperCamelCase__    :   Union[str, Any]					=				nn.Parameter(old_model.in_proj_weight[embed_dim : 2 * embed_dim, :]		)
                              UpperCamelCase__    :   Optional[int]					=				nn.Parameter(old_model.in_proj_bias[embed_dim : 2 * embed_dim]		)
                        elif attribute == "value_proj":
                              UpperCamelCase__    :   List[str]					=				nn.Parameter(old_model.in_proj_weight[2 * embed_dim :, :]		)
                              UpperCamelCase__    :   Tuple					=				nn.Parameter(old_model.in_proj_bias[2 * embed_dim :]		)
                        UpperCamelCase__    :   Tuple					=				True
                        break
                  elif attribute == "position_embeddings":
                        assert (
                            model.position_embeddings.weight.shape[-1] == old_model.embed_positions.weight.shape[-1]
                        ), "Hidden size has to match"
                        assert model.position_embeddings.weight.shape[0] == 5_1_2, "We want 512 position_embeddings."
                        UpperCamelCase__    :   Optional[Any]					=				nn.Parameter(old_model.embed_positions.weight[:5_1_2, :]		)
                        UpperCamelCase__    :   List[Any]					=				True
                        break

                  if attribute.isdigit():
                        UpperCamelCase__    :   int					=				model[int(UpperCamelCase__		)]
                        UpperCamelCase__    :   int					=				old_model[int(UpperCamelCase__		)]
                  else:
                        UpperCamelCase__    :   str					=				getattr(UpperCamelCase__ ,		UpperCamelCase__		)

                        if old_attribute == "":
                              UpperCamelCase__    :   List[str]					=				old_model
                        else:
                              if not hasattr(UpperCamelCase__ ,		UpperCamelCase__		):
                                    raise ValueError(f'''{old_model} does not have {old_attribute}'''		)
                              UpperCamelCase__    :   List[Any]					=				getattr(UpperCamelCase__ ,		UpperCamelCase__		)

            if not is_key_init:
                  raise ValueError(f'''{key} was not correctly initialized!'''		)

      print(f'''Saving model to {pytorch_dump_folder_path}'''		)
      prophet.save_pretrained(UpperCamelCase__		)


if __name__ == "__main__":
    lowerCamelCase      =argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        "--prophetnet_checkpoint_path", default=None, type=str, required=True, help="Path the official PyTorch dump."
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
    )
    lowerCamelCase      =parser.parse_args()
    convert_prophetnet_checkpoint_to_pytorch(args.prophetnet_checkpoint_path, args.pytorch_dump_folder_path)


