


'''simple docstring'''






import argparse
import os
import re

import tensorflow as tf
import torch

from transformers import BertConfig, BertModel
from transformers.utils import logging


logging.set_verbosity_info()
__A	:	Optional[Any]     =     logging.get_logger(__name__)





def      UpperCAmelCase      (		lowerCamelCase_  :List[str]				,      lowerCamelCase_  :Dict				,      lowerCamelCase_  :List[Any] ):







		'''simple docstring'''
		snake_case_		:							Optional[Any]			       =						os.path.abspath(lowerCamelCase_ )
		logger.info(F'''Converting TensorFlow checkpoint from {tf_path}''' )
		# Load weights from TF model
		snake_case_		:							Tuple			       =						tf.train.list_variables(lowerCamelCase_ )
		snake_case_		:							List[Any]			       =						[]
		snake_case_		:							Any			       =						[]
		snake_case_		:							Union[str, Any]			       =						[]
		for full_name, shape in init_vars:
				# logger.info(f"Loading TF weight {name} with shape {shape}")
				snake_case_		:							Any			       =						full_name.split("""/""" )
				if full_name == "_CHECKPOINTABLE_OBJECT_GRAPH" or name[0] in ["global_step", "save_counter"]:
						logger.info(F'''Skipping non-model layer {full_name}''' )
						continue
				if "optimizer" in full_name:
						logger.info(F'''Skipping optimization layer {full_name}''' )
						continue
				if name[0] == "model":
						# ignore initial 'model'
						snake_case_		:							List[str]			       =						name[1:]
				# figure out how many levels deep the name is
				snake_case_		:							List[Any]			       =						0
				for _name in name:
						if _name.startswith("""layer_with_weights""" ):
								depth += 1
						else:
								break
				layer_depth.append(lowerCamelCase_ )
				# read data
				snake_case_		:							List[str]			       =						tf.train.load_variable(lowerCamelCase_				,      lowerCamelCase_ )
				names.append("""/""".join(lowerCamelCase_ ) )
				arrays.append(lowerCamelCase_ )
		logger.info(F'''Read a total of {len(lowerCamelCase_ ):,} layers''' )

		# Sanity check
		if len(set(lowerCamelCase_ ) ) != 1:
				raise ValueError(F'''Found layer names with different depths (layer depth {list(set(lowerCamelCase_ ) )})''' )
		snake_case_		:							Any			       =						list(set(lowerCamelCase_ ) )[0]
		if layer_depth != 1:
				raise ValueError(
				    """The model contains more than just the embedding/encoder layers. This script does not handle MLM/NSP"""
				    """ heads.""" )

		# convert layers
		logger.info("""Converting weights...""" )
		for full_name, array in zip(lowerCamelCase_				,      lowerCamelCase_ ):
				snake_case_		:							List[Any]			       =						full_name.split("""/""" )
				snake_case_		:							str			       =						model
				snake_case_		:							str			       =						[]
				for i, m_name in enumerate(lowerCamelCase_ ):
						if m_name == ".ATTRIBUTES":
								# variable names end with .ATTRIBUTES/VARIABLE_VALUE
								break
						if m_name.startswith("""layer_with_weights""" ):
								snake_case_		:							List[Any]			       =						int(m_name.split("""-""" )[-1] )
								if layer_num <= 2:
										# embedding layers
										# layer_num 0: word_embeddings
										# layer_num 1: position_embeddings
										# layer_num 2: token_type_embeddings
										continue
								elif layer_num == 3:
										# embedding LayerNorm
										trace.extend(["""embeddings""", """LayerNorm"""] )
										snake_case_		:							List[Any]			       =						getattr(lowerCamelCase_				,      """embeddings""" )
										snake_case_		:							Optional[Any]			       =						getattr(lowerCamelCase_				,      """LayerNorm""" )
								elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
										# encoder layers
										trace.extend(["""encoder""", """layer""", str(layer_num - 4 )] )
										snake_case_		:							Dict			       =						getattr(lowerCamelCase_				,      """encoder""" )
										snake_case_		:							int			       =						getattr(lowerCamelCase_				,      """layer""" )
										snake_case_		:							Optional[int]			       =						pointer[layer_num - 4]
								elif layer_num == config.num_hidden_layers + 4:
										# pooler layer
										trace.extend(["""pooler""", """dense"""] )
										snake_case_		:							Optional[int]			       =						getattr(lowerCamelCase_				,      """pooler""" )
										snake_case_		:							List[str]			       =						getattr(lowerCamelCase_				,      """dense""" )
						elif m_name == "embeddings":
								trace.append("""embeddings""" )
								snake_case_		:							Union[str, Any]			       =						getattr(lowerCamelCase_				,      """embeddings""" )
								if layer_num == 0:
										trace.append("""word_embeddings""" )
										snake_case_		:							Optional[int]			       =						getattr(lowerCamelCase_				,      """word_embeddings""" )
								elif layer_num == 1:
										trace.append("""position_embeddings""" )
										snake_case_		:							List[str]			       =						getattr(lowerCamelCase_				,      """position_embeddings""" )
								elif layer_num == 2:
										trace.append("""token_type_embeddings""" )
										snake_case_		:							List[Any]			       =						getattr(lowerCamelCase_				,      """token_type_embeddings""" )
								else:
										raise ValueError(F'''Unknown embedding layer with name {full_name}''' )
								trace.append("""weight""" )
								snake_case_		:							List[str]			       =						getattr(lowerCamelCase_				,      """weight""" )
						elif m_name == "_attention_layer":
								# self-attention layer
								trace.extend(["""attention""", """self"""] )
								snake_case_		:							List[Any]			       =						getattr(lowerCamelCase_				,      """attention""" )
								snake_case_		:							str			       =						getattr(lowerCamelCase_				,      """self""" )
						elif m_name == "_attention_layer_norm":
								# output attention norm
								trace.extend(["""attention""", """output""", """LayerNorm"""] )
								snake_case_		:							str			       =						getattr(lowerCamelCase_				,      """attention""" )
								snake_case_		:							Optional[Any]			       =						getattr(lowerCamelCase_				,      """output""" )
								snake_case_		:							Any			       =						getattr(lowerCamelCase_				,      """LayerNorm""" )
						elif m_name == "_attention_output_dense":
								# output attention dense
								trace.extend(["""attention""", """output""", """dense"""] )
								snake_case_		:							List[str]			       =						getattr(lowerCamelCase_				,      """attention""" )
								snake_case_		:							str			       =						getattr(lowerCamelCase_				,      """output""" )
								snake_case_		:							Optional[int]			       =						getattr(lowerCamelCase_				,      """dense""" )
						elif m_name == "_output_dense":
								# output dense
								trace.extend(["""output""", """dense"""] )
								snake_case_		:							int			       =						getattr(lowerCamelCase_				,      """output""" )
								snake_case_		:							List[str]			       =						getattr(lowerCamelCase_				,      """dense""" )
						elif m_name == "_output_layer_norm":
								# output dense
								trace.extend(["""output""", """LayerNorm"""] )
								snake_case_		:							List[Any]			       =						getattr(lowerCamelCase_				,      """output""" )
								snake_case_		:							Optional[Any]			       =						getattr(lowerCamelCase_				,      """LayerNorm""" )
						elif m_name == "_key_dense":
								# attention key
								trace.append("""key""" )
								snake_case_		:							str			       =						getattr(lowerCamelCase_				,      """key""" )
						elif m_name == "_query_dense":
								# attention query
								trace.append("""query""" )
								snake_case_		:							Any			       =						getattr(lowerCamelCase_				,      """query""" )
						elif m_name == "_value_dense":
								# attention value
								trace.append("""value""" )
								snake_case_		:							Union[str, Any]			       =						getattr(lowerCamelCase_				,      """value""" )
						elif m_name == "_intermediate_dense":
								# attention intermediate dense
								trace.extend(["""intermediate""", """dense"""] )
								snake_case_		:							Dict			       =						getattr(lowerCamelCase_				,      """intermediate""" )
								snake_case_		:							Tuple			       =						getattr(lowerCamelCase_				,      """dense""" )
						elif m_name == "_output_layer_norm":
								# output layer norm
								trace.append("""output""" )
								snake_case_		:							List[Any]			       =						getattr(lowerCamelCase_				,      """output""" )
						# weights & biases
						elif m_name in ["bias", "beta"]:
								trace.append("""bias""" )
								snake_case_		:							Tuple			       =						getattr(lowerCamelCase_				,      """bias""" )
						elif m_name in ["kernel", "gamma"]:
								trace.append("""weight""" )
								snake_case_		:							List[Any]			       =						getattr(lowerCamelCase_				,      """weight""" )
						else:
								logger.warning(F'''Ignored {m_name}''' )
        # for certain layers reshape is necessary
				snake_case_		:							int			       =						""".""".join(lowerCamelCase_ )
				if re.match(R"""(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)"""				,      lowerCamelCase_ ) or re.match(
				    R"""(\S+)\.attention\.output\.dense\.weight"""				,      lowerCamelCase_ ):
						snake_case_		:							List[str]			       =						array.reshape(pointer.data.shape )
				if "kernel" in full_name:
						snake_case_		:							Dict			       =						array.transpose()
				if pointer.shape == array.shape:
						snake_case_		:							Union[str, Any]			       =						torch.from_numpy(lowerCamelCase_ )
				else:
						raise ValueError(
						    F'''Shape mismatch in layer {full_name}: Model expects shape {pointer.shape} but layer contains shape:'''
						    F''' {array.shape}''' )
				logger.info(F'''Successfully set variable {full_name} to PyTorch layer {trace}''' )
		return model





def      UpperCAmelCase      (		lowerCamelCase_  :Any				,      lowerCamelCase_  :Optional[Any]				,      lowerCamelCase_  :List[str] ):







		'''simple docstring'''
		# Instantiate model
		logger.info(F'''Loading model based on config from {config_path}...''' )
		snake_case_		:							Tuple			       =						BertConfig.from_json_file(lowerCamelCase_ )
		snake_case_		:							Tuple			       =						BertModel(lowerCamelCase_ )

		# Load weights from checkpoint
		logger.info(F'''Loading weights from checkpoint {tf_checkpoint_path}...''' )
		load_tfa_weights_in_bert(lowerCamelCase_				,      lowerCamelCase_				,      lowerCamelCase_ )

		# Save pytorch-model
		logger.info(F'''Saving PyTorch model to {pytorch_dump_path}...''' )
		torch.save(model.state_dict()				,      lowerCamelCase_ )


if __name__ == "__main__":
							__A	:	Optional[int]     =     argparse.ArgumentParser()
							parser.add_argument(
							    '--tf_checkpoint_path', type=str, required=True, help='Path to the TensorFlow 2.x checkpoint path.'
							)
							parser.add_argument(
							    '--bert_config_file',
							    type=str,
							    required=True,
							    help='The config json file corresponding to the BERT model. This specifies the model architecture.',
							)
							parser.add_argument(
							    '--pytorch_dump_path',
							    type=str,
							    required=True,
							    help='Path to the output PyTorch model (must include filename).',
							)
							__A	:	Any     =     parser.parse_args()
							convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)



'''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


__A	:	List[Any]     =     logging.get_logger(__name__)





def      UpperCAmelCase      (		lowerCamelCase_  :nn.ModuleList				,      lowerCamelCase_  :nn.ModuleList				,      lowerCamelCase_  :List[int] ):







		'''simple docstring'''
		snake_case_		:							Optional[Any]			       =						nn.ModuleList([src_layers[i] for i in layers_to_copy] )
		assert len(lowerCamelCase_ ) == len(lowerCamelCase_ ), F'''{len(lowerCamelCase_ )} != {len(lowerCamelCase_ )}'''
		dest_layers.load_state_dict(layers_to_copy.state_dict() )


__A	:	Optional[int]     =     {
    # 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))},
}
__A	:	Dict     =     {
    # 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_  :Union[str, Any]				,      lowerCamelCase_  :int ):







		'''simple docstring'''
		try:
				snake_case_		:							int			       =						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(lowerCamelCase_ ) )





def      UpperCAmelCase      (		lowerCamelCase_  :List[Any]				,      lowerCamelCase_  :List[Any] ):







		'''simple docstring'''
		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(lowerCamelCase_ ) )
		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[str]=False				,      lowerCamelCase_  :Optional[int]=None				,      lowerCamelCase_  :Optional[Any]=None				,      **lowerCamelCase_  :Dict				,      ):







		'''simple docstring'''
		snake_case_		:							Optional[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(lowerCamelCase_				,      lowerCamelCase_ ):
				AutoTokenizer.from_pretrained(lowerCamelCase_ ).save_pretrained(lowerCamelCase_ )  # purely for convenience
				snake_case_		:							List[Any]			       =						AutoModelForSeqaSeqLM.from_pretrained(lowerCamelCase_ ).eval()
		else:
				assert isinstance(lowerCamelCase_				,      lowerCamelCase_ ), F'''teacher must be a model or string got type {type(lowerCamelCase_ )}'''
		snake_case_		:							Any			       =						teacher.config.to_diff_dict()

		try:
				snake_case_       ,      snake_case_		:							List[str]			       =						teacher.config.encoder_layers, teacher.config.decoder_layers
				if e is None:
						snake_case_		:							Dict			       =						teacher_e
				if d is None:
						snake_case_		:							List[Any]			       =						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_		:							Tuple			       =						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_		:							Optional[int]			       =						teacher_e
				if d is None:
						snake_case_		:							List[str]			       =						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(lowerCamelCase_ )

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

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