
'''simple docstring'''



from __future__ import annotations

import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import is_tf_available, is_vision_available

from ...test_modeling_tf_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_tf_bert import TFBertModelTester
from ..clip.test_modeling_tf_clip import TFCLIPVisionModelTester
from ..deit.test_modeling_tf_deit import TFDeiTModelTester
from ..roberta.test_modeling_tf_roberta import TFRobertaModelTester
from ..vit.test_modeling_tf_vit import TFViTModelTester


if is_tf_available():
						from transformers import (
						    TFBertModel,
						    TFCLIPVisionModel,
						    TFDeiTModel,
						    TFRobertaModel,
						    TFVisionTextDualEncoderModel,
						    TFViTModel,
						    VisionTextDualEncoderConfig,
						)

if is_vision_available():
						from PIL import Image

						from transformers import VisionTextDualEncoderProcessor


def A (__lowerCamelCase :int						):
						if isinstance(__lowerCamelCase					,      collections.abc.Iterable						):
												return x
						return (x, x)




@require_tf
class     UpperCAmelCase_	:






						'''simple docstring'''






						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase ):

												"""simple docstring"""



												pass




						def    _lowercase      (   self ):

												"""simple docstring"""



												pass




						def    _lowercase      (   self ):

												"""simple docstring"""



												pass




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase=None       ,  **_lowercase ):

												"""simple docstring"""



												_lowerCAmelCase    		=					VisionTextDualEncoderConfig.from_vision_text_configs(_lowercase       ,  _lowercase )

												_lowerCAmelCase    		=					TFVisionTextDualEncoderModel(_lowercase )

												_lowerCAmelCase    		=					model(input_ids=_lowercase       ,  pixel_values=_lowercase       ,  attention_mask=_lowercase )

												self.assertEqual(output["""text_embeds"""].shape       ,  (input_ids.shape[0], config.projection_dim) )
												self.assertEqual(output["""image_embeds"""].shape       ,  (pixel_values.shape[0], config.projection_dim) )




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase=None       ,  **_lowercase ):

												"""simple docstring"""



												_lowerCAmelCase		,       _lowerCAmelCase    		=					self.get_vision_text_model(_lowercase       ,  _lowercase )
												_lowerCAmelCase    		=					TFVisionTextDualEncoderModel(vision_model=_lowercase       ,  text_model=_lowercase )

												_lowerCAmelCase    		=					model(input_ids=_lowercase       ,  pixel_values=_lowercase       ,  attention_mask=_lowercase )

												self.assertEqual(output["""text_embeds"""].shape       ,  (input_ids.shape[0], model.config.projection_dim) )
												self.assertEqual(output["""image_embeds"""].shape       ,  (pixel_values.shape[0], model.config.projection_dim) )




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase=None       ,  **_lowercase ):

												"""simple docstring"""



												_lowerCAmelCase		,       _lowerCAmelCase    		=					self.get_vision_text_model(_lowercase       ,  _lowercase )
												_lowerCAmelCase    		=					{"""vision_model""": vision_model, """text_model""": text_model}
												_lowerCAmelCase    		=					TFVisionTextDualEncoderModel.from_vision_text_pretrained(**_lowercase )

												_lowerCAmelCase    		=					model(input_ids=_lowercase       ,  pixel_values=_lowercase       ,  attention_mask=_lowercase )

												self.assertEqual(output["""text_embeds"""].shape       ,  (input_ids.shape[0], model.config.projection_dim) )
												self.assertEqual(output["""image_embeds"""].shape       ,  (pixel_values.shape[0], model.config.projection_dim) )




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase=None       ,  **_lowercase ):

												"""simple docstring"""



												_lowerCAmelCase		,       _lowerCAmelCase    		=					self.get_vision_text_model(_lowercase       ,  _lowercase )
												_lowerCAmelCase    		=					TFVisionTextDualEncoderModel(vision_model=_lowercase       ,  text_model=_lowercase )

												_lowerCAmelCase    		=					model(input_ids=_lowercase       ,  pixel_values=_lowercase       ,  attention_mask=_lowercase )
												_lowerCAmelCase    		=					output[0].numpy()

												with tempfile.TemporaryDirectory() as tmpdirname:
																		model.save_pretrained(_lowercase )
																		_lowerCAmelCase    		=					TFVisionTextDualEncoderModel.from_pretrained(_lowercase )

																		_lowerCAmelCase    		=					model(input_ids=_lowercase       ,  pixel_values=_lowercase       ,  attention_mask=_lowercase )
																		_lowerCAmelCase    		=					after_output[0].numpy()
																		_lowerCAmelCase    		=					np.amax(np.abs(out_a - out_a ) )
																		self.assertLessEqual(_lowercase       ,  1e-5 )




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase=None       ,  **_lowercase ):

												"""simple docstring"""



												_lowerCAmelCase		,       _lowerCAmelCase    		=					self.get_vision_text_model(_lowercase       ,  _lowercase )
												_lowerCAmelCase    		=					TFVisionTextDualEncoderModel(vision_model=_lowercase       ,  text_model=_lowercase )

												_lowerCAmelCase    		=					model(
												    input_ids=_lowercase       ,  pixel_values=_lowercase       ,  attention_mask=_lowercase       ,  output_attentions=_lowercase )

												_lowerCAmelCase    		=					output.vision_model_output.attentions
												self.assertEqual(len(_lowercase )       ,  vision_config.num_hidden_layers )

												# in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
												_lowerCAmelCase    		=					to_atuple(vision_model.config.image_size )
												_lowerCAmelCase    		=					to_atuple(vision_model.config.patch_size )
												_lowerCAmelCase    		=					(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
												_lowerCAmelCase    		=					num_patches + 1
												self.assertEqual(vision_attentions[0].shape[-3:]       ,  (vision_config.num_attention_heads, seq_len, seq_len) )

												_lowerCAmelCase    		=					output.text_model_output.attentions
												self.assertEqual(len(_lowercase )       ,  text_config.num_hidden_layers )

												self.assertEqual(
												    text_attentions[0].shape[-3:]       ,  (text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])       ,  )




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase       ,  _lowercase ):

												"""simple docstring"""



												_lowerCAmelCase    		=					np.abs((a - b) ).max()
												self.assertLessEqual(_lowercase       ,  _lowercase       ,  F'Difference between torch and flax is {diff} (>= {tol}).' )




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					self.prepare_config_and_inputs()
												self.check_vision_text_dual_encoder_model(**_lowercase )




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					self.prepare_config_and_inputs()
												self.check_model_from_pretrained_configs(**_lowercase )




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					self.prepare_config_and_inputs()
												self.check_vision_text_dual_encoder_from_pretrained(**_lowercase )




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					self.prepare_config_and_inputs()
												self.check_save_load(**_lowercase )




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					self.prepare_config_and_inputs()
												self.check_vision_text_output_attention(**_lowercase )




						@slow
						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase		,       _lowerCAmelCase    		=					self.get_pretrained_model_and_inputs()

												_lowerCAmelCase    		=					model_a(**_lowercase )
												_lowerCAmelCase    		=					outputs[0].numpy()

												with tempfile.TemporaryDirectory() as tmp_dirname:
																		model_a.save_pretrained(_lowercase )
																		_lowerCAmelCase    		=					TFVisionTextDualEncoderModel.from_pretrained(_lowercase )

																		_lowerCAmelCase    		=					model_a(**_lowercase )
																		_lowerCAmelCase    		=					after_outputs[0].numpy()
																		_lowerCAmelCase    		=					np.amax(np.abs(out_a - out_a ) )
																		self.assertLessEqual(_lowercase       ,  1e-5 )






@require_tf
class     UpperCAmelCase_	(						_SCREAMING_SNAKE_CASE			,				unittest.TestCase					):






						'''simple docstring'''






						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					TFVisionTextDualEncoderModel.from_vision_text_pretrained(
												    """hf-internal-testing/tiny-random-vit"""       ,  """hf-internal-testing/tiny-random-bert""" )
												_lowerCAmelCase    		=					13
												_lowerCAmelCase    		=					floats_tensor(
												    [
												        batch_size,
												        model.vision_model.config.num_channels,
												        model.vision_model.config.image_size,
												        model.vision_model.config.image_size,
												    ] )
												_lowerCAmelCase    		=					ids_tensor([batch_size, 4]       ,  model.text_model.config.vocab_size )
												_lowerCAmelCase    		=					random_attention_mask([batch_size, 4] )
												_lowerCAmelCase    		=					{"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

												return model, inputs




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase ):

												"""simple docstring"""



												_lowerCAmelCase    		=					TFViTModel(_lowercase       ,  name="""vision_model""" )
												_lowerCAmelCase    		=					TFBertModel(_lowercase       ,  name="""text_model""" )
												return vision_model, text_model




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					TFViTModelTester(self )
												_lowerCAmelCase    		=					TFBertModelTester(self )
												_lowerCAmelCase    		=					vit_model_tester.prepare_config_and_inputs()
												_lowerCAmelCase    		=					bert_model_tester.prepare_config_and_inputs()

												_lowerCAmelCase		,       _lowerCAmelCase		,       _lowerCAmelCase    		=					vision_config_and_inputs

												(
												    (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       
												)    		=					text_config_and_inputs

												return {
												    "text_config": text_config,
												    "vision_config": vision_config,
												    "pixel_values": pixel_values,
												    "attention_mask": input_mask,
												    "input_ids": input_ids,
												    "text_token_type_ids": token_type_ids,
												    "text_sequence_labels": sequence_labels,
												    "text_token_labels": token_labels,
												    "text_choice_labels": choice_labels,
												}






@require_tf
class     UpperCAmelCase_	(						_SCREAMING_SNAKE_CASE			,				unittest.TestCase					):






						'''simple docstring'''






						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					TFVisionTextDualEncoderModel.from_vision_text_pretrained(
												    """Rocketknight1/tiny-random-deit-tf"""       ,  """hf-internal-testing/tiny-random-roberta""" )
												_lowerCAmelCase    		=					13
												_lowerCAmelCase    		=					floats_tensor(
												    [
												        batch_size,
												        model.vision_model.config.num_channels,
												        model.vision_model.config.image_size,
												        model.vision_model.config.image_size,
												    ] )
												_lowerCAmelCase    		=					ids_tensor([batch_size, 4]       ,  model.text_model.config.vocab_size )
												_lowerCAmelCase    		=					random_attention_mask([batch_size, 4] )
												_lowerCAmelCase    		=					{"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

												return model, inputs




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase       ,  _lowercase=None       ,  **_lowercase ):

												"""simple docstring"""



												_lowerCAmelCase		,       _lowerCAmelCase    		=					self.get_vision_text_model(_lowercase       ,  _lowercase )
												_lowerCAmelCase    		=					TFVisionTextDualEncoderModel(vision_model=_lowercase       ,  text_model=_lowercase )

												_lowerCAmelCase    		=					model(
												    input_ids=_lowercase       ,  pixel_values=_lowercase       ,  attention_mask=_lowercase       ,  output_attentions=_lowercase )

												_lowerCAmelCase    		=					output.vision_model_output.attentions
												self.assertEqual(len(_lowercase )       ,  vision_config.num_hidden_layers )

												# in DEiT, the seq_len equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
												_lowerCAmelCase    		=					to_atuple(vision_model.config.image_size )
												_lowerCAmelCase    		=					to_atuple(vision_model.config.patch_size )
												_lowerCAmelCase    		=					(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
												_lowerCAmelCase    		=					num_patches + 2
												self.assertEqual(vision_attentions[0].shape[-3:]       ,  (vision_config.num_attention_heads, seq_len, seq_len) )

												_lowerCAmelCase    		=					output.text_model_output.attentions
												self.assertEqual(len(_lowercase )       ,  text_config.num_hidden_layers )

												self.assertEqual(
												    text_attentions[0].shape[-3:]       ,  (text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])       ,  )




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase ):

												"""simple docstring"""



												_lowerCAmelCase    		=					TFDeiTModel(_lowercase       ,  name="""vision_model""" )
												_lowerCAmelCase    		=					TFRobertaModel(_lowercase       ,  name="""text_model""" )
												return vision_model, text_model




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					TFDeiTModelTester(self )
												_lowerCAmelCase    		=					TFRobertaModelTester(self )
												_lowerCAmelCase    		=					vit_model_tester.prepare_config_and_inputs()
												_lowerCAmelCase    		=					bert_model_tester.prepare_config_and_inputs()

												_lowerCAmelCase		,       _lowerCAmelCase		,       _lowerCAmelCase    		=					vision_config_and_inputs

												(
												    (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       
												)    		=					text_config_and_inputs

												return {
												    "text_config": text_config,
												    "vision_config": vision_config,
												    "pixel_values": pixel_values,
												    "attention_mask": input_mask,
												    "input_ids": input_ids,
												    "text_token_type_ids": token_type_ids,
												    "text_sequence_labels": sequence_labels,
												    "text_token_labels": token_labels,
												    "text_choice_labels": choice_labels,
												}






@require_tf
class     UpperCAmelCase_	(						_SCREAMING_SNAKE_CASE			,				unittest.TestCase					):






						'''simple docstring'''






						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					TFVisionTextDualEncoderModel.from_vision_text_pretrained(
												    """Rocketknight1/tiny-random-clip-tf"""       ,  """hf-internal-testing/tiny-random-bert""" )
												_lowerCAmelCase    		=					13
												_lowerCAmelCase    		=					floats_tensor(
												    [
												        batch_size,
												        model.vision_model.config.num_channels,
												        model.vision_model.config.image_size,
												        model.vision_model.config.image_size,
												    ] )
												_lowerCAmelCase    		=					ids_tensor([batch_size, 4]       ,  model.text_model.config.vocab_size )
												_lowerCAmelCase    		=					random_attention_mask([batch_size, 4] )
												_lowerCAmelCase    		=					{"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

												return model, inputs




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase ):

												"""simple docstring"""



												_lowerCAmelCase    		=					TFCLIPVisionModel(_lowercase       ,  name="""vision_model""" )
												_lowerCAmelCase    		=					TFBertModel(_lowercase       ,  name="""text_model""" )
												return vision_model, text_model




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					TFCLIPVisionModelTester(self )
												_lowerCAmelCase    		=					TFBertModelTester(self )
												_lowerCAmelCase    		=					clip_model_tester.prepare_config_and_inputs()
												_lowerCAmelCase    		=					bert_model_tester.prepare_config_and_inputs()

												_lowerCAmelCase		,       _lowerCAmelCase    		=					vision_config_and_inputs

												(
												    (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       (
												    _lowerCAmelCase
												)		,       
												)    		=					text_config_and_inputs

												return {
												    "text_config": text_config,
												    "vision_config": vision_config,
												    "pixel_values": pixel_values,
												    "attention_mask": input_mask,
												    "input_ids": input_ids,
												    "text_token_type_ids": token_type_ids,
												    "text_sequence_labels": sequence_labels,
												    "text_token_labels": token_labels,
												    "text_choice_labels": choice_labels,
												}






@require_vision
@require_tf
class     UpperCAmelCase_	(						unittest.TestCase					):






						'''simple docstring'''






						@slow
						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					TFVisionTextDualEncoderModel.from_pretrained(
												    """clip-italian/clip-italian"""       ,  logit_scale_init_value=1.0       ,  from_pt=_lowercase )
												_lowerCAmelCase    		=					VisionTextDualEncoderProcessor.from_pretrained("""clip-italian/clip-italian""" )

												_lowerCAmelCase    		=					Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png""" )
												_lowerCAmelCase    		=					processor(
												    text=["""una foto di un gatto""", """una foto di un cane"""]       ,  images=_lowercase       ,  padding=_lowercase       ,  return_tensors="""np""" )

												_lowerCAmelCase    		=					model(**_lowercase )

												# verify the logits
												self.assertEqual(outputs.logits_per_image.shape       ,  (inputs.pixel_values.shape[0], inputs.input_ids.shape[0]) )
												self.assertEqual(
												    outputs.logits_per_text.shape       ,  (inputs.input_ids.shape[0], inputs.pixel_values.shape[0])       ,  )

												_lowerCAmelCase    		=					np.array([[1.228_4727, 0.310_4122]] )

												self.assertTrue(np.allclose(outputs.logits_per_image.numpy()       ,  _lowercase       ,  atol=1e-3 ) )







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





