

'''simple docstring'''

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

import fire
from torch import nn

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


lowercase__         =      logging.get_logger(__name__)




def 				UpperCamelCase(					UpperCAmelCase_  ,	UpperCAmelCase_  ,	UpperCAmelCase_     ):
 UpperCAmelCase	:     List[Any]						 =			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()     )


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




def 				UpperCamelCase(					UpperCAmelCase_  ,	UpperCAmelCase_     ):
 try:
  UpperCAmelCase	:     Optional[Any]						 =			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 				UpperCamelCase(					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 				UpperCamelCase(					UpperCAmelCase_  ,	UpperCAmelCase_ = "student"  ,	UpperCAmelCase_ = None  ,	UpperCAmelCase_ = None  ,	UpperCAmelCase_=False  ,	UpperCAmelCase_=None  ,	UpperCAmelCase_=None  ,	**UpperCAmelCase_  ,	):
 UpperCAmelCase	:     Any						 =			'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	:     Optional[int]						 =			AutoModelForSeqaSeqLM.from_pretrained(UpperCAmelCase_     ).eval()
 else:
  assert isinstance(UpperCAmelCase_  ,	UpperCAmelCase_     ), F"""teacher must be a model or string got type {type(UpperCAmelCase_     )}"""
 UpperCAmelCase	:     Optional[Any]						 =			teacher.config.to_diff_dict()

 try:
  UpperCAmelCase ,			UpperCAmelCase	:     str						 =			teacher.config.encoder_layers, teacher.config.decoder_layers
  if e is None:
   UpperCAmelCase	:     Union[str, Any]						 =			teacher_e
  if d is None:
   UpperCAmelCase	:     Optional[Any]						 =			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	:     Dict						 =			teacher.config.num_layers, teacher.config.num_decoder_layers
  if e is None:
   UpperCAmelCase	:     List[str]						 =			teacher_e
  if d is None:
   UpperCAmelCase	:     Tuple						 =			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	:     Dict						 =			teacher.config_class(**UpperCAmelCase_     )
 UpperCAmelCase	:     List[Any]						 =			AutoModelForSeqaSeqLM.from_config(UpperCAmelCase_     )
 # Start by copying the full teacher state dict this will copy the first N teacher layers to the student.
 UpperCAmelCase	:     List[Any]						 =			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	:     List[str]						 =			{
     '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)







'''simple docstring'''

from collections import OrderedDict
from typing import Any, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...feature_extraction_utils import FeatureExtractionMixin
from ...onnx import OnnxConfig
from ...onnx.utils import compute_effective_axis_dimension
from ...tokenization_utils_base import PreTrainedTokenizerBase
from ...utils import TensorType, logging


lowercase__         =      logging.get_logger(__name__)

lowercase__         =      {
    "deepmind/language-perceiver": "https://huggingface.co/deepmind/language-perceiver/resolve/main/config.json",
    # See all Perceiver models at https://huggingface.co/models?filter=perceiver
}




class 	A_			(    _snake_case							):






     '''simple docstring'''



     UpperCAmelCase_						:     Optional[int]  		=  """perceiver"""


     def __init__(		self			:					str					,      lowercase_			:					List[str]=256					,      lowercase_			:					List[str]=1_280					,      lowercase_			:					str=768					,      lowercase_			:					Tuple=1					,      lowercase_			:					str=26					,      lowercase_			:					List[Any]=8					,      lowercase_			:					int=8					,      lowercase_			:					List[str]=None					,      lowercase_			:					Dict=None					,      lowercase_			:					int="kv"					,      lowercase_			:					Union[str, Any]=1					,      lowercase_			:					List[str]=1					,      lowercase_			:					Any="gelu"					,      lowercase_			:					Optional[Any]=0.1					,      lowercase_			:					str=0.02					,      lowercase_			:					Optional[Any]=1E-12					,      lowercase_			:					Any=True					,      lowercase_			:					Union[str, Any]=262					,      lowercase_			:					Union[str, Any]=2_048					,      lowercase_			:					Optional[int]=56					,      lowercase_			:					int=[368, 496]					,      lowercase_			:					str=16					,      lowercase_			:					Optional[int]=1_920					,      lowercase_			:					Tuple=16					,      lowercase_			:					int=[1, 16, 224, 224]					,      **lowercase_			:					Union[str, Any]					,      )						->     List[str]:
      super().__init__(**lowercase_	)

      UpperCAmelCase	:     Union[str, Any]						 =			num_latents
      UpperCAmelCase	:     List[Any]						 =			d_latents
      UpperCAmelCase	:     Dict						 =			d_model
      UpperCAmelCase	:     Dict						 =			num_blocks
      UpperCAmelCase	:     Optional[int]						 =			num_self_attends_per_block
      UpperCAmelCase	:     Optional[Any]						 =			num_self_attention_heads
      UpperCAmelCase	:     Optional[Any]						 =			num_cross_attention_heads
      UpperCAmelCase	:     Tuple						 =			qk_channels
      UpperCAmelCase	:     Optional[int]						 =			v_channels
      UpperCAmelCase	:     str						 =			cross_attention_shape_for_attention
      UpperCAmelCase	:     Union[str, Any]						 =			self_attention_widening_factor
      UpperCAmelCase	:     List[Any]						 =			cross_attention_widening_factor
      UpperCAmelCase	:     Optional[Any]						 =			hidden_act
      UpperCAmelCase	:     int						 =			attention_probs_dropout_prob
      UpperCAmelCase	:     Optional[int]						 =			initializer_range
      UpperCAmelCase	:     Dict						 =			layer_norm_eps
      UpperCAmelCase	:     Any						 =			use_query_residual
      # masked language modeling attributes
      UpperCAmelCase	:     Any						 =			vocab_size
      UpperCAmelCase	:     List[Any]						 =			max_position_embeddings
      # image classification attributes
      UpperCAmelCase	:     str						 =			image_size
      # flow attributes
      UpperCAmelCase	:     Any						 =			train_size
      # multimodal autoencoding attributes
      UpperCAmelCase	:     Any						 =			num_frames
      UpperCAmelCase	:     List[Any]						 =			audio_samples_per_frame
      UpperCAmelCase	:     Tuple						 =			samples_per_patch
      UpperCAmelCase	:     Union[str, Any]						 =			output_shape




class 	A_			(    _snake_case							):






     '''simple docstring'''



     @property
     def        UpperCAmelCase_	(		self			:					Union[str, Any]	)						->     Mapping[str, Mapping[int, str]]:
      if self.task == "multiple-choice":
       UpperCAmelCase	:     Tuple						 =			{0: 'batch', 1: 'choice', 2: 'sequence'}
      else:
       UpperCAmelCase	:     Tuple						 =			{0: 'batch', 1: 'sequence'}
      return OrderedDict(
          [
              ('inputs', dynamic_axis),
              ('attention_mask', dynamic_axis),
          ]	)


     @property
     def        UpperCAmelCase_	(		self			:					str	)						->     float:
      return 1E-4




     def        UpperCAmelCase_	(		self			:					str					,      lowercase_			:					Union["PreTrainedTokenizerBase", "FeatureExtractionMixin"]					,      lowercase_			:					int = -1					,      lowercase_			:					int = -1					,      lowercase_			:					int = -1					,      lowercase_			:					bool = False					,      lowercase_			:					Optional[TensorType] = None					,      lowercase_			:					int = 3					,      lowercase_			:					int = 40					,      lowercase_			:					int = 40					,      )						->     Mapping[str, Any]:
      # copied from `transformers.onnx.config.OnnxConfig` and slightly altered/simplified

      if isinstance(lowercase_					,      lowercase_	):
       # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
       UpperCAmelCase	:     Tuple						 =			compute_effective_axis_dimension(
           lowercase_					,      fixed_dimension=OnnxConfig.default_fixed_batch					,      num_token_to_add=0	)
       # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
       UpperCAmelCase	:     int						 =			preprocessor.num_special_tokens_to_add(lowercase_	)
       UpperCAmelCase	:     Union[str, Any]						 =			compute_effective_axis_dimension(
           lowercase_					,      fixed_dimension=OnnxConfig.default_fixed_sequence					,      num_token_to_add=lowercase_	)
       # Generate dummy inputs according to compute batch and sequence
       UpperCAmelCase	:     List[Any]						 =			[' '.join(['a']	) * seq_length] * batch_size
       UpperCAmelCase	:     Union[str, Any]						 =			dict(preprocessor(lowercase_					,      return_tensors=lowercase_	)	)
       UpperCAmelCase	:     Union[str, Any]						 =			inputs.pop('input_ids'	)
       return inputs
      elif isinstance(lowercase_					,      lowercase_	) and preprocessor.model_input_names[0] == "pixel_values":
       # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
       UpperCAmelCase	:     Tuple						 =			compute_effective_axis_dimension(lowercase_					,      fixed_dimension=OnnxConfig.default_fixed_batch	)
       UpperCAmelCase	:     Any						 =			self._generate_dummy_images(lowercase_					,      lowercase_					,      lowercase_					,      lowercase_	)
       UpperCAmelCase	:     Tuple						 =			dict(preprocessor(images=lowercase_					,      return_tensors=lowercase_	)	)
       UpperCAmelCase	:     Dict						 =			inputs.pop('pixel_values'	)
       return inputs
      else:
       raise ValueError(
           'Unable to generate dummy inputs for the model. Please provide a tokenizer or a preprocessor.'	)




