





"""simple docstring"""
import unittest
from pathlib import Path
from tempfile import NamedTemporaryFile, TemporaryDirectory

from transformers import BertConfig, BertTokenizerFast, FeatureExtractionPipeline
from transformers.convert_graph_to_onnx import (
    convert,
    ensure_valid_input,
    generate_identified_filename,
    infer_shapes,
    quantize,
)
from transformers.testing_utils import require_tf, require_tokenizers, require_torch, slow







class       a							:



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  return None







class       a							:



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  return None







class       a							(   unittest.TestCase  ):
 snake_case__             =       [
     # (model_name, model_kwargs)
     ('''bert-base-cased''', {}),
     ('''gpt2''', {'''use_cache''': False}),  # We don't support exporting GPT2 past keys anymore
 ]



 @require_tf
 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
   self._test_export(_snake_case				,       'tf'				,       12				,       **_snake_case     )



 @require_torch
 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
   self._test_export(_snake_case				,       'pt'				,       12				,       **_snake_case     )



 @require_torch
 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  from transformers import BertModel

  lowerCAmelCase    				= ['[UNK]', '[SEP]', '[CLS]', '[PAD]', '[MASK]', 'some', 'other', 'words']
  with NamedTemporaryFile(mode='w+t'     ) as vocab_file:
   vocab_file.write('\n'.join(_snake_case     )     )
   vocab_file.flush()
   lowerCAmelCase    				= BertTokenizerFast(vocab_file.name     )

  with TemporaryDirectory() as bert_save_dir:
   lowerCAmelCase    				= BertModel(BertConfig(vocab_size=len(_snake_case     )     )     )
   model.save_pretrained(_snake_case     )
   self._test_export(_snake_case				,       'pt'				,       12				,       _snake_case     )



 @require_tf
 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
   lowerCAmelCase    				= self._test_export(_snake_case				,       'tf'				,       12				,       **_snake_case     )
   lowerCAmelCase    				= quantize(Path(_snake_case     )     )

   # Ensure the actual quantized model is not bigger than the original one
   if quantized_path.stat().st_size >= Path(_snake_case     ).stat().st_size:
    self.fail('Quantized model is bigger than initial ONNX model'     )



 @require_torch
 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
   lowerCAmelCase    				= self._test_export(_snake_case				,       'pt'				,       12				,       **_snake_case     )
   lowerCAmelCase    				= quantize(_snake_case     )

   # Ensure the actual quantized model is not bigger than the original one
   if quantized_path.stat().st_size >= Path(_snake_case     ).stat().st_size:
    self.fail('Quantized model is bigger than initial ONNX model'     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case=None				,       **_snake_case     ):







  """simple docstring"""






  try:
   # Compute path
   with TemporaryDirectory() as tempdir:
    lowerCAmelCase    				= Path(_snake_case     ).joinpath('model.onnx'     )

   # Remove folder if exists
   if path.parent.exists():
    path.parent.rmdir()

   # Export
   convert(_snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       **_snake_case     )

   return path
  except Exception as e:
   self.fail(_snake_case     )



 @require_torch
 @require_tokenizers
 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  from transformers import BertModel

  lowerCAmelCase    				= BertModel(BertConfig.from_pretrained('lysandre/tiny-bert-random'     )     )
  lowerCAmelCase    				= BertTokenizerFast.from_pretrained('lysandre/tiny-bert-random'     )
  self._test_infer_dynamic_axis(_snake_case				,       _snake_case				,       'pt'     )



 @require_tf
 @require_tokenizers
 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  from transformers import TFBertModel

  lowerCAmelCase    				= TFBertModel(BertConfig.from_pretrained('lysandre/tiny-bert-random'     )     )
  lowerCAmelCase    				= BertTokenizerFast.from_pretrained('lysandre/tiny-bert-random'     )
  self._test_infer_dynamic_axis(_snake_case				,       _snake_case				,       'tf'     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= FeatureExtractionPipeline(_snake_case				,       _snake_case     )

  lowerCAmelCase    				= ['input_ids', 'token_type_ids', 'attention_mask', 'output_0', 'output_1']
  lowerCAmelCase				,lowerCAmelCase				,lowerCAmelCase				,lowerCAmelCase    				= infer_shapes(_snake_case				,       _snake_case     )

  # Assert all variables are present
  self.assertEqual(len(_snake_case     )				,       len(_snake_case     )     )
  self.assertTrue(all(var_name in shapes for var_name in variable_names     )     )
  self.assertSequenceEqual(variable_names[:3]				,       _snake_case     )
  self.assertSequenceEqual(variable_names[3:]				,       _snake_case     )

  # Assert inputs are {0: batch, 1: sequence}
  for var_name in ["input_ids", "token_type_ids", "attention_mask"]:
   self.assertDictEqual(shapes[var_name]				,       {0: 'batch', 1: 'sequence'}     )

  # Assert outputs are {0: batch, 1: sequence} and {0: batch}
  self.assertDictEqual(shapes['output_0']				,       {0: 'batch', 1: 'sequence'}     )
  self.assertDictEqual(shapes['output_1']				,       {0: 'batch'}     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= ['input_ids', 'attention_mask', 'token_type_ids']
  lowerCAmelCase    				= {'input_ids': [1, 2, 3, 4], 'attention_mask': [0, 0, 0, 0], 'token_type_ids': [1, 1, 1, 1]}
  lowerCAmelCase				,lowerCAmelCase    				= ensure_valid_input(FuncContiguousArgs()				,       _snake_case				,       _snake_case     )

  # Should have exactly the same number of args (all are valid)
  self.assertEqual(len(_snake_case     )				,       3     )

  # Should have exactly the same input names
  self.assertEqual(set(_snake_case     )				,       set(_snake_case     )     )

  # Parameter should be reordered according to their respective place in the function:
  # (input_ids, token_type_ids, attention_mask)
  self.assertEqual(_snake_case				,       (tokens['input_ids'], tokens['token_type_ids'], tokens['attention_mask'])     )

  # Generated args are interleaved with another args (for instance parameter "past" in GPT2)
  lowerCAmelCase				,lowerCAmelCase    				= ensure_valid_input(FuncNonContiguousArgs()				,       _snake_case				,       _snake_case     )

  # Should have exactly the one arg (all before the one not provided "some_other_args")
  self.assertEqual(len(_snake_case     )				,       1     )
  self.assertEqual(len(_snake_case     )				,       1     )

  # Should have only "input_ids"
  self.assertEqual(inputs_args[0]				,       tokens['input_ids']     )
  self.assertEqual(ordered_input_names[0]				,       'input_ids'     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= generate_identified_filename(Path('/home/something/my_fake_model.onnx'     )				,       '-test'     )
  self.assertEqual('/home/something/my_fake_model-test.onnx'				,       generated.as_posix()     )












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





