
"""simple docstring"""
import os
import tempfile
import unittest

from transformers import DistilBertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
			import torch

			from transformers import (
			    DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
			    DistilBertForMaskedLM,
			    DistilBertForMultipleChoice,
			    DistilBertForQuestionAnswering,
			    DistilBertForSequenceClassification,
			    DistilBertForTokenClassification,
			    DistilBertModel,
			)
class 	lowerCAmelCase_						(							UpperCAmelCase__      ):



							'''simple docstring'''





							def __init__(      self  ,							snake_case_  ,							snake_case_=13  ,							snake_case_=7  ,							snake_case_=True  ,							snake_case_=True  ,							snake_case_=False  ,							snake_case_=True  ,							snake_case_=99  ,							snake_case_=32  ,							snake_case_=5  ,							snake_case_=4  ,							snake_case_=37  ,							snake_case_="gelu"  ,							snake_case_=0.1  ,							snake_case_=0.1  ,							snake_case_=512  ,							snake_case_=16  ,							snake_case_=2  ,							snake_case_=0.02  ,							snake_case_=3  ,							snake_case_=4  ,							snake_case_=None  ,							)    ->   Dict:
													__lowerCAmelCase   					=   parent
													__lowerCAmelCase   					=   batch_size
													__lowerCAmelCase   					=   seq_length
													__lowerCAmelCase   					=   is_training
													__lowerCAmelCase   					=   use_input_mask
													__lowerCAmelCase   					=   use_token_type_ids
													__lowerCAmelCase   					=   use_labels
													__lowerCAmelCase   					=   vocab_size
													__lowerCAmelCase   					=   hidden_size
													__lowerCAmelCase   					=   num_hidden_layers
													__lowerCAmelCase   					=   num_attention_heads
													__lowerCAmelCase   					=   intermediate_size
													__lowerCAmelCase   					=   hidden_act
													__lowerCAmelCase   					=   hidden_dropout_prob
													__lowerCAmelCase   					=   attention_probs_dropout_prob
													__lowerCAmelCase   					=   max_position_embeddings
													__lowerCAmelCase   					=   type_vocab_size
													__lowerCAmelCase   					=   type_sequence_label_size
													__lowerCAmelCase   					=   initializer_range
													__lowerCAmelCase   					=   num_labels
													__lowerCAmelCase   					=   num_choices
													__lowerCAmelCase   					=   scope




							def 					A__		(      self							)    ->   Optional[int]:
													__lowerCAmelCase   					=   ids_tensor([self.batch_size, self.seq_length]  ,							self.vocab_size							)

													__lowerCAmelCase   					=   None
													if self.use_input_mask:
																			__lowerCAmelCase   					=   random_attention_mask([self.batch_size, self.seq_length]							)

													__lowerCAmelCase   					=   None
													__lowerCAmelCase   					=   None
													__lowerCAmelCase   					=   None
													if self.use_labels:
																			__lowerCAmelCase   					=   ids_tensor([self.batch_size]  ,							self.type_sequence_label_size							)
																			__lowerCAmelCase   					=   ids_tensor([self.batch_size, self.seq_length]  ,							self.num_labels							)
																			__lowerCAmelCase   					=   ids_tensor([self.batch_size]  ,							self.num_choices							)

													__lowerCAmelCase   					=   self.get_config()

													return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels




							def 					A__		(      self							)    ->   str:
													return DistilBertConfig(
													    vocab_size=self.vocab_size  ,							dim=self.hidden_size  ,							n_layers=self.num_hidden_layers  ,							n_heads=self.num_attention_heads  ,							hidden_dim=self.intermediate_size  ,							hidden_act=self.hidden_act  ,							dropout=self.hidden_dropout_prob  ,							attention_dropout=self.attention_probs_dropout_prob  ,							max_position_embeddings=self.max_position_embeddings  ,							initializer_range=self.initializer_range  ,							)




							def 					A__		(      self  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_							)    ->   Dict:
													__lowerCAmelCase   					=   DistilBertModel(config=UpperCAmelCase__							)
													model.to(UpperCAmelCase__							)
													model.eval()
													__lowerCAmelCase   					=   model(UpperCAmelCase__  ,							UpperCAmelCase__							)
													__lowerCAmelCase   					=   model(UpperCAmelCase__							)
													self.parent.assertEqual(result.last_hidden_state.shape  ,							(self.batch_size, self.seq_length, self.hidden_size)							)




							def 					A__		(      self  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_							)    ->   int:
													__lowerCAmelCase   					=   DistilBertForMaskedLM(config=UpperCAmelCase__							)
													model.to(UpperCAmelCase__							)
													model.eval()
													__lowerCAmelCase   					=   model(UpperCAmelCase__  ,							attention_mask=UpperCAmelCase__  ,							labels=UpperCAmelCase__							)
													self.parent.assertEqual(result.logits.shape  ,							(self.batch_size, self.seq_length, self.vocab_size)							)




							def 					A__		(      self  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_							)    ->   List[str]:
													__lowerCAmelCase   					=   DistilBertForQuestionAnswering(config=UpperCAmelCase__							)
													model.to(UpperCAmelCase__							)
													model.eval()
													__lowerCAmelCase   					=   model(
													    UpperCAmelCase__  ,							attention_mask=UpperCAmelCase__  ,							start_positions=UpperCAmelCase__  ,							end_positions=UpperCAmelCase__							)
													self.parent.assertEqual(result.start_logits.shape  ,							(self.batch_size, self.seq_length)							)
													self.parent.assertEqual(result.end_logits.shape  ,							(self.batch_size, self.seq_length)							)




							def 					A__		(      self  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_							)    ->   Union[str, Any]:
													__lowerCAmelCase   					=   self.num_labels
													__lowerCAmelCase   					=   DistilBertForSequenceClassification(UpperCAmelCase__							)
													model.to(UpperCAmelCase__							)
													model.eval()
													__lowerCAmelCase   					=   model(UpperCAmelCase__  ,							attention_mask=UpperCAmelCase__  ,							labels=UpperCAmelCase__							)
													self.parent.assertEqual(result.logits.shape  ,							(self.batch_size, self.num_labels)							)




							def 					A__		(      self  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_							)    ->   Optional[Any]:
													__lowerCAmelCase   					=   self.num_labels
													__lowerCAmelCase   					=   DistilBertForTokenClassification(config=UpperCAmelCase__							)
													model.to(UpperCAmelCase__							)
													model.eval()

													__lowerCAmelCase   					=   model(UpperCAmelCase__  ,							attention_mask=UpperCAmelCase__  ,							labels=UpperCAmelCase__							)
													self.parent.assertEqual(result.logits.shape  ,							(self.batch_size, self.seq_length, self.num_labels)							)




							def 					A__		(      self  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_  ,							snake_case_							)    ->   Dict:
													__lowerCAmelCase   					=   self.num_choices
													__lowerCAmelCase   					=   DistilBertForMultipleChoice(config=UpperCAmelCase__							)
													model.to(UpperCAmelCase__							)
													model.eval()
													__lowerCAmelCase   					=   input_ids.unsqueeze(1							).expand(-1  ,							self.num_choices  ,							-1							).contiguous()
													__lowerCAmelCase   					=   input_mask.unsqueeze(1							).expand(-1  ,							self.num_choices  ,							-1							).contiguous()
													__lowerCAmelCase   					=   model(
													    UpperCAmelCase__  ,							attention_mask=UpperCAmelCase__  ,							labels=UpperCAmelCase__  ,							)
													self.parent.assertEqual(result.logits.shape  ,							(self.batch_size, self.num_choices)							)




							def 					A__		(      self							)    ->   List[Any]:
													__lowerCAmelCase   					=   self.prepare_config_and_inputs()
													((__lowerCAmelCase)					,			(__lowerCAmelCase)					,			(__lowerCAmelCase)					,			(__lowerCAmelCase)					,			(__lowerCAmelCase)					,			(__lowerCAmelCase))   					=   config_and_inputs
													__lowerCAmelCase   					=   {"""input_ids""": input_ids, """attention_mask""": input_mask}
													return config, inputs_dict
@require_torch
class 	lowerCAmelCase_						(							UpperCAmelCase__							,							UpperCAmelCase__							,							unittest.TestCase      ):



							'''simple docstring'''





							_snake_case			=						(
							    (
							        DistilBertModel,
							        DistilBertForMaskedLM,
							        DistilBertForMultipleChoice,
							        DistilBertForQuestionAnswering,
							        DistilBertForSequenceClassification,
							        DistilBertForTokenClassification,
							    )
							    if is_torch_available()
							    else None
							)
							_snake_case			=						(
							    {
							        '''feature-extraction''': DistilBertModel,
							        '''fill-mask''': DistilBertForMaskedLM,
							        '''question-answering''': DistilBertForQuestionAnswering,
							        '''text-classification''': DistilBertForSequenceClassification,
							        '''token-classification''': DistilBertForTokenClassification,
							        '''zero-shot''': DistilBertForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							_snake_case			=						True
							_snake_case			=						True
							_snake_case			=						True
							_snake_case			=						True




							def 					A__		(      self							)    ->   List[str]:
													__lowerCAmelCase   					=   DistilBertModelTester(self							)
													__lowerCAmelCase   					=   ConfigTester(self  ,							config_class=UpperCAmelCase__  ,							dim=37							)




							def 					A__		(      self							)    ->   List[Any]:
													self.config_tester.run_common_tests()




							def 					A__		(      self							)    ->   Optional[int]:
													__lowerCAmelCase   					=   self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_distilbert_model(*UpperCAmelCase__							)




							def 					A__		(      self							)    ->   Union[str, Any]:
													__lowerCAmelCase   					=   self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_distilbert_for_masked_lm(*UpperCAmelCase__							)




							def 					A__		(      self							)    ->   Any:
													__lowerCAmelCase   					=   self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_distilbert_for_question_answering(*UpperCAmelCase__							)




							def 					A__		(      self							)    ->   Dict:
													__lowerCAmelCase   					=   self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_distilbert_for_sequence_classification(*UpperCAmelCase__							)




							def 					A__		(      self							)    ->   List[Any]:
													__lowerCAmelCase   					=   self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_distilbert_for_token_classification(*UpperCAmelCase__							)




							def 					A__		(      self							)    ->   str:
													__lowerCAmelCase   					=   self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_distilbert_for_multiple_choice(*UpperCAmelCase__							)




							@slow
							def 					A__		(      self							)    ->   Tuple:

													for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			__lowerCAmelCase   					=   DistilBertModel.from_pretrained(UpperCAmelCase__							)
																			self.assertIsNotNone(UpperCAmelCase__							)




							@slow
							@require_torch_gpu
							def 					A__		(      self							)    ->   Dict:
													__lowerCAmelCase					,			__lowerCAmelCase   					=   self.model_tester.prepare_config_and_inputs_for_common()
													for model_class in self.all_model_classes:
																			# BertForMultipleChoice behaves incorrectly in JIT environments.
																			if model_class == DistilBertForMultipleChoice:
																									return

																			__lowerCAmelCase   					=   True
																			__lowerCAmelCase   					=   model_class(config=UpperCAmelCase__							)

																			__lowerCAmelCase   					=   self._prepare_for_class(UpperCAmelCase__  ,							UpperCAmelCase__							)
																			__lowerCAmelCase   					=   torch.jit.trace(
																			    UpperCAmelCase__  ,							(inputs_dict["""input_ids"""].to("""cpu"""							), inputs_dict["""attention_mask"""].to("""cpu"""							))							)

																			with tempfile.TemporaryDirectory() as tmp:
																									torch.jit.save(UpperCAmelCase__  ,							os.path.join(UpperCAmelCase__  ,							"""traced_model.pt"""							)							)
																									__lowerCAmelCase   					=   torch.jit.load(os.path.join(UpperCAmelCase__  ,							"""traced_model.pt"""							)  ,							map_location=UpperCAmelCase__							)
																									loaded(inputs_dict["""input_ids"""].to(UpperCAmelCase__							)  ,							inputs_dict["""attention_mask"""].to(UpperCAmelCase__							)							)
@require_torch
class 	lowerCAmelCase_						(							unittest.TestCase      ):



							'''simple docstring'''





							@slow
							def 					A__		(      self							)    ->   int:
													__lowerCAmelCase   					=   DistilBertModel.from_pretrained("""distilbert-base-uncased"""							)
													__lowerCAmelCase   					=   torch.tensor([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]							)
													__lowerCAmelCase   					=   torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]							)
													with torch.no_grad():
																			__lowerCAmelCase   					=   model(UpperCAmelCase__  ,							attention_mask=UpperCAmelCase__							)[0]
													__lowerCAmelCase   					=   torch.Size((1, 11, 768)							)
													self.assertEqual(output.shape  ,							UpperCAmelCase__							)
													__lowerCAmelCase   					=   torch.tensor(
													    [[[-0.1_639, 0.3_299, 0.1_648], [-0.1_746, 0.3_289, 0.1_710], [-0.1_884, 0.3_357, 0.1_810]]]							)

													self.assertTrue(torch.allclose(output[:, 1:4, 1:4]  ,							UpperCAmelCase__  ,							atol=1e-4							)							)






import tempfile
import unittest

from make_student import create_student_by_copying_alternating_layers

from transformers import AutoConfig
from transformers.file_utils import cached_property
from transformers.testing_utils import require_torch


_lowerCamelCase	:		List[Any]							       =   """sshleifer/bart-tiny-random"""
_lowerCamelCase	:		List[Any]							       =   """patrickvonplaten/t5-tiny-random"""


@require_torch
class 				UpperCamelCase_    ( unittest.TestCase     ):

    '''simple docstring'''






    @cached_property
    def        SCREAMING_SNAKE_CASE							(			self      :	Union[str, Any])						->int:





           '''simple docstring'''

           return AutoConfig.from_pretrained(UpperCAmelCase__)


    def        SCREAMING_SNAKE_CASE							(			self      :	List[str])						->Any:





           '''simple docstring'''

           A__ ,							*A__       =    create_student_by_copying_alternating_layers(UpperCAmelCase__					, tempfile.mkdtemp()					, e=1					, d=1)
           self.assertEqual(student.config.num_hidden_layers					, 1)


    def        SCREAMING_SNAKE_CASE							(			self      :	int)						->Any:





           '''simple docstring'''

           A__ ,							*A__       =    create_student_by_copying_alternating_layers(UpperCAmelCase__					, tempfile.mkdtemp()					, e=1					, d=UpperCAmelCase__)


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





           '''simple docstring'''

           A__ ,							*A__       =    create_student_by_copying_alternating_layers(UpperCAmelCase__					, tempfile.mkdtemp()					, e=1					, d=UpperCAmelCase__)
           self.assertEqual(student.config.encoder_layers					, 1)
           self.assertEqual(student.config.decoder_layers					, self.teacher_config.encoder_layers)


    def        SCREAMING_SNAKE_CASE							(			self      :	Dict)						->int:





           '''simple docstring'''

           A__ ,							*A__       =    create_student_by_copying_alternating_layers(UpperCAmelCase__					, tempfile.mkdtemp()					, e=1					, d=1)
           self.assertEqual(student.config.encoder_layers					, 1)
           self.assertEqual(student.config.decoder_layers					, 1)





    def        SCREAMING_SNAKE_CASE							(			self      :	str)						->List[Any]:





           '''simple docstring'''

           with self.assertRaises(UpperCAmelCase__):
                  create_student_by_copying_alternating_layers(UpperCAmelCase__					, tempfile.mkdtemp()					, e=UpperCAmelCase__					, d=UpperCAmelCase__)


