
"""simple docstring"""




# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch

from ..models.speechta import SpeechTaForTextToSpeech, SpeechTaHifiGan, SpeechTaProcessor
from ..utils import is_datasets_available
from .base import PipelineTool


if is_datasets_available():
    from datasets import load_dataset







class   _lowercase						(       lowerCAmelCase_				):


     '''simple docstring'''





     _A								=							'microsoft/speecht5_tts'
     _A								=							(
         'This is a tool that reads an English text out loud. It takes an input named `text` which should contain the '
         'text to read (in English) and returns a waveform object containing the sound.'
     )
     _A								=							'text_reader'
     _A								=							SpeechTaProcessor
     _A								=							SpeechTaForTextToSpeech
     _A								=							SpeechTaHifiGan

     _A								=							['text']
     _A								=							['audio']



     def  lowerCAmelCase__       (      self  )->					Tuple:
          if self.post_processor is None:
               UpperCAmelCase__							: Optional[int] 							=						"microsoft/speecht5_hifigan"
          super().setup()



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase=None  )->					Optional[Any]:
          UpperCAmelCase__							: Tuple 							=						self.pre_processor(text=__UpperCamelCase ,			return_tensors="pt" ,			truncation=__UpperCamelCase  )

          if speaker_embeddings is None:
               if not is_datasets_available():
                    raise ImportError("Datasets needs to be installed if not passing speaker embeddings."  )

               UpperCAmelCase__							: Optional[int] 							=						load_dataset("Matthijs/cmu-arctic-xvectors" ,			split="validation"  )
               UpperCAmelCase__							: Optional[Any] 							=						torch.tensor(embeddings_dataset[73_05]["xvector"]  ).unsqueeze(0  )

          return {"input_ids": inputs["input_ids"], "speaker_embeddings": speaker_embeddings}



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase  )->					List[Any]:
          with torch.no_grad():
               return self.model.generate_speech(**__UpperCamelCase  )



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase  )->					Dict:
          with torch.no_grad():
               return self.post_processor(__UpperCamelCase  ).cpu().detach()





"""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__      :	Tuple          =    logging.get_logger(__name__)



def 							a__ (				lowerCAmelCase	:       nn.ModuleList       ,    lowerCAmelCase	:       nn.ModuleList       ,    lowerCAmelCase	:       List[int]		):






     '''simple docstring'''
     UpperCAmelCase__							: 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__      :	List[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[Any]          =    {
    # 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 							a__ (				lowerCAmelCase	:       Dict       ,    lowerCAmelCase	:       Dict		):






     '''simple docstring'''
     try:
          UpperCAmelCase__							: 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(lowerCAmelCase		)		)



def 							a__ (				lowerCAmelCase	:       int       ,    lowerCAmelCase	:       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(lowerCAmelCase		)		)
     elif n_student == 1:
          return [n_teacher - 1]
     else:
          return LAYERS_TO_SUPERVISE[n_teacher][n_student]



def 							a__ (				lowerCAmelCase	:       Union[str, PreTrainedModel]       ,    lowerCAmelCase	:       Union[str, Path] = "student"       ,    lowerCAmelCase	:       Union[int, None] = None       ,    lowerCAmelCase	:       Union[int, None] = None       ,    lowerCAmelCase	:       List[str]=False       ,    lowerCAmelCase	:       List[str]=None       ,    lowerCAmelCase	:       List[str]=None       ,    **lowerCAmelCase	:       List[str]       ,    ):






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

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



