





import itertools
import random
import unittest

import numpy as np

from transformers import is_speech_available
from transformers.testing_utils import require_torch, require_torchaudio

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
				from transformers import SpeechaTextFeatureExtractor

UpperCamelCase__   : Dict					  = random.Random()




def    A_(			A				,      A=1.0				,      A=None				,      A=None   ):
	if rng is None:
		UpperCAmelCase_								=     global_rng

	UpperCAmelCase_								=     []
	for batch_idx in range(shape[0]   ):
		values.append([]   )
		for _ in range(shape[1]   ):
			values[-1].append(rng.random() * scale   )

	return values







@require_torch
@require_torchaudio
class   _UpperCamelCase  (       unittest.TestCase  ):







					'''simple docstring'''

					def __init__(						self    :  Tuple ,					__lowercase    :  List[Any] ,					__lowercase    :  List[str]=7 ,					__lowercase    :  Optional[Any]=4_00 ,					__lowercase    :  Any=20_00 ,					__lowercase    :  List[Any]=24 ,					__lowercase    :  Any=24 ,					__lowercase    :  Any=0.0 ,					__lowercase    :  Dict=1_60_00 ,					__lowercase    :  int=True ,					__lowercase    :  Any=True ,					):






						'''simple docstring'''







						UpperCAmelCase_								=     parent
						UpperCAmelCase_								=     batch_size
						UpperCAmelCase_								=     min_seq_length
						UpperCAmelCase_								=     max_seq_length
						UpperCAmelCase_								=     (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
						UpperCAmelCase_								=     feature_size
						UpperCAmelCase_								=     num_mel_bins
						UpperCAmelCase_								=     padding_value
						UpperCAmelCase_								=     sampling_rate
						UpperCAmelCase_								=     return_attention_mask
						UpperCAmelCase_								=     do_normalize





					def        SCREAMING_SNAKE_CASE							(						self    :  Union[str, Any]		):






						'''simple docstring'''







						return {
						    "feature_size": self.feature_size,
						    "num_mel_bins": self.num_mel_bins,
						    "padding_value": self.padding_value,
						    "sampling_rate": self.sampling_rate,
						    "return_attention_mask": self.return_attention_mask,
						    "do_normalize": self.do_normalize,
						}





					def        SCREAMING_SNAKE_CASE							(						self    :  List[str] ,					__lowercase    :  Dict=False ,					__lowercase    :  Tuple=False		):






						'''simple docstring'''







						def _flatten(__lowercase    :  Any		):
							return list(itertools.chain(*__lowercase		)		)

						if equal_length:
							UpperCAmelCase_								=     [floats_list((self.max_seq_length, self.feature_size)		) for _ in range(self.batch_size		)]
						else:
							# make sure that inputs increase in size
							UpperCAmelCase_								=     [
							    floats_list((x, self.feature_size)		)
							    for x in range(self.min_seq_length ,					self.max_seq_length ,					self.seq_length_diff		)
							]
						if numpify:
							UpperCAmelCase_								=     [np.asarray(__lowercase		) for x in speech_inputs]
						return speech_inputs







@require_torch
@require_torchaudio
class   _UpperCamelCase  (       A_						,       unittest.TestCase  ):







					'''simple docstring'''

					lowerCamelCase     :     Dict			    =       SpeechaTextFeatureExtractor if is_speech_available() else None





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     SpeechaTextFeatureExtractionTester(self		)





					def        SCREAMING_SNAKE_CASE							(						self    :  str ,					__lowercase    :  Tuple		):






						'''simple docstring'''







						self.assertTrue(np.all(np.mean(__lowercase ,					axis=0		) < 1e-3		)		)
						self.assertTrue(np.all(np.abs(np.var(__lowercase ,					axis=0		) - 1		) < 1e-3		)		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Any		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
						# create three inputs of length 800, 1000, and 1200
						UpperCAmelCase_								=     [floats_list((1, x)		)[0] for x in range(8_00 ,					14_00 ,					2_00		)]
						UpperCAmelCase_								=     [np.asarray(__lowercase		) for speech_input in speech_inputs]

						# Test feature size
						UpperCAmelCase_								=     feature_extractor(__lowercase ,					padding=__lowercase ,					return_tensors="""np"""		).input_features
						self.assertTrue(input_features.ndim == 3		)
						self.assertTrue(input_features.shape[-1] == feature_extractor.feature_size		)

						# Test not batched input
						UpperCAmelCase_								=     feature_extractor(speech_inputs[0] ,					return_tensors="""np"""		).input_features
						UpperCAmelCase_								=     feature_extractor(np_speech_inputs[0] ,					return_tensors="""np"""		).input_features
						self.assertTrue(np.allclose(__lowercase ,					__lowercase ,					atol=1e-3		)		)

						# Test batched
						UpperCAmelCase_								=     feature_extractor(__lowercase ,					return_tensors="""np"""		).input_features
						UpperCAmelCase_								=     feature_extractor(__lowercase ,					return_tensors="""np"""		).input_features
						for enc_seq_a, enc_seq_a in zip(__lowercase ,					__lowercase		):
							self.assertTrue(np.allclose(__lowercase ,					__lowercase ,					atol=1e-3		)		)

						# Test 2-D numpy arrays are batched.
						UpperCAmelCase_								=     [floats_list((1, x)		)[0] for x in (8_00, 8_00, 8_00)]
						UpperCAmelCase_								=     np.asarray(__lowercase		)
						UpperCAmelCase_								=     feature_extractor(__lowercase ,					return_tensors="""np"""		).input_features
						UpperCAmelCase_								=     feature_extractor(__lowercase ,					return_tensors="""np"""		).input_features
						for enc_seq_a, enc_seq_a in zip(__lowercase ,					__lowercase		):
							self.assertTrue(np.allclose(__lowercase ,					__lowercase ,					atol=1e-3		)		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Union[str, Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
						UpperCAmelCase_								=     [floats_list((1, x)		)[0] for x in range(8_00 ,					14_00 ,					2_00		)]

						UpperCAmelCase_								=     ["""longest""", """max_length""", """do_not_pad"""]
						UpperCAmelCase_								=     [None, 16, None]
						for max_length, padding in zip(__lowercase ,					__lowercase		):
							UpperCAmelCase_								=     feature_extractor(
							    __lowercase ,					padding=__lowercase ,					max_length=__lowercase ,					return_attention_mask=__lowercase		)
							UpperCAmelCase_								=     inputs.input_features
							UpperCAmelCase_								=     inputs.attention_mask
							UpperCAmelCase_								=     [np.sum(__lowercase		) for x in attention_mask]

							self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]		)
							self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]		)
							self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]		)





					def        SCREAMING_SNAKE_CASE							(						self    :  List[str]		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
						UpperCAmelCase_								=     [floats_list((1, x)		)[0] for x in range(8_00 ,					14_00 ,					2_00		)]

						UpperCAmelCase_								=     ["""longest""", """max_length""", """do_not_pad"""]
						UpperCAmelCase_								=     [None, 16, None]
						for max_length, padding in zip(__lowercase ,					__lowercase		):
							UpperCAmelCase_								=     feature_extractor(
							    __lowercase ,					max_length=__lowercase ,					padding=__lowercase ,					return_tensors="""np""" ,					return_attention_mask=__lowercase		)
							UpperCAmelCase_								=     inputs.input_features
							UpperCAmelCase_								=     inputs.attention_mask
							UpperCAmelCase_								=     [np.sum(__lowercase		) for x in attention_mask]

							self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]		)
							self.assertTrue(input_features[0][fbank_feat_lengths[0] :].sum() < 1e-6		)
							self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]		)
							self.assertTrue(input_features[0][fbank_feat_lengths[1] :].sum() < 1e-6		)
							self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Any		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
						UpperCAmelCase_								=     [floats_list((1, x)		)[0] for x in range(8_00 ,					14_00 ,					2_00		)]
						UpperCAmelCase_								=     feature_extractor(
						    __lowercase ,					padding="""max_length""" ,					max_length=4 ,					truncation=__lowercase ,					return_tensors="""np""" ,					return_attention_mask=__lowercase ,					)
						UpperCAmelCase_								=     inputs.input_features
						UpperCAmelCase_								=     inputs.attention_mask
						UpperCAmelCase_								=     np.sum(attention_mask == 1 ,					axis=1		)

						self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]		)
						self._check_zero_mean_unit_variance(input_features[1]		)
						self._check_zero_mean_unit_variance(input_features[2]		)





					def        SCREAMING_SNAKE_CASE							(						self    :  str		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
						UpperCAmelCase_								=     [floats_list((1, x)		)[0] for x in range(8_00 ,					14_00 ,					2_00		)]
						UpperCAmelCase_								=     feature_extractor(
						    __lowercase ,					padding="""longest""" ,					max_length=4 ,					truncation=__lowercase ,					return_tensors="""np""" ,					return_attention_mask=__lowercase ,					)
						UpperCAmelCase_								=     inputs.input_features
						UpperCAmelCase_								=     inputs.attention_mask
						UpperCAmelCase_								=     np.sum(attention_mask == 1 ,					axis=1		)

						self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]		)
						self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]		)
						self._check_zero_mean_unit_variance(input_features[2]		)

						# make sure that if max_length < longest -> then pad to max_length
						self.assertEqual(input_features.shape ,					(3, 4, 24)		)

						UpperCAmelCase_								=     [floats_list((1, x)		)[0] for x in range(8_00 ,					14_00 ,					2_00		)]
						UpperCAmelCase_								=     feature_extractor(
						    __lowercase ,					padding="""longest""" ,					max_length=16 ,					truncation=__lowercase ,					return_tensors="""np""" ,					return_attention_mask=__lowercase ,					)
						UpperCAmelCase_								=     inputs.input_features
						UpperCAmelCase_								=     inputs.attention_mask
						UpperCAmelCase_								=     np.sum(attention_mask == 1 ,					axis=1		)

						self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]		)
						self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]		)
						self._check_zero_mean_unit_variance(input_features[2]		)

						# make sure that if max_length < longest -> then pad to max_length
						self.assertEqual(input_features.shape ,					(3, 6, 24)		)





					def        SCREAMING_SNAKE_CASE							(						self    :  str		):






						'''simple docstring'''







						import torch

						UpperCAmelCase_								=     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
						UpperCAmelCase_								=     np.random.rand(1_00 ,					32		).astype(np.floataa		)
						UpperCAmelCase_								=     np_speech_inputs.tolist()

						for inputs in [py_speech_inputs, np_speech_inputs]:
							UpperCAmelCase_								=     feature_extractor.pad([{"""input_features""": inputs}] ,					return_tensors="""np"""		)
							self.assertTrue(np_processed.input_features.dtype == np.floataa		)
							UpperCAmelCase_								=     feature_extractor.pad([{"""input_features""": inputs}] ,					return_tensors="""pt"""		)
							self.assertTrue(pt_processed.input_features.dtype == torch.floataa		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[int] ,					__lowercase    :  Dict		):






						'''simple docstring'''







						from datasets import load_dataset

						UpperCAmelCase_								=     load_dataset("""hf-internal-testing/librispeech_asr_dummy""" ,					"""clean""" ,					split="""validation"""		)
						# automatic decoding with librispeech
						UpperCAmelCase_								=     ds.sort("""id"""		).select(range(__lowercase		)		)[:num_samples]["""audio"""]

						return [x["array"] for x in speech_samples]





					def        SCREAMING_SNAKE_CASE							(						self    :  List[Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     np.array([
						    -1.5745, -1.7713, -1.7020, -1.6069, -1.2250, -1.1105, -0.9072, -0.8241,
						    -1.2310, -0.8098, -0.3320, -0.4101, -0.7985, -0.4996, -0.8213, -0.9128,
						    -1.0420, -1.1286, -1.0440, -0.7999, -0.8405, -1.2275, -1.5443, -1.4625,
						]		)
						# fmt: on

						UpperCAmelCase_								=     self._load_datasamples(1		)
						UpperCAmelCase_								=     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
						UpperCAmelCase_								=     feature_extractor(__lowercase ,					return_tensors="""pt"""		).input_features
						self.assertEquals(input_features.shape ,					(1, 5_84, 24)		)
						self.assertTrue(np.allclose(input_features[0, 0, :30] ,					__lowercase ,					atol=1e-4		)		)









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




def    A_(			A				,      A				,      A   ):
	UpperCAmelCase_								=     nn.ModuleList([src_layers[i] for i in layers_to_copy]   )
	assert len(A   ) == len(A   ), f"""{len(A   )} != {len(A   )}"""
	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
    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))},
}
UpperCamelCase__   : 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_(			A				,      A   ):
	try:
		UpperCAmelCase_								=     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(A   )   )




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




def    A_(			A				,      A = "student"				,      A = None				,      A = None				,      A=False				,      A=None				,      A=None				,      **A				,      ):
	UpperCAmelCase_								=     """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(A				,      A   ):
		AutoTokenizer.from_pretrained(A   ).save_pretrained(A   )  # purely for convenience
		UpperCAmelCase_								=     AutoModelForSeqaSeqLM.from_pretrained(A   ).eval()
	else:
		assert isinstance(A				,      A   ), f"""teacher must be a model or string got type {type(A   )}"""
	UpperCAmelCase_								=     teacher.config.to_diff_dict()

	try:
		UpperCAmelCase_						,	UpperCAmelCase_								=     teacher.config.encoder_layers, teacher.config.decoder_layers
		if e is None:
			UpperCAmelCase_								=     teacher_e
		if d is None:
			UpperCAmelCase_								=     teacher_d
		init_kwargs.update({"""encoder_layers""": e, """decoder_layers""": d}   )
	except AttributeError:  # T5
		if hasattr(teacher.config				,      """num_encoder_layers"""   ):
			UpperCAmelCase_						,	UpperCAmelCase_								=     teacher.config.num_encoder_layers, teacher.config.num_decoder_layers
		else:
			UpperCAmelCase_						,	UpperCAmelCase_								=     teacher.config.num_layers, teacher.config.num_decoder_layers
		if e is None:
			UpperCAmelCase_								=     teacher_e
		if d is None:
			UpperCAmelCase_								=     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(A   )

	# Copy weights
	UpperCAmelCase_								=     teacher.config_class(**A   )
	UpperCAmelCase_								=     AutoModelForSeqaSeqLM.from_config(A   )
	# Start by copying the full teacher state dict this will copy the first N teacher layers to the student.
	UpperCAmelCase_								=     student.load_state_dict(teacher.state_dict()				,      strict=A   )
	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_								=     list(range(A   )   ), list(range(A   )   )
		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(A   )
		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_								=     pick_layers_to_copy(A				,      A   )
	if d_layers_to_copy is None:
		UpperCAmelCase_								=     pick_layers_to_copy(A				,      A   )

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


