



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

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							str     ,		lowerCamelCase_:							List[str]     ,		lowerCamelCase_:							int ):
						"""simple docstring"""






						snake_case   :			Dict						    = 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() )


A		=						{
    # 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))},
}
A		=						{
    # 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  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							str     ,		lowerCamelCase_:							Union[str, Any] ):
						"""simple docstring"""






						try:
												snake_case   :			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(_a ) )

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Optional[Any]     ,		lowerCamelCase_:							Optional[int] ):
						"""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(_a ) )
						elif n_student == 1:
												return [n_teacher - 1]
						else:
												return LAYERS_TO_SUPERVISE[n_teacher][n_student]

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Dict     ,		lowerCamelCase_:							Union[str, Any] = "student"     ,		lowerCamelCase_:							Dict = None     ,		lowerCamelCase_:							Dict = None     ,		lowerCamelCase_:							List[Any]=False     ,		lowerCamelCase_:							Any=None     ,		lowerCamelCase_:							int=None     ,		**lowerCamelCase_:							str     ,		):
						"""simple docstring"""






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

						try:
												snake_case   :			str						    = teacher.config.encoder_layers, teacher.config.decoder_layers
												if e is None:
																		snake_case   :			List[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   :			str						    = teacher.config.num_encoder_layers, teacher.config.num_decoder_layers
												else:
																		snake_case   :			Tuple						    = teacher.config.num_layers, teacher.config.num_decoder_layers
												if e is None:
																		snake_case   :			Optional[Any]						    = teacher_e
												if d is None:
																		snake_case   :			List[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(_a )

						# Copy weights
						snake_case   :			Tuple						    = teacher.config_class(**_a )
						snake_case   :			Any						    = AutoModelForSeqaSeqLM.from_config(_a )
						# Start by copying the full teacher state dict this will copy the first N teacher layers to the student.
						snake_case   :			List[str]						    = 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
												snake_case   :			Optional[int]						    = 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:
												snake_case   :			List[int]						    = pick_layers_to_copy(_a     ,		_a )
						if d_layers_to_copy is None:
												snake_case   :			List[int]						    = 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}''' )
						snake_case   :			Optional[int]						    = {
						    "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)






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 _UpperCamelCase				(				__A ):
			'''simple docstring'''







			def __init__(      self				:			Dict   ,   a				:			Tuple   ,   a				:			Any=13   ,   a				:			Any=7   ,   a				:			Union[str, Any]=True   ,   a				:			List[Any]=True   ,   a				:			List[str]=False   ,   a				:			List[str]=True   ,   a				:			Any=99   ,   a				:			str=32   ,   a				:			Any=5   ,   a				:			Optional[int]=4   ,   a				:			Union[str, Any]=37   ,   a				:			Dict="gelu"   ,   a				:			List[Any]=0.1   ,   a				:			Optional[Any]=0.1   ,   a				:			List[str]=512   ,   a				:			Union[str, Any]=16   ,   a				:			str=2   ,   a				:			Dict=0.02   ,   a				:			Optional[int]=3   ,   a				:			Union[str, Any]=4   ,   a				:			int=None   ,   )				->							Tuple:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							parent
				SCREAMING_SNAKE_CASE :   Any					  =							batch_size
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							seq_length
				SCREAMING_SNAKE_CASE :   List[Any]					  =							is_training
				SCREAMING_SNAKE_CASE :   int					  =							use_input_mask
				SCREAMING_SNAKE_CASE :   Tuple					  =							use_token_type_ids
				SCREAMING_SNAKE_CASE :   str					  =							use_labels
				SCREAMING_SNAKE_CASE :   Any					  =							vocab_size
				SCREAMING_SNAKE_CASE :   List[Any]					  =							hidden_size
				SCREAMING_SNAKE_CASE :   str					  =							num_hidden_layers
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							num_attention_heads
				SCREAMING_SNAKE_CASE :   Tuple					  =							intermediate_size
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							hidden_act
				SCREAMING_SNAKE_CASE :   Dict					  =							hidden_dropout_prob
				SCREAMING_SNAKE_CASE :   str					  =							attention_probs_dropout_prob
				SCREAMING_SNAKE_CASE :   Tuple					  =							max_position_embeddings
				SCREAMING_SNAKE_CASE :   List[str]					  =							type_vocab_size
				SCREAMING_SNAKE_CASE :   List[str]					  =							type_sequence_label_size
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							initializer_range
				SCREAMING_SNAKE_CASE :   Tuple					  =							num_labels
				SCREAMING_SNAKE_CASE :   Tuple					  =							num_choices
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							scope
			def 				__UpperCamelCase			(      self				:			Union[str, Any]						)				->							List[str]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							ids_tensor([self.batch_size, self.seq_length]   ,   self.vocab_size						)

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							None
				if self.use_input_mask:
					SCREAMING_SNAKE_CASE :   str					  =							random_attention_mask([self.batch_size, self.seq_length]						)

				SCREAMING_SNAKE_CASE :   int					  =							None
				SCREAMING_SNAKE_CASE :   List[Any]					  =							None
				SCREAMING_SNAKE_CASE :   List[str]					  =							None
				if self.use_labels:
					SCREAMING_SNAKE_CASE :   Tuple					  =							ids_tensor([self.batch_size]   ,   self.type_sequence_label_size						)
					SCREAMING_SNAKE_CASE :   str					  =							ids_tensor([self.batch_size, self.seq_length]   ,   self.num_labels						)
					SCREAMING_SNAKE_CASE :   Optional[Any]					  =							ids_tensor([self.batch_size]   ,   self.num_choices						)

				SCREAMING_SNAKE_CASE :   int					  =							self.get_config()

				return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
			def 				__UpperCamelCase			(      self				:			Dict						)				->							str:

				"""simple docstring"""

				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 				__UpperCamelCase			(      self				:			Optional[Any]   ,   a				:			int   ,   a				:			Optional[int]   ,   a				:			Optional[int]   ,   a				:			Dict   ,   a				:			str   ,   a				:			str						)				->							int:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   int					  =							DistilBertModel(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							model(a   ,   a						)
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							model(a						)
				self.parent.assertEqual(result.last_hidden_state.shape   ,   (self.batch_size, self.seq_length, self.hidden_size)						)
			def 				__UpperCamelCase			(      self				:			Tuple   ,   a				:			Optional[int]   ,   a				:			Dict   ,   a				:			Tuple   ,   a				:			int   ,   a				:			int   ,   a				:			Any						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[int]					  =							DistilBertForMaskedLM(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   str					  =							model(a   ,   attention_mask=a   ,   labels=a						)
				self.parent.assertEqual(result.logits.shape   ,   (self.batch_size, self.seq_length, self.vocab_size)						)
			def 				__UpperCamelCase			(      self				:			List[Any]   ,   a				:			int   ,   a				:			Optional[Any]   ,   a				:			Optional[Any]   ,   a				:			str   ,   a				:			str   ,   a				:			Tuple						)				->							List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							DistilBertForQuestionAnswering(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   List[str]					  =							model(
				    a   ,   attention_mask=a   ,   start_positions=a   ,   end_positions=a						)
				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 				__UpperCamelCase			(      self				:			Optional[int]   ,   a				:			str   ,   a				:			Any   ,   a				:			int   ,   a				:			Optional[Any]   ,   a				:			int   ,   a				:			str						)				->							Optional[int]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							self.num_labels
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							DistilBertForSequenceClassification(a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							model(a   ,   attention_mask=a   ,   labels=a						)
				self.parent.assertEqual(result.logits.shape   ,   (self.batch_size, self.num_labels)						)
			def 				__UpperCamelCase			(      self				:			Optional[Any]   ,   a				:			List[Any]   ,   a				:			Optional[int]   ,   a				:			Union[str, Any]   ,   a				:			Dict   ,   a				:			Any   ,   a				:			Optional[Any]						)				->							Any:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   str					  =							self.num_labels
				SCREAMING_SNAKE_CASE :   List[str]					  =							DistilBertForTokenClassification(config=a						)
				model.to(a						)
				model.eval()

				SCREAMING_SNAKE_CASE :   List[str]					  =							model(a   ,   attention_mask=a   ,   labels=a						)
				self.parent.assertEqual(result.logits.shape   ,   (self.batch_size, self.seq_length, self.num_labels)						)
			def 				__UpperCamelCase			(      self				:			int   ,   a				:			Any   ,   a				:			Optional[int]   ,   a				:			Union[str, Any]   ,   a				:			Tuple   ,   a				:			Optional[int]   ,   a				:			Tuple						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[str]					  =							self.num_choices
				SCREAMING_SNAKE_CASE :   Any					  =							DistilBertForMultipleChoice(config=a						)
				model.to(a						)
				model.eval()
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							input_ids.unsqueeze(1						).expand(-1   ,   self.num_choices   ,   -1						).contiguous()
				SCREAMING_SNAKE_CASE :   Dict					  =							input_mask.unsqueeze(1						).expand(-1   ,   self.num_choices   ,   -1						).contiguous()
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							model(
				    a   ,   attention_mask=a   ,   labels=a   ,   )
				self.parent.assertEqual(result.logits.shape   ,   (self.batch_size, self.num_choices)						)






			def 				__UpperCamelCase			(      self				:			Union[str, Any]						)				->							Optional[int]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							self.prepare_config_and_inputs()
				((SCREAMING_SNAKE_CASE)  ,(SCREAMING_SNAKE_CASE)  ,(SCREAMING_SNAKE_CASE)  ,(SCREAMING_SNAKE_CASE)  ,(SCREAMING_SNAKE_CASE)  ,(SCREAMING_SNAKE_CASE)) :   Tuple					  =							config_and_inputs
				SCREAMING_SNAKE_CASE :   int					  =							{"input_ids": input_ids, "attention_mask": input_mask}
				return config, inputs_dict






@require_torch
class _UpperCamelCase				(				__A				,						__A				,						unittest.TestCase ):
			'''simple docstring'''







			lowerCamelCase__							     =(
			    (
			        DistilBertModel,
			        DistilBertForMaskedLM,
			        DistilBertForMultipleChoice,
			        DistilBertForQuestionAnswering,
			        DistilBertForSequenceClassification,
			        DistilBertForTokenClassification,
			    )
			    if is_torch_available()
			    else None
			)
			lowerCamelCase__							     =(
			    {
			        'feature-extraction': DistilBertModel,
			        'fill-mask': DistilBertForMaskedLM,
			        'question-answering': DistilBertForQuestionAnswering,
			        'text-classification': DistilBertForSequenceClassification,
			        'token-classification': DistilBertForTokenClassification,
			        'zero-shot': DistilBertForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			lowerCamelCase__							     =True
			lowerCamelCase__							     =True
			lowerCamelCase__							     =True
			lowerCamelCase__							     =True
			def 				__UpperCamelCase			(      self				:			Optional[Any]						)				->							Tuple:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							DistilBertModelTester(self						)
				SCREAMING_SNAKE_CASE :   List[str]					  =							ConfigTester(self   ,   config_class=a   ,   dim=37						)
			def 				__UpperCamelCase			(      self				:			List[Any]						)				->							Tuple:

				"""simple docstring"""

				self.config_tester.run_common_tests()
			def 				__UpperCamelCase			(      self				:			Union[str, Any]						)				->							List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   int					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_model(*a						)
			def 				__UpperCamelCase			(      self				:			Tuple						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_masked_lm(*a						)
			def 				__UpperCamelCase			(      self				:			List[Any]						)				->							Optional[int]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[Any]					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_question_answering(*a						)
			def 				__UpperCamelCase			(      self				:			Union[str, Any]						)				->							Any:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[int]					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_sequence_classification(*a						)
			def 				__UpperCamelCase			(      self				:			str						)				->							List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[str]					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_token_classification(*a						)
			def 				__UpperCamelCase			(      self				:			List[Any]						)				->							str:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[Any]					  =							self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_distilbert_for_multiple_choice(*a						)
			@slow
			def 				__UpperCamelCase			(      self				:			int						)				->							Any:

				"""simple docstring"""

				for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					SCREAMING_SNAKE_CASE :   Optional[Any]					  =							DistilBertModel.from_pretrained(a						)
					self.assertIsNotNone(a						)






			@slow
			@require_torch_gpu
			def 				__UpperCamelCase			(      self				:			List[Any]						)				->							Optional[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE :   Dict					  =							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

					SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							True
					SCREAMING_SNAKE_CASE :   Any					  =							model_class(config=a						)

					SCREAMING_SNAKE_CASE :   List[Any]					  =							self._prepare_for_class(a   ,   a						)
					SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							torch.jit.trace(
					    a   ,   (inputs_dict["input_ids"].to("cpu"						), inputs_dict["attention_mask"].to("cpu"						))						)

					with tempfile.TemporaryDirectory() as tmp:
						torch.jit.save(a   ,   os.path.join(a   ,   "traced_model.pt"						)						)
						SCREAMING_SNAKE_CASE :   Tuple					  =							torch.jit.load(os.path.join(a   ,   "traced_model.pt"						)   ,   map_location=a						)
						loaded(inputs_dict["input_ids"].to(a						)   ,   inputs_dict["attention_mask"].to(a						)						)






@require_torch
class _UpperCamelCase				(				unittest.TestCase ):
			'''simple docstring'''







			@slow
			def 				__UpperCamelCase			(      self				:			int						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							DistilBertModel.from_pretrained("distilbert-base-uncased"						)
				SCREAMING_SNAKE_CASE :   List[str]					  =							torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]						)
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]						)
				with torch.no_grad():
					SCREAMING_SNAKE_CASE :   Optional[Any]					  =							model(a   ,   attention_mask=a						)[0]
				SCREAMING_SNAKE_CASE :   List[str]					  =							torch.Size((1, 11, 768)						)
				self.assertEqual(output.shape   ,   a						)
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							torch.tensor(
				    [[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]						)

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