
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      UpperCamelCase   (       __lowerCamelCase	:					nn.ModuleList							,   __lowerCamelCase	:					nn.ModuleList							,   __lowerCamelCase	:					List[int]       ):
       snake_case			:  Dict     							= nn.ModuleList([src_layers[i] for i in layers_to_copy]       )
       assert len(_lowercase       ) == len(_lowercase       ), f"""{len(_lowercase       )} != {len(_lowercase       )}"""
       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
    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))},
}
__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]},
    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   (       __lowerCamelCase	:					List[Any]							,   __lowerCamelCase	:					List[str]       ):




       try:
              snake_case			:  Tuple     							= 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(_lowercase       )       )


def      UpperCamelCase   (       __lowerCamelCase	:					List[Any]							,   __lowerCamelCase	:					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(_lowercase       )       )
       elif n_student == 1:
              return [n_teacher - 1]
       else:
              return LAYERS_TO_SUPERVISE[n_teacher][n_student]


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

       try:
              snake_case							,   snake_case			:  Optional[int]     							= teacher.config.encoder_layers, teacher.config.decoder_layers
              if e is None:
                     snake_case			:  Optional[Any]     							= teacher_e
              if d is None:
                     snake_case			:  str     							= teacher_d
              init_kwargs.update({"encoder_layers": e, "decoder_layers": d}       )
       except AttributeError:  # T5
              if hasattr(teacher.config							,   "num_encoder_layers"       ):
                     snake_case							,   snake_case			:  Union[str, Any]     							= teacher.config.num_encoder_layers, teacher.config.num_decoder_layers
              else:
                     snake_case							,   snake_case			:  Union[str, Any]     							= teacher.config.num_layers, teacher.config.num_decoder_layers
              if e is None:
                     snake_case			:  str     							= teacher_e
              if d is None:
                     snake_case			:  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(_lowercase       )

       # Copy weights
       snake_case			:  Dict     							= teacher.config_class(**_lowercase       )
       snake_case			:  Any     							= AutoModelForSeqaSeqLM.from_config(_lowercase       )
       # Start by copying the full teacher state dict this will copy the first N teacher layers to the student.
       snake_case			:  Optional[int]     							= student.load_state_dict(teacher.state_dict()							,   strict=_lowercase       )
       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
              snake_case							,   snake_case			:  Dict     							= list(range(_lowercase       )       ), list(range(_lowercase       )       )
              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(_lowercase       )
              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:
              snake_case			:  Optional[int]     							= pick_layers_to_copy(_lowercase							,   _lowercase       )
       if d_layers_to_copy is None:
              snake_case			:  Any     							= pick_layers_to_copy(_lowercase							,   _lowercase       )

       try:
              if hasattr(
                  _lowercase							,   "prophetnet"       ):  # For ProphetNet, student.model.encoder.layers is called student.prophetnet.encoder.layers
                     copy_layers(teacher.prophetnet.encoder.layers							,   student.prophetnet.encoder.layers							,   _lowercase       )
                     copy_layers(teacher.prophetnet.decoder.layers							,   student.prophetnet.decoder.layers							,   _lowercase       )
              else:
                     copy_layers(teacher.model.encoder.layers							,   student.model.encoder.layers							,   _lowercase       )
                     copy_layers(teacher.model.decoder.layers							,   student.model.decoder.layers							,   _lowercase       )
       except AttributeError:  # For t5, student.model.encoder.layers is called student.encoder.block
              copy_layers(teacher.encoder.block							,   student.encoder.block							,   _lowercase       )
              copy_layers(teacher.decoder.block							,   student.decoder.block							,   _lowercase       )
       logger.info(
           f"""Copied encoder layers {e_layers_to_copy} and decoder layers {d_layers_to_copy}. Saving them to {save_path}"""       )
       snake_case			:  str     							= {
           "teacher_type": teacher.config.model_type,
           "copied_encoder_layers": e_layers_to_copy,
           "copied_decoder_layers": d_layers_to_copy,
       }
       student.save_pretrained(_lowercase       )
       # 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"""



import argparse
import json
import os

import numpy as np
import PIL
import requests
import tensorflow.keras.applications.efficientnet as efficientnet
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from tensorflow.keras.preprocessing import image

from transformers import (
    EfficientNetConfig,
    EfficientNetForImageClassification,
    EfficientNetImageProcessor,
)
from transformers.utils import logging


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

lowercase_							=							{
    'b0': efficientnet.EfficientNetBa,
    'b1': efficientnet.EfficientNetBa,
    'b2': efficientnet.EfficientNetBa,
    'b3': efficientnet.EfficientNetBa,
    'b4': efficientnet.EfficientNetBa,
    'b5': efficientnet.EfficientNetBa,
    'b6': efficientnet.EfficientNetBa,
    'b7': efficientnet.EfficientNetBa,
}

lowercase_							=							{
    'b0': {
        'hidden_dim': 12_80,
        'width_coef': 1.0,
        'depth_coef': 1.0,
        'image_size': 2_24,
        'dropout_rate': 0.2,
        'dw_padding': [],
    },
    'b1': {
        'hidden_dim': 12_80,
        'width_coef': 1.0,
        'depth_coef': 1.1,
        'image_size': 2_40,
        'dropout_rate': 0.2,
        'dw_padding': [16],
    },
    'b2': {
        'hidden_dim': 14_08,
        'width_coef': 1.1,
        'depth_coef': 1.2,
        'image_size': 2_60,
        'dropout_rate': 0.3,
        'dw_padding': [5, 8, 16],
    },
    'b3': {
        'hidden_dim': 15_36,
        'width_coef': 1.2,
        'depth_coef': 1.4,
        'image_size': 3_00,
        'dropout_rate': 0.3,
        'dw_padding': [5, 18],
    },
    'b4': {
        'hidden_dim': 17_92,
        'width_coef': 1.4,
        'depth_coef': 1.8,
        'image_size': 3_80,
        'dropout_rate': 0.4,
        'dw_padding': [6],
    },
    'b5': {
        'hidden_dim': 20_48,
        'width_coef': 1.6,
        'depth_coef': 2.2,
        'image_size': 4_56,
        'dropout_rate': 0.4,
        'dw_padding': [13, 27],
    },
    'b6': {
        'hidden_dim': 23_04,
        'width_coef': 1.8,
        'depth_coef': 2.6,
        'image_size': 5_28,
        'dropout_rate': 0.5,
        'dw_padding': [31],
    },
    'b7': {
        'hidden_dim': 25_60,
        'width_coef': 2.0,
        'depth_coef': 3.1,
        'image_size': 6_00,
        'dropout_rate': 0.5,
        'dw_padding': [18],
    },
}
def      UpperCAmelCase     (       _lowercase :   Union[str, Any]	)						-> List[Any]:


	"""simple docstring"""




	lowerCAmelCase_      							=			EfficientNetConfig()
	lowerCAmelCase_      							=			CONFIG_MAP[model_name]['''hidden_dim''']
	lowerCAmelCase_      							=			CONFIG_MAP[model_name]['''width_coef''']
	lowerCAmelCase_      							=			CONFIG_MAP[model_name]['''depth_coef''']
	lowerCAmelCase_      							=			CONFIG_MAP[model_name]['''image_size''']
	lowerCAmelCase_      							=			CONFIG_MAP[model_name]['''dropout_rate''']
	lowerCAmelCase_      							=			CONFIG_MAP[model_name]['''dw_padding''']

	lowerCAmelCase_      							=			'''huggingface/label-files'''
	lowerCAmelCase_      							=			'''imagenet-1k-id2label.json'''
	lowerCAmelCase_      							=			1_0_0_0
	lowerCAmelCase_      							=			json.load(open(hf_hub_download(_lowercase			,   _lowercase			,   repo_type='''dataset'''	)			,   '''r'''	)	)
	lowerCAmelCase_      							=			{int(_lowercase	): v for k, v in idalabel.items()}

	lowerCAmelCase_      							=			idalabel
	lowerCAmelCase_      							=			{v: k for k, v in idalabel.items()}
	return config
def      UpperCAmelCase     (       )						-> List[str]:


	"""simple docstring"""




	lowerCAmelCase_      							=			'''http://images.cocodataset.org/val2017/000000039769.jpg'''
	lowerCAmelCase_      							=			Image.open(requests.get(_lowercase			,   stream=_lowercase	).raw	)
	return im
def      UpperCAmelCase     (       _lowercase :   Dict	)						-> int:


	"""simple docstring"""




	lowerCAmelCase_      							=			CONFIG_MAP[model_name]['''image_size''']
	lowerCAmelCase_      							=			EfficientNetImageProcessor(
	    size={'''height''': size, '''width''': size}			,   image_mean=[0.4_8_5, 0.4_5_6, 0.4_0_6]			,   image_std=[0.4_7_8_5_3_9_4_4, 0.4_7_3_2_8_6_4, 0.4_7_4_3_4_1_6_3]			,   do_center_crop=_lowercase			,   )
	return preprocessor
def      UpperCAmelCase     (       _lowercase :   List[Any]	)						-> int:


	"""simple docstring"""




	lowerCAmelCase_      							=			[v.split('''_'''	)[0].split('''block'''	)[1] for v in original_param_names if v.startswith('''block'''	)]
	lowerCAmelCase_      							=			sorted(set(_lowercase	)	)
	lowerCAmelCase_      							=			len(_lowercase	)
	lowerCAmelCase_      							=			{b: str(_lowercase	) for b, i in zip(_lowercase			,   range(_lowercase	)	)}

	lowerCAmelCase_      							=			[]
	rename_keys.append(('''stem_conv/kernel:0''', '''embeddings.convolution.weight''')	)
	rename_keys.append(('''stem_bn/gamma:0''', '''embeddings.batchnorm.weight''')	)
	rename_keys.append(('''stem_bn/beta:0''', '''embeddings.batchnorm.bias''')	)
	rename_keys.append(('''stem_bn/moving_mean:0''', '''embeddings.batchnorm.running_mean''')	)
	rename_keys.append(('''stem_bn/moving_variance:0''', '''embeddings.batchnorm.running_var''')	)

	for b in block_names:
		lowerCAmelCase_      							=			block_name_mapping[b]
		rename_keys.append((F"""block{b}_expand_conv/kernel:0""", F"""encoder.blocks.{hf_b}.expansion.expand_conv.weight""")	)
		rename_keys.append((F"""block{b}_expand_bn/gamma:0""", F"""encoder.blocks.{hf_b}.expansion.expand_bn.weight""")	)
		rename_keys.append((F"""block{b}_expand_bn/beta:0""", F"""encoder.blocks.{hf_b}.expansion.expand_bn.bias""")	)
		rename_keys.append(
		    (F"""block{b}_expand_bn/moving_mean:0""", F"""encoder.blocks.{hf_b}.expansion.expand_bn.running_mean""")	)
		rename_keys.append(
		    (F"""block{b}_expand_bn/moving_variance:0""", F"""encoder.blocks.{hf_b}.expansion.expand_bn.running_var""")	)
		rename_keys.append(
		    (F"""block{b}_dwconv/depthwise_kernel:0""", F"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight""")	)
		rename_keys.append((F"""block{b}_bn/gamma:0""", F"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight""")	)
		rename_keys.append((F"""block{b}_bn/beta:0""", F"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias""")	)
		rename_keys.append(
		    (F"""block{b}_bn/moving_mean:0""", F"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean""")	)
		rename_keys.append(
		    (F"""block{b}_bn/moving_variance:0""", F"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var""")	)

		rename_keys.append((F"""block{b}_se_reduce/kernel:0""", F"""encoder.blocks.{hf_b}.squeeze_excite.reduce.weight""")	)
		rename_keys.append((F"""block{b}_se_reduce/bias:0""", F"""encoder.blocks.{hf_b}.squeeze_excite.reduce.bias""")	)
		rename_keys.append((F"""block{b}_se_expand/kernel:0""", F"""encoder.blocks.{hf_b}.squeeze_excite.expand.weight""")	)
		rename_keys.append((F"""block{b}_se_expand/bias:0""", F"""encoder.blocks.{hf_b}.squeeze_excite.expand.bias""")	)
		rename_keys.append(
		    (F"""block{b}_project_conv/kernel:0""", F"""encoder.blocks.{hf_b}.projection.project_conv.weight""")	)
		rename_keys.append((F"""block{b}_project_bn/gamma:0""", F"""encoder.blocks.{hf_b}.projection.project_bn.weight""")	)
		rename_keys.append((F"""block{b}_project_bn/beta:0""", F"""encoder.blocks.{hf_b}.projection.project_bn.bias""")	)
		rename_keys.append(
		    (F"""block{b}_project_bn/moving_mean:0""", F"""encoder.blocks.{hf_b}.projection.project_bn.running_mean""")	)
		rename_keys.append(
		    (F"""block{b}_project_bn/moving_variance:0""", F"""encoder.blocks.{hf_b}.projection.project_bn.running_var""")	)

	rename_keys.append(('''top_conv/kernel:0''', '''encoder.top_conv.weight''')	)
	rename_keys.append(('''top_bn/gamma:0''', '''encoder.top_bn.weight''')	)
	rename_keys.append(('''top_bn/beta:0''', '''encoder.top_bn.bias''')	)
	rename_keys.append(('''top_bn/moving_mean:0''', '''encoder.top_bn.running_mean''')	)
	rename_keys.append(('''top_bn/moving_variance:0''', '''encoder.top_bn.running_var''')	)

	lowerCAmelCase_      							=			{}
	for item in rename_keys:
		if item[0] in original_param_names:
			lowerCAmelCase_      							=			'''efficientnet.''' + item[1]

	lowerCAmelCase_      							=			'''classifier.weight'''
	lowerCAmelCase_      							=			'''classifier.bias'''
	return key_mapping
def      UpperCAmelCase     (       _lowercase :   Dict			,   _lowercase :   List[str]			,   _lowercase :   str	)						-> Any:


	"""simple docstring"""




	for key, value in tf_params.items():
		if "normalization" in key:
			continue

		lowerCAmelCase_      							=			key_mapping[key]
		if "_conv" in key and "kernel" in key:
			lowerCAmelCase_      							=			torch.from_numpy(_lowercase	).permute(3			,   2			,   0			,   1	)
		elif "depthwise_kernel" in key:
			lowerCAmelCase_      							=			torch.from_numpy(_lowercase	).permute(2			,   3			,   0			,   1	)
		elif "kernel" in key:
			lowerCAmelCase_      							=			torch.from_numpy(np.transpose(_lowercase	)	)
		else:
			lowerCAmelCase_      							=			torch.from_numpy(_lowercase	)

		# Replace HF parameters with original TF model parameters
		assert hf_params[hf_key].shape == new_hf_value.shape
		hf_params[hf_key].copy_(_lowercase	)
@torch.no_grad()
def      UpperCAmelCase     (       _lowercase :   str			,   _lowercase :   Any			,   _lowercase :   List[Any]			,   _lowercase :   Union[str, Any]	)						-> Optional[int]:


	"""simple docstring"""




	lowerCAmelCase_      							=			model_classes[model_name](
	    include_top=_lowercase			,   weights='''imagenet'''			,   input_tensor=_lowercase			,   input_shape=_lowercase			,   pooling=_lowercase			,   classes=1_0_0_0			,   classifier_activation='''softmax'''			,   )

	lowerCAmelCase_      							=			original_model.trainable_variables
	lowerCAmelCase_      							=			original_model.non_trainable_variables
	lowerCAmelCase_      							=			{param.name: param.numpy() for param in tf_params}
	for param in tf_non_train_params:
		lowerCAmelCase_      							=			param.numpy()
	lowerCAmelCase_      							=			list(tf_params.keys()	)

	# Load HuggingFace model
	lowerCAmelCase_      							=			get_efficientnet_config(_lowercase	)
	lowerCAmelCase_      							=			EfficientNetForImageClassification(_lowercase	).eval()
	lowerCAmelCase_      							=			hf_model.state_dict()

	# Create src-to-dst parameter name mapping dictionary
	print('''Converting parameters...'''	)
	lowerCAmelCase_      							=			rename_keys(_lowercase	)
	replace_params(_lowercase			,   _lowercase			,   _lowercase	)

	# Initialize preprocessor and preprocess input image
	lowerCAmelCase_      							=			convert_image_processor(_lowercase	)
	lowerCAmelCase_      							=			preprocessor(images=prepare_img()			,   return_tensors='''pt'''	)

	# HF model inference
	hf_model.eval()
	with torch.no_grad():
		lowerCAmelCase_      							=			hf_model(**_lowercase	)
	lowerCAmelCase_      							=			outputs.logits.detach().numpy()

	# Original model inference
	lowerCAmelCase_      							=			False
	lowerCAmelCase_      							=			CONFIG_MAP[model_name]['''image_size''']
	lowerCAmelCase_      							=			prepare_img().resize((image_size, image_size)			,   resample=PIL.Image.NEAREST	)
	lowerCAmelCase_      							=			image.img_to_array(_lowercase	)
	lowerCAmelCase_      							=			np.expand_dims(_lowercase			,   axis=0	)
	lowerCAmelCase_      							=			original_model.predict(_lowercase	)

	# Check whether original and HF model outputs match  -> np.allclose
	assert np.allclose(_lowercase			,   _lowercase			,   atol=1E-3	), "The predicted logits are not the same."
	print('''Model outputs match!'''	)

	if save_model:
		# Create folder to save model
		if not os.path.isdir(_lowercase	):
			os.mkdir(_lowercase	)
		# Save converted model and image processor
		hf_model.save_pretrained(_lowercase	)
		preprocessor.save_pretrained(_lowercase	)

	if push_to_hub:
		# Push model and image processor to hub
		print(F"""Pushing converted {model_name} to the hub..."""	)
		lowerCAmelCase_      							=			F"""efficientnet-{model_name}"""
		preprocessor.push_to_hub(_lowercase	)
		hf_model.push_to_hub(_lowercase	)


if __name__ == "__main__":
		lowercase_							=							argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '--model_name',
		    default='b0',
		    type=str,
		    help='Version name of the EfficientNet model you want to convert, select from [b0, b1, b2, b3, b4, b5, b6, b7].',
		)
		parser.add_argument(
		    '--pytorch_dump_folder_path',
		    default='hf_model',
		    type=str,
		    help='Path to the output PyTorch model directory.',
		)
		parser.add_argument('--save_model', action='store_true', help='Save model to local')
		parser.add_argument('--push_to_hub', action='store_true', help='Push model and image processor to the hub')

		lowercase_							=							parser.parse_args()
		convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)