


'''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__					(	UpperCAmelCase_    :				nn.ModuleList      ,    UpperCAmelCase_    :				nn.ModuleList      ,    UpperCAmelCase_    :				List[int]							)     ->		None:
     __lowerCamelCase   :	List[Any]								=		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__		:					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 			UpperCAmelCase__					(	UpperCAmelCase_    :				Optional[Any]      ,    UpperCAmelCase_    :				Tuple							)     ->		Union[str, Any]:
     try:
          __lowerCamelCase   :	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(UpperCAmelCase_							)							)



def 			UpperCAmelCase__					(	UpperCAmelCase_    :				List[Any]      ,    UpperCAmelCase_    :				int							)     ->		List[int]:
     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 			UpperCAmelCase__					(	UpperCAmelCase_    :				Union[str, PreTrainedModel]      ,    UpperCAmelCase_    :				Union[str, Path] = "student"      ,    UpperCAmelCase_    :				Union[int, None] = None      ,    UpperCAmelCase_    :				Union[int, None] = None      ,    UpperCAmelCase_    :				Union[str, Any]=False      ,    UpperCAmelCase_    :				List[Any]=None      ,    UpperCAmelCase_    :				Union[str, Any]=None      ,    **UpperCAmelCase_    :				Optional[Any]      ,    )     ->		Tuple[PreTrainedModel, List[int], List[int]]:
     __lowerCamelCase   :	int								=		'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   :	List[Any]								=		AutoModelForSeqaSeqLM.from_pretrained(UpperCAmelCase_							).eval()
     else:
          assert isinstance(UpperCAmelCase_      ,    UpperCAmelCase_							), F'teacher must be a model or string got type {type(UpperCAmelCase_							)}'
     __lowerCamelCase   :	Any								=		teacher.config.to_diff_dict()

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

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



from . import (
    albert,
    align,
    altclip,
    audio_spectrogram_transformer,
    auto,
    autoformer,
    bark,
    bart,
    barthez,
    bartpho,
    beit,
    bert,
    bert_generation,
    bert_japanese,
    bertweet,
    big_bird,
    bigbird_pegasus,
    biogpt,
    bit,
    blenderbot,
    blenderbot_small,
    blip,
    blip_a,
    bloom,
    bridgetower,
    byta,
    camembert,
    canine,
    chinese_clip,
    clap,
    clip,
    clipseg,
    codegen,
    conditional_detr,
    convbert,
    convnext,
    convnextva,
    cpm,
    cpmant,
    ctrl,
    cvt,
    dataavec,
    deberta,
    deberta_va,
    decision_transformer,
    deformable_detr,
    deit,
    deprecated,
    deta,
    detr,
    dialogpt,
    dinat,
    distilbert,
    dit,
    donut,
    dpr,
    dpt,
    efficientformer,
    efficientnet,
    electra,
    encodec,
    encoder_decoder,
    ernie,
    ernie_m,
    esm,
    falcon,
    flaubert,
    flava,
    fnet,
    focalnet,
    fsmt,
    funnel,
    git,
    glpn,
    gpta,
    gpt_bigcode,
    gpt_neo,
    gpt_neox,
    gpt_neox_japanese,
    gpt_swa,
    gptj,
    gptsan_japanese,
    graphormer,
    groupvit,
    herbert,
    hubert,
    ibert,
    imagegpt,
    informer,
    instructblip,
    jukebox,
    layoutlm,
    layoutlmva,
    layoutlmva,
    layoutxlm,
    led,
    levit,
    lilt,
    llama,
    longformer,
    longta,
    luke,
    lxmert,
    mam_aaa,
    marian,
    markuplm,
    maskaformer,
    maskformer,
    mbart,
    mbartaa,
    mega,
    megatron_bert,
    megatron_gpta,
    mgp_str,
    mluke,
    mobilebert,
    mobilenet_va,
    mobilenet_va,
    mobilevit,
    mobilevitva,
    mpnet,
    mra,
    mta,
    musicgen,
    mvp,
    nat,
    nezha,
    nllb,
    nllb_moe,
    nystromformer,
    oneformer,
    open_llama,
    openai,
    opt,
    owlvit,
    pegasus,
    pegasus_x,
    perceiver,
    phobert,
    pixastruct,
    plbart,
    poolformer,
    prophetnet,
    qdqbert,
    rag,
    realm,
    reformer,
    regnet,
    rembert,
    resnet,
    roberta,
    roberta_prelayernorm,
    roc_bert,
    roformer,
    rwkv,
    sam,
    segformer,
    sew,
    sew_d,
    speech_encoder_decoder,
    speech_to_text,
    speech_to_text_a,
    speechta,
    splinter,
    squeezebert,
    swiftformer,
    swin,
    swinasr,
    swinva,
    switch_transformers,
    ta,
    table_transformer,
    tapas,
    time_series_transformer,
    timesformer,
    timm_backbone,
    transfo_xl,
    trocr,
    tvlt,
    umta,
    unispeech,
    unispeech_sat,
    upernet,
    videomae,
    vilt,
    vision_encoder_decoder,
    vision_text_dual_encoder,
    visual_bert,
    vit,
    vit_hybrid,
    vit_mae,
    vit_msn,
    vivit,
    wavaveca,
    wavaveca_conformer,
    wavaveca_phoneme,
    wavaveca_with_lm,
    wavlm,
    whisper,
    x_clip,
    xglm,
    xlm,
    xlm_prophetnet,
    xlm_roberta,
    xlm_roberta_xl,
    xlnet,
    xmod,
    yolos,
    yoso,
)






