


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




def      _a     (						_lowercase :     nn.ModuleList      , _lowercase :     nn.ModuleList      , _lowercase :     List[int]  ):





     '''simple docstring'''






     __UpperCAmelCase     :			Tuple  							=      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()  )


__UpperCAmelCase      :str			    =       {
    # maps  num layers in teacher -> num_layers in student -> which teacher layers to copy.
    # 12: bart, 16: pegasus, 6: marian/Helsinki-NLP
    1_2: {
        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, 1_1],
        4: [0, 4, 8, 1_1],
        6: [0, 2, 4, 7, 9, 1_1],
        9: [0, 1, 2, 4, 5, 7, 9, 1_0, 1_1],
        1_2: list(range(1_2)),
    },
    1_6: {  # maps  num layers in student -> which teacher layers to copy
        1: [0],
        2: [0, 1_5],
        3: [0, 8, 1_5],
        4: [0, 5, 1_0, 1_5],
        6: [0, 3, 6, 9, 1_2, 1_5],
        8: [0, 2, 4, 6, 8, 1_0, 1_2, 1_5],
        9: [0, 1, 3, 5, 7, 9, 1_1, 1_3, 1_5],
        1_2: [0, 1, 2, 3, 4, 5, 6, 7, 9, 1_1, 1_3, 1_5],
        1_6: list(range(1_6)),
    },
    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]},
    1_2: {1: [1_1], 2: [5, 1_1], 3: [3, 7, 1_1], 6: [1, 3, 5, 8, 1_0, 1_1]},
    1_6: {1: [1_5], 4: [4, 9, 1_2, 1_5], 8: [1, 3, 5, 7, 9, 1_1, 1_3, 1_5]},
}




def      _a     (						_lowercase :     List[Any]      , _lowercase :     str  ):





     '''simple docstring'''






     try:
          __UpperCAmelCase     :			Optional[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(_lowercase  )  )




def      _a     (						_lowercase :     int      , _lowercase :     Tuple  ):





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




def      _a     (						_lowercase :     Union[str, PreTrainedModel]      , _lowercase :     Union[str, Path] = "student"      , _lowercase :     Union[int, None] = None      , _lowercase :     Union[int, None] = None      , _lowercase :     str=False      , _lowercase :     Optional[Any]=None      , _lowercase :     List[str]=None      , **_lowercase :     int      , ):





     '''simple docstring'''






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

     try:
          __UpperCAmelCase							,       __UpperCAmelCase     :			Optional[int]  							=      teacher.config.encoder_layers, teacher.config.decoder_layers
          if e is None:
               __UpperCAmelCase     :			Any  							=      teacher_e
          if d is None:
               __UpperCAmelCase     :			int  							=      teacher_d
          init_kwargs.update({'''encoder_layers''': e, '''decoder_layers''': d}  )
     except AttributeError:  # T5
          if hasattr(teacher.config      , '''num_encoder_layers'''  ):
               __UpperCAmelCase							,       __UpperCAmelCase     :			Optional[Any]  							=      teacher.config.num_encoder_layers, teacher.config.num_decoder_layers
          else:
               __UpperCAmelCase							,       __UpperCAmelCase     :			Dict  							=      teacher.config.num_layers, teacher.config.num_decoder_layers
          if e is None:
               __UpperCAmelCase     :			List[str]  							=      teacher_e
          if d is None:
               __UpperCAmelCase     :			Tuple  							=      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
     __UpperCAmelCase     :			Union[str, Any]  							=      teacher.config_class(**_lowercase  )
     __UpperCAmelCase     :			Optional[int]  							=      AutoModelForSeqaSeqLM.from_config(_lowercase  )
     # Start by copying the full teacher state dict this will copy the first N teacher layers to the student.
     __UpperCAmelCase     :			Union[str, Any]  							=      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
          __UpperCAmelCase							,       __UpperCAmelCase     :			Union[str, Any]  							=      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:
          __UpperCAmelCase     :			List[int]  							=      pick_layers_to_copy(_lowercase      , _lowercase  )
     if d_layers_to_copy is None:
          __UpperCAmelCase     :			List[int]  							=      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}'  )
     __UpperCAmelCase     :			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 os
import sys
import unittest


__UpperCAmelCase      :Optional[int]			    =       os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, "utils"))

import check_dummies  # noqa: E402
from check_dummies import create_dummy_files, create_dummy_object, find_backend, read_init  # noqa: E402


# Align TRANSFORMERS_PATH in check_dummies with the current path
__UpperCAmelCase      :Dict			    =       os.path.join(git_repo_path, "src", "diffusers")



class      a    (		unittest.TestCase	):


     """simple docstring"""






     def 			lowerCamelCase__			(    self					:		Any      )						->							int:
          __UpperCAmelCase     :			Optional[Any]  							=      find_backend('''    if not is_torch_available():'''      )
          self.assertEqual(snake_case							,			'''torch'''      )

          # backend_with_underscore = find_backend("    if not is_tensorflow_text_available():")
          # self.assertEqual(backend_with_underscore, "tensorflow_text")

          __UpperCAmelCase     :			Union[str, Any]  							=      find_backend('''    if not (is_torch_available() and is_transformers_available()):'''      )
          self.assertEqual(snake_case							,			'''torch_and_transformers'''      )

          # double_backend_with_underscore = find_backend(
          #    "    if not (is_sentencepiece_available() and is_tensorflow_text_available()):"
          # )
          # self.assertEqual(double_backend_with_underscore, "sentencepiece_and_tensorflow_text")

          __UpperCAmelCase     :			List[str]  							=      find_backend(
              '''    if not (is_torch_available() and is_transformers_available() and is_onnx_available()):'''      )
          self.assertEqual(snake_case							,			'''torch_and_transformers_and_onnx'''      )

     def 			lowerCamelCase__			(    self					:		Optional[int]      )						->							int:
          __UpperCAmelCase     :			Tuple  							=      read_init()
          # We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects
          self.assertIn('''torch'''							,			snake_case      )
          self.assertIn('''torch_and_transformers'''							,			snake_case      )
          self.assertIn('''flax_and_transformers'''							,			snake_case      )
          self.assertIn('''torch_and_transformers_and_onnx'''							,			snake_case      )

          # Likewise, we can't assert on the exact content of a key
          self.assertIn('''UNet2DModel'''							,			objects['''torch''']      )
          self.assertIn('''FlaxUNet2DConditionModel'''							,			objects['''flax''']      )
          self.assertIn('''StableDiffusionPipeline'''							,			objects['''torch_and_transformers''']      )
          self.assertIn('''FlaxStableDiffusionPipeline'''							,			objects['''flax_and_transformers''']      )
          self.assertIn('''LMSDiscreteScheduler'''							,			objects['''torch_and_scipy''']      )
          self.assertIn('''OnnxStableDiffusionPipeline'''							,			objects['''torch_and_transformers_and_onnx''']      )

     def 			lowerCamelCase__			(    self					:		Optional[int]      )						->							List[Any]:
          __UpperCAmelCase     :			str  							=      create_dummy_object('''CONSTANT'''							,			'''\'torch\''''      )
          self.assertEqual(snake_case							,			'''\nCONSTANT = None\n'''      )

          __UpperCAmelCase     :			Union[str, Any]  							=      create_dummy_object('''function'''							,			'''\'torch\''''      )
          self.assertEqual(
              snake_case							,			'''\ndef function(*args, **kwargs):\n    requires_backends(function, \'torch\')\n'''      )

          __UpperCAmelCase     :			Optional[int]  							=      '''
class FakeClass(metaclass=DummyObject):
    _backends = \'torch\'

    def __init__(self, *args, **kwargs):
        requires_backends(self, \'torch\')

    @classmethod
    def from_config(cls, *args, **kwargs):
        requires_backends(cls, \'torch\')

    @classmethod
    def from_pretrained(cls, *args, **kwargs):
        requires_backends(cls, \'torch\')
'''
          __UpperCAmelCase     :			Optional[Any]  							=      create_dummy_object('''FakeClass'''							,			'''\'torch\''''      )
          self.assertEqual(snake_case							,			snake_case      )






     def 			lowerCamelCase__			(    self					:		int      )						->							List[Any]:
          __UpperCAmelCase     :			List[str]  							=      '''# This file is autogenerated by the command `make fix-copies`, do not edit.
from ..utils import DummyObject, requires_backends


CONSTANT = None


def function(*args, **kwargs):
    requires_backends(function, ["torch"])


class FakeClass(metaclass=DummyObject):
    _backends = ["torch"]

    def __init__(self, *args, **kwargs):
        requires_backends(self, ["torch"])

    @classmethod
    def from_config(cls, *args, **kwargs):
        requires_backends(cls, ["torch"])

    @classmethod
    def from_pretrained(cls, *args, **kwargs):
        requires_backends(cls, ["torch"])
'''
          __UpperCAmelCase     :			Optional[int]  							=      create_dummy_files({'''torch''': ['''CONSTANT''', '''function''', '''FakeClass''']}      )
          self.assertEqual(dummy_files['''torch''']							,			snake_case      )