




'''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  :				Optional[Any]									=  logging.get_logger(__name__)


def 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase,      __UpperCAmelCase	)		->					None:

					'''simple docstring'''

					snake_case_									=				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()	)


a  :				Union[str, 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))},
}
a  :				List[str]									=  {
    # 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 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase	)		->					Any:

					'''simple docstring'''

					try:
										snake_case_									=				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 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase	)		->					List[int]:

					'''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(__UpperCAmelCase	)	)
					elif n_student == 1:
										return [n_teacher - 1]
					else:
										return LAYERS_TO_SUPERVISE[n_teacher][n_student]


def 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase = "student",      __UpperCAmelCase = None,      __UpperCAmelCase = None,      __UpperCAmelCase=False,      __UpperCAmelCase=None,      __UpperCAmelCase=None,      **__UpperCAmelCase,      )		->					Tuple[PreTrainedModel, List[int], List[int]]:

					'''simple docstring'''

					snake_case_									=				'''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
										snake_case_									=				AutoModelForSeqaSeqLM.from_pretrained(__UpperCAmelCase	).eval()
					else:
										assert isinstance(__UpperCAmelCase,      __UpperCAmelCase	), F"teacher must be a model or string got type {type(__UpperCAmelCase	)}"
					snake_case_									=				teacher.config.to_diff_dict()

					try:
										snake_case_			,snake_case_									=				teacher.config.encoder_layers, teacher.config.decoder_layers
										if e is None:
															snake_case_									=				teacher_e
										if d is None:
															snake_case_									=				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_									=				teacher.config.num_encoder_layers, teacher.config.num_decoder_layers
										else:
															snake_case_			,snake_case_									=				teacher.config.num_layers, teacher.config.num_decoder_layers
										if e is None:
															snake_case_									=				teacher_e
										if d is None:
															snake_case_									=				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
					snake_case_									=				teacher.config_class(**__UpperCAmelCase	)
					snake_case_									=				AutoModelForSeqaSeqLM.from_config(__UpperCAmelCase	)
					# Start by copying the full teacher state dict this will copy the first N teacher layers to the student.
					snake_case_									=				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
										snake_case_			,snake_case_									=				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:
										snake_case_									=				pick_layers_to_copy(__UpperCAmelCase,      __UpperCAmelCase	)
					if d_layers_to_copy is None:
										snake_case_									=				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}"	)
					snake_case_									=				{
					    '''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'''






import torch
from torch import nn






class 			lowercase_	(nn.Module      ):




  """simple docstring"""




  def __init__(      self   :     Optional[int]					,lowercase__   :     List[str]					,lowercase__   :     Any					,lowercase__   :     Union[str, Any]					,lowercase__   :     Optional[int]					,lowercase__   :     Optional[int]=1					,lowercase__   :     Optional[Any]=False			):
   super().__init__()

   __lowercase							  = n_token
   __lowercase							  = d_embed
   __lowercase							  = d_proj

   __lowercase							  = cutoffs + [n_token]
   __lowercase							  = [0] + self.cutoffs
   __lowercase							  = div_val

   __lowercase							  = self.cutoffs[0]
   __lowercase							  = len(self.cutoffs			) - 1
   __lowercase							  = self.shortlist_size + self.n_clusters

   if self.n_clusters > 0:
    __lowercase							  = nn.Parameter(torch.zeros(self.n_clusters					,self.d_embed			)			)
    __lowercase							  = nn.Parameter(torch.zeros(self.n_clusters			)			)

   __lowercase							  = nn.ModuleList()
   __lowercase							  = nn.ParameterList()

   if div_val == 1:
    for i in range(len(self.cutoffs			)			):
     if d_proj != d_embed:
      self.out_projs.append(nn.Parameter(torch.FloatTensor(lowercase__					,lowercase__			)			)			)
     else:
      self.out_projs.append(lowercase__			)

    self.out_layers.append(nn.Linear(lowercase__					,lowercase__			)			)
   else:
    for i in range(len(self.cutoffs			)			):
     __lowercase  ,	__lowercase							  = self.cutoff_ends[i], self.cutoff_ends[i + 1]
     __lowercase							  = d_embed // (div_val**i)

     self.out_projs.append(nn.Parameter(torch.FloatTensor(lowercase__					,lowercase__			)			)			)

     self.out_layers.append(nn.Linear(lowercase__					,r_idx - l_idx			)			)

   __lowercase							  = keep_order






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]					,lowercase__   :     List[Any]					,lowercase__   :     Dict					,lowercase__   :     List[Any]					,lowercase__   :     Any			):
   if proj is None:
    __lowercase							  = nn.functional.linear(lowercase__					,lowercase__					,bias=lowercase__			)
   else:
    # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
    __lowercase							  = nn.functional.linear(lowercase__					,proj.t().contiguous()			)
    __lowercase							  = nn.functional.linear(lowercase__					,lowercase__					,bias=lowercase__			)
    # else:
    #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
    #     if bias is not None:
    #         logit = logit + bias

   return logit






  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]					,lowercase__   :     Optional[Any]					,lowercase__   :     Any=None					,lowercase__   :     List[str]=False			):

   if labels is not None:
    # Shift so that tokens < n predict n
    __lowercase							  = hidden[..., :-1, :].contiguous()
    __lowercase							  = labels[..., 1:].contiguous()
    __lowercase							  = hidden.view(-1					,hidden.size(-1			)			)
    __lowercase							  = labels.view(-1			)
    if hidden.size(0			) != labels.size(0			):
     raise RuntimeError('''Input and labels should have the same size in the batch dimension.'''			)
   else:
    __lowercase							  = hidden.view(-1					,hidden.size(-1			)			)

   if self.n_clusters == 0:
    __lowercase							  = self._compute_logit(lowercase__					,self.out_layers[0].weight					,self.out_layers[0].bias					,self.out_projs[0]			)
    if labels is not None:
     __lowercase							  = labels != -1_0_0
     __lowercase							  = torch.zeros_like(lowercase__					,dtype=hidden.dtype					,device=hidden.device			)
     __lowercase							  = (
         -nn.functional.log_softmax(lowercase__					,dim=-1			)[mask].gather(1					,labels[mask].unsqueeze(1			)			).squeeze(1			)
     )
    else:
     __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=-1			)
   else:
    # construct weights and biases
    __lowercase  ,	__lowercase							  = [], []
    for i in range(len(self.cutoffs			)			):
     if self.div_val == 1:
      __lowercase  ,	__lowercase							  = self.cutoff_ends[i], self.cutoff_ends[i + 1]
      __lowercase							  = self.out_layers[0].weight[l_idx:r_idx]
      __lowercase							  = self.out_layers[0].bias[l_idx:r_idx]
     else:
      __lowercase							  = self.out_layers[i].weight
      __lowercase							  = self.out_layers[i].bias

     if i == 0:
      __lowercase							  = torch.cat([weight_i, self.cluster_weight]					,dim=0			)
      __lowercase							  = torch.cat([bias_i, self.cluster_bias]					,dim=0			)

     weights.append(lowercase__			)
     biases.append(lowercase__			)

    __lowercase  ,	__lowercase  ,	__lowercase							  = weights[0], biases[0], self.out_projs[0]

    __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)
    __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=1			)

    if labels is None:
     __lowercase							  = hidden.new_empty((head_logit.size(0			), self.n_token)			)
    else:
     __lowercase							  = torch.zeros_like(lowercase__					,dtype=hidden.dtype					,device=hidden.device			)

    __lowercase							  = 0
    __lowercase							  = [0] + self.cutoffs
    for i in range(len(lowercase__			) - 1			):
     __lowercase  ,	__lowercase							  = cutoff_values[i], cutoff_values[i + 1]

     if labels is not None:
      __lowercase							  = (labels >= l_idx) & (labels < r_idx)
      __lowercase							  = mask_i.nonzero().squeeze()

      if indices_i.numel() == 0:
       continue

      __lowercase							  = labels.index_select(0					,lowercase__			) - l_idx
      __lowercase							  = head_logprob.index_select(0					,lowercase__			)
      __lowercase							  = hidden.index_select(0					,lowercase__			)
     else:
      __lowercase							  = hidden

     if i == 0:
      if labels is not None:
       __lowercase							  = head_logprob_i.gather(1					,target_i[:, None]			).squeeze(1			)
      else:
       __lowercase							  = head_logprob[:, : self.cutoffs[0]]
     else:
      __lowercase  ,	__lowercase  ,	__lowercase							  = weights[i], biases[i], self.out_projs[i]

      __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)
      __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=1			)
      __lowercase							  = self.cutoffs[0] + i - 1  # No probability for the head cluster
      if labels is not None:
       __lowercase							  = head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
           1					,target_i[:, None]			).squeeze(1			)
      else:
       __lowercase							  = head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
       __lowercase							  = logprob_i

     if labels is not None:
      if (hasattr(self					,'''keep_order'''			) and self.keep_order) or keep_order:
       out.index_copy_(0					,lowercase__					,-logprob_i			)
      else:
       out[offset : offset + logprob_i.size(0			)].copy_(-logprob_i			)
      offset += logprob_i.size(0			)

   return out






  def 					SCREAMING_SNAKE_CASE      (      self   :     Any					,lowercase__   :     Union[str, Any]			):
   if self.n_clusters == 0:
    __lowercase							  = self._compute_logit(lowercase__					,self.out_layers[0].weight					,self.out_layers[0].bias					,self.out_projs[0]			)
    return nn.functional.log_softmax(lowercase__					,dim=-1			)
   else:
    # construct weights and biases
    __lowercase  ,	__lowercase							  = [], []
    for i in range(len(self.cutoffs			)			):
     if self.div_val == 1:
      __lowercase  ,	__lowercase							  = self.cutoff_ends[i], self.cutoff_ends[i + 1]
      __lowercase							  = self.out_layers[0].weight[l_idx:r_idx]
      __lowercase							  = self.out_layers[0].bias[l_idx:r_idx]
     else:
      __lowercase							  = self.out_layers[i].weight
      __lowercase							  = self.out_layers[i].bias

     if i == 0:
      __lowercase							  = torch.cat([weight_i, self.cluster_weight]					,dim=0			)
      __lowercase							  = torch.cat([bias_i, self.cluster_bias]					,dim=0			)

     weights.append(lowercase__			)
     biases.append(lowercase__			)

    __lowercase  ,	__lowercase  ,	__lowercase							  = weights[0], biases[0], self.out_projs[0]
    __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)

    __lowercase							  = hidden.new_empty((head_logit.size(0			), self.n_token)			)
    __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=1			)

    __lowercase							  = [0] + self.cutoffs
    for i in range(len(lowercase__			) - 1			):
     __lowercase  ,	__lowercase							  = cutoff_values[i], cutoff_values[i + 1]

     if i == 0:
      __lowercase							  = head_logprob[:, : self.cutoffs[0]]
     else:
      __lowercase  ,	__lowercase  ,	__lowercase							  = weights[i], biases[i], self.out_projs[i]

      __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)
      __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=1			)

      __lowercase							  = head_logprob[:, -i] + tail_logprob_i
      __lowercase							  = logprob_i

    return out

