






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


lowercase  :							List[str]			  =    logging.get_logger(__name__)

def      lowerCAmelCase__	(    _a       :			nn.ModuleList  ,      _a       :			nn.ModuleList  ,      _a       :			List[int]    ):
  snake_case_				:							str						 =       nn.ModuleList([src_layers[i] for i in layers_to_copy]    )
  assert len(snake_case_    ) == len(snake_case_    ), F'''{len(snake_case_    )} != {len(snake_case_    )}'''
  dest_layers.load_state_dict(layers_to_copy.state_dict()    )


lowercase  :							Dict			  =    {
    # 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))},
}
lowercase  :							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      lowerCAmelCase__	(    _a       :			Any  ,      _a       :			Optional[Any]    ):
  try:
    snake_case_				:							Optional[Any]						 =       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(snake_case_    )    )

def      lowerCAmelCase__	(    _a       :			List[str]  ,      _a       :			str    ):
  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(snake_case_    )    )
  elif n_student == 1:
    return [n_teacher - 1]
  else:
    return LAYERS_TO_SUPERVISE[n_teacher][n_student]

def      lowerCAmelCase__	(    _a       :			Union[str, PreTrainedModel]  ,      _a       :			Union[str, Path] = "student"  ,      _a       :			Union[int, None] = None  ,      _a       :			Union[int, None] = None  ,      _a       :			Optional[Any]=False  ,      _a       :			Union[str, Any]=None  ,      _a       :			Optional[Any]=None  ,      **_a       :			Tuple  ,      ):
  snake_case_				:							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(snake_case_  ,      snake_case_    ):
    AutoTokenizer.from_pretrained(snake_case_    ).save_pretrained(snake_case_    )  # purely for convenience
    snake_case_				:							str						 =       AutoModelForSeqaSeqLM.from_pretrained(snake_case_    ).eval()
  else:
    assert isinstance(snake_case_  ,      snake_case_    ), F'''teacher must be a model or string got type {type(snake_case_    )}'''
  snake_case_				:							Dict						 =       teacher.config.to_diff_dict()

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

  # Copy weights
  snake_case_				:							Any						 =       teacher.config_class(**snake_case_    )
  snake_case_				:							Optional[int]						 =       AutoModelForSeqaSeqLM.from_config(snake_case_    )
  # 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=snake_case_    )
  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_				:							List[Any]						 =       list(range(snake_case_    )    ), list(range(snake_case_    )    )
    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(snake_case_    )
    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(snake_case_  ,      snake_case_    )
  if d_layers_to_copy is None:
    snake_case_				:							List[int]						 =       pick_layers_to_copy(snake_case_  ,      snake_case_    )

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








from typing import List, Union

import numpy as np

from ..tokenization_utils import TruncationStrategy
from ..utils import add_end_docstrings, logging
from .base import PIPELINE_INIT_ARGS, ArgumentHandler, ChunkPipeline


lowercase  :							Dict			  =    logging.get_logger(__name__)




class 		UpperCAmelCase_						(	SCREAMING_SNAKE_CASE__					):


						'''simple docstring'''




						def       _lowerCAmelCase       (	self		,				_SCREAMING_SNAKE_CASE				) ->				List[Any]:
								if isinstance(_SCREAMING_SNAKE_CASE		,				_SCREAMING_SNAKE_CASE				):
										snake_case_				:							List[str]						 =       [label.strip() for label in labels.split(","				) if label.strip()]
								return labels






						def __call__(	self		,				_SCREAMING_SNAKE_CASE		,				_SCREAMING_SNAKE_CASE		,				_SCREAMING_SNAKE_CASE				) ->				Optional[Any]:
								if len(_SCREAMING_SNAKE_CASE				) == 0 or len(_SCREAMING_SNAKE_CASE				) == 0:
										raise ValueError("You must include at least one label and at least one sequence."				)
								if hypothesis_template.format(labels[0]				) == hypothesis_template:
										raise ValueError(
										    (
										        "The provided hypothesis_template \"{}\" was not able to be formatted with the target labels. "
										        "Make sure the passed template includes formatting syntax such as {{}} where the label should go."
										    ).format(_SCREAMING_SNAKE_CASE				)				)

								if isinstance(_SCREAMING_SNAKE_CASE		,				_SCREAMING_SNAKE_CASE				):
										snake_case_				:							List[str]						 =       [sequences]

								snake_case_				:							List[str]						 =       []
								for sequence in sequences:
										sequence_pairs.extend([[sequence, hypothesis_template.format(_SCREAMING_SNAKE_CASE				)] for label in labels]				)

								return sequence_pairs, sequences




@add_end_docstrings(SCREAMING_SNAKE_CASE__					)
class 		UpperCAmelCase_						(	SCREAMING_SNAKE_CASE__					):


						'''simple docstring'''




						def __init__(	self		,				_SCREAMING_SNAKE_CASE=ZeroShotClassificationArgumentHandler()		,				*_SCREAMING_SNAKE_CASE		,				**_SCREAMING_SNAKE_CASE				) ->				Optional[Any]:
								snake_case_				:							Dict						 =       args_parser
								super().__init__(*_SCREAMING_SNAKE_CASE		,				**_SCREAMING_SNAKE_CASE				)
								if self.entailment_id == -1:
										logger.warning(
										    "Failed to determine 'entailment' label id from the label2id mapping in the model config. Setting to "
										    "-1. Define a descriptive label2id mapping in the model config to ensure correct outputs."				)






						@property
						def       _lowerCAmelCase       (	self				) ->				str:
								for label, ind in self.model.config.labelaid.items():
										if label.lower().startswith("entail"				):
												return ind
								return -1






						def       _lowerCAmelCase       (	self		,				_SCREAMING_SNAKE_CASE		,				_SCREAMING_SNAKE_CASE=True		,				_SCREAMING_SNAKE_CASE=True		,				_SCREAMING_SNAKE_CASE=TruncationStrategy.ONLY_FIRST		,				**_SCREAMING_SNAKE_CASE				) ->				Dict:
								snake_case_				:							List[Any]						 =       self.framework
								if self.tokenizer.pad_token is None:
										# Override for tokenizers not supporting padding
										logger.error(
										    "Tokenizer was not supporting padding necessary for zero-shot, attempting to use "
										    " `pad_token=eos_token`"				)
										snake_case_				:							List[str]						 =       self.tokenizer.eos_token
								try:
										snake_case_				:							Union[str, Any]						 =       self.tokenizer(
										    _SCREAMING_SNAKE_CASE		,				add_special_tokens=_SCREAMING_SNAKE_CASE		,				return_tensors=_SCREAMING_SNAKE_CASE		,				padding=_SCREAMING_SNAKE_CASE		,				truncation=_SCREAMING_SNAKE_CASE		,				)
								except Exception as e:
										if "too short" in str(_SCREAMING_SNAKE_CASE				):
												# tokenizers might yell that we want to truncate
												# to a value that is not even reached by the input.
												# In that case we don't want to truncate.
												# It seems there's not a really better way to catch that
												# exception.

												snake_case_				:							Union[str, Any]						 =       self.tokenizer(
												    _SCREAMING_SNAKE_CASE		,				add_special_tokens=_SCREAMING_SNAKE_CASE		,				return_tensors=_SCREAMING_SNAKE_CASE		,				padding=_SCREAMING_SNAKE_CASE		,				truncation=TruncationStrategy.DO_NOT_TRUNCATE		,				)
										else:
												raise e

								return inputs






						def       _lowerCAmelCase       (	self		,				**_SCREAMING_SNAKE_CASE				) ->				Optional[int]:
								if kwargs.get("multi_class"		,				_SCREAMING_SNAKE_CASE				) is not None:
										snake_case_				:							Any						 =       kwargs["multi_class"]
										logger.warning(
										    "The `multi_class` argument has been deprecated and renamed to `multi_label`. "
										    "`multi_class` will be removed in a future version of Transformers."				)
								snake_case_				:							Any						 =       {}
								if "candidate_labels" in kwargs:
										snake_case_				:							Tuple						 =       self._args_parser._parse_labels(kwargs["candidate_labels"]				)
								if "hypothesis_template" in kwargs:
										snake_case_				:							Optional[Any]						 =       kwargs["hypothesis_template"]

								snake_case_				:							Dict						 =       {}
								if "multi_label" in kwargs:
										snake_case_				:							List[Any]						 =       kwargs["multi_label"]
								return preprocess_params, {}, postprocess_params






						def __call__(	self		,				_SCREAMING_SNAKE_CASE		,				*_SCREAMING_SNAKE_CASE		,				**_SCREAMING_SNAKE_CASE		,				) ->				str:
								if len(_SCREAMING_SNAKE_CASE				) == 0:
										pass
								elif len(_SCREAMING_SNAKE_CASE				) == 1 and "candidate_labels" not in kwargs:
										snake_case_				:							int						 =       args[0]
								else:
										raise ValueError(f'''Unable to understand extra arguments {args}'''				)

								return super().__call__(_SCREAMING_SNAKE_CASE		,				**_SCREAMING_SNAKE_CASE				)






						def       _lowerCAmelCase       (	self		,				_SCREAMING_SNAKE_CASE		,				_SCREAMING_SNAKE_CASE=None		,				_SCREAMING_SNAKE_CASE="This example is {}."				) ->				str:
								snake_case_ ,       snake_case_				:							Optional[int]						 =       self._args_parser(_SCREAMING_SNAKE_CASE		,				_SCREAMING_SNAKE_CASE		,				_SCREAMING_SNAKE_CASE				)

								for i, (candidate_label, sequence_pair) in enumerate(zip(_SCREAMING_SNAKE_CASE		,				_SCREAMING_SNAKE_CASE				)				):
										snake_case_				:							str						 =       self._parse_and_tokenize([sequence_pair]				)

										yield {
										    "candidate_label": candidate_label,
										    "sequence": sequences[0],
										    "is_last": i == len(_SCREAMING_SNAKE_CASE				) - 1,
										    **model_input,
										}






						def       _lowerCAmelCase       (	self		,				_SCREAMING_SNAKE_CASE				) ->				Any:
								snake_case_				:							Optional[Any]						 =       inputs["candidate_label"]
								snake_case_				:							Dict						 =       inputs["sequence"]
								snake_case_				:							Optional[int]						 =       {k: inputs[k] for k in self.tokenizer.model_input_names}
								snake_case_				:							Dict						 =       self.model(**_SCREAMING_SNAKE_CASE				)

								snake_case_				:							Optional[int]						 =       {
								    "candidate_label": candidate_label,
								    "sequence": sequence,
								    "is_last": inputs["is_last"],
								    **outputs,
								}
								return model_outputs






						def       _lowerCAmelCase       (	self		,				_SCREAMING_SNAKE_CASE		,				_SCREAMING_SNAKE_CASE=False				) ->				Any:
								snake_case_				:							str						 =       [outputs["candidate_label"] for outputs in model_outputs]
								snake_case_				:							Union[str, Any]						 =       [outputs["sequence"] for outputs in model_outputs]
								snake_case_				:							List[Any]						 =       np.concatenate([output["logits"].numpy() for output in model_outputs]				)
								snake_case_				:							Tuple						 =       logits.shape[0]
								snake_case_				:							List[Any]						 =       len(_SCREAMING_SNAKE_CASE				)
								snake_case_				:							int						 =       N // n
								snake_case_				:							Optional[Any]						 =       logits.reshape((num_sequences, n, -1)				)

								if multi_label or len(_SCREAMING_SNAKE_CASE				) == 1:
										# softmax over the entailment vs. contradiction dim for each label independently
										snake_case_				:							Any						 =       self.entailment_id
										snake_case_				:							List[str]						 =       -1 if entailment_id == 0 else 0
										snake_case_				:							Dict						 =       reshaped_outputs[..., [contradiction_id, entailment_id]]
										snake_case_				:							List[str]						 =       np.exp(_SCREAMING_SNAKE_CASE				) / np.exp(_SCREAMING_SNAKE_CASE				).sum(-1		,				keepdims=_SCREAMING_SNAKE_CASE				)
										snake_case_				:							Union[str, Any]						 =       scores[..., 1]
								else:
										# softmax the "entailment" logits over all candidate labels
										snake_case_				:							List[str]						 =       reshaped_outputs[..., self.entailment_id]
										snake_case_				:							Optional[Any]						 =       np.exp(_SCREAMING_SNAKE_CASE				) / np.exp(_SCREAMING_SNAKE_CASE				).sum(-1		,				keepdims=_SCREAMING_SNAKE_CASE				)

								snake_case_				:							int						 =       list(reversed(scores[0].argsort()				)				)
								return {
								    "sequence": sequences[0],
								    "labels": [candidate_labels[i] for i in top_inds],
								    "scores": scores[0, top_inds].tolist(),
								}
