



"""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    __lowercase     (					lowercase_):







						"""simple docstring"""







						def __init__(self   ,		lowercase__   ,		lowercase__=13   ,		lowercase__=7   ,		lowercase__=True   ,		lowercase__=True   ,		lowercase__=False   ,		lowercase__=True   ,		lowercase__=99   ,		lowercase__=32   ,		lowercase__=5   ,		lowercase__=4   ,		lowercase__=37   ,		lowercase__="gelu"   ,		lowercase__=0.1   ,		lowercase__=0.1   ,		lowercase__=5_12   ,		lowercase__=16   ,		lowercase__=2   ,		lowercase__=0.02   ,		lowercase__=3   ,		lowercase__=4   ,		lowercase__=None   ,		):
											snake_case_					:		Union[str, Any] 					=					parent
											snake_case_					:		Optional[Any] 					=					batch_size
											snake_case_					:		str 					=					seq_length
											snake_case_					:		Dict 					=					is_training
											snake_case_					:		Union[str, Any] 					=					use_input_mask
											snake_case_					:		Tuple 					=					use_token_type_ids
											snake_case_					:		List[str] 					=					use_labels
											snake_case_					:		str 					=					vocab_size
											snake_case_					:		Optional[int] 					=					hidden_size
											snake_case_					:		int 					=					num_hidden_layers
											snake_case_					:		int 					=					num_attention_heads
											snake_case_					:		int 					=					intermediate_size
											snake_case_					:		List[Any] 					=					hidden_act
											snake_case_					:		Optional[int] 					=					hidden_dropout_prob
											snake_case_					:		Tuple 					=					attention_probs_dropout_prob
											snake_case_					:		Optional[Any] 					=					max_position_embeddings
											snake_case_					:		Tuple 					=					type_vocab_size
											snake_case_					:		List[Any] 					=					type_sequence_label_size
											snake_case_					:		Optional[int] 					=					initializer_range
											snake_case_					:		int 					=					num_labels
											snake_case_					:		Tuple 					=					num_choices
											snake_case_					:		int 					=					scope







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		List[str] 					=					ids_tensor([self.batch_size, self.seq_length]   ,		self.vocab_size       )

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

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

											snake_case_					:		List[Any] 					=					self.get_config()

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







						def 						__UpperCamelCase							(self       ):
											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   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Union[str, Any] 					=					DistilBertModel(config=lowercase__       )
											model.to(lowercase__       )
											model.eval()
											snake_case_					:		Optional[int] 					=					model(lowercase__   ,		lowercase__       )
											snake_case_					:		Union[str, Any] 					=					model(lowercase__       )
											self.parent.assertEqual(result.last_hidden_state.shape   ,		(self.batch_size, self.seq_length, self.hidden_size)       )







						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Optional[int] 					=					DistilBertForMaskedLM(config=lowercase__       )
											model.to(lowercase__       )
											model.eval()
											snake_case_					:		List[Any] 					=					model(lowercase__   ,		attention_mask=lowercase__   ,		labels=lowercase__       )
											self.parent.assertEqual(result.logits.shape   ,		(self.batch_size, self.seq_length, self.vocab_size)       )







						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		List[Any] 					=					DistilBertForQuestionAnswering(config=lowercase__       )
											model.to(lowercase__       )
											model.eval()
											snake_case_					:		Optional[int] 					=					model(
											    lowercase__   ,		attention_mask=lowercase__   ,		start_positions=lowercase__   ,		end_positions=lowercase__       )
											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   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		List[Any] 					=					self.num_labels
											snake_case_					:		int 					=					DistilBertForSequenceClassification(lowercase__       )
											model.to(lowercase__       )
											model.eval()
											snake_case_					:		str 					=					model(lowercase__   ,		attention_mask=lowercase__   ,		labels=lowercase__       )
											self.parent.assertEqual(result.logits.shape   ,		(self.batch_size, self.num_labels)       )







						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		List[str] 					=					self.num_labels
											snake_case_					:		Union[str, Any] 					=					DistilBertForTokenClassification(config=lowercase__       )
											model.to(lowercase__       )
											model.eval()

											snake_case_					:		Optional[Any] 					=					model(lowercase__   ,		attention_mask=lowercase__   ,		labels=lowercase__       )
											self.parent.assertEqual(result.logits.shape   ,		(self.batch_size, self.seq_length, self.num_labels)       )







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







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Any 					=					self.prepare_config_and_inputs()
											((snake_case_)					,			(snake_case_)					,			(snake_case_)					,			(snake_case_)					,			(snake_case_)					,			(snake_case_))					:		Optional[Any] 					=					config_and_inputs
											snake_case_					:		Union[str, Any] 					=					{"""input_ids""": input_ids, """attention_mask""": input_mask}
											return config, inputs_dict







@require_torch
class    __lowercase     (					lowercase_					,    lowercase_					,    unittest.TestCase):







						"""simple docstring"""







						_A		:							str							       =		(
						    (
						        DistilBertModel,
						        DistilBertForMaskedLM,
						        DistilBertForMultipleChoice,
						        DistilBertForQuestionAnswering,
						        DistilBertForSequenceClassification,
						        DistilBertForTokenClassification,
						    )
						    if is_torch_available()
						    else None
						)
						_A		:							List[str]							       =		(
						    {
						        '''feature-extraction''': DistilBertModel,
						        '''fill-mask''': DistilBertForMaskedLM,
						        '''question-answering''': DistilBertForQuestionAnswering,
						        '''text-classification''': DistilBertForSequenceClassification,
						        '''token-classification''': DistilBertForTokenClassification,
						        '''zero-shot''': DistilBertForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)
						_A		:							List[str]							       =		True
						_A		:							Dict							       =		True
						_A		:							Optional[Any]							       =		True
						_A		:							Any							       =		True







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		str 					=					DistilBertModelTester(self       )
											snake_case_					:		Any 					=					ConfigTester(self   ,		config_class=lowercase__   ,		dim=37       )







						def 						__UpperCamelCase							(self       ):
											self.config_tester.run_common_tests()







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		List[str] 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_distilbert_model(*lowercase__       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		str 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_distilbert_for_masked_lm(*lowercase__       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		int 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_distilbert_for_question_answering(*lowercase__       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Optional[Any] 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowercase__       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Any 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_distilbert_for_token_classification(*lowercase__       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		List[str] 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowercase__       )







						@slow
						def 						__UpperCamelCase							(self       ):
											for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																snake_case_					:		int 					=					DistilBertModel.from_pretrained(lowercase__       )
																self.assertIsNotNone(lowercase__       )







						@slow
						@require_torch_gpu
						def 						__UpperCamelCase							(self       ):
											snake_case_					,			snake_case_					:		List[str] 					=					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

																snake_case_					:		Any 					=					True
																snake_case_					:		Union[str, Any] 					=					model_class(config=lowercase__       )

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

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







@require_torch
class    __lowercase     (					unittest.TestCase):







						"""simple docstring"""







						@slow
						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Dict 					=					DistilBertModel.from_pretrained("""distilbert-base-uncased"""       )
											snake_case_					:		Tuple 					=					torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]       )
											snake_case_					:		Dict 					=					torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]       )
											with torch.no_grad():
																snake_case_					:		int 					=					model(lowercase__   ,		attention_mask=lowercase__       )[0]
											snake_case_					:		int 					=					torch.Size((1, 11, 7_68)       )
											self.assertEqual(output.shape   ,		lowercase__       )
											snake_case_					:		Dict 					=					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]   ,		lowercase__   ,		atol=1e-4       )       )









"""simple docstring"""

from manim import *







class    __lowercase     (					_UpperCAmelCase):







						"""simple docstring"""







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Union[str, Any] 					=					Rectangle(height=0.5   ,		width=0.5       )
											snake_case_					:		str 					=					Rectangle(height=0.46   ,		width=0.46       ).set_stroke(width=0       )

											snake_case_					:		Optional[Any] 					=					[mem.copy() for i in range(6       )]
											snake_case_					:		str 					=					[mem.copy() for i in range(6       )]
											snake_case_					:		str 					=					VGroup(*lowercase__       ).arrange(lowercase__   ,		buff=0       )
											snake_case_					:		Any 					=					VGroup(*lowercase__       ).arrange(lowercase__   ,		buff=0       )
											snake_case_					:		List[str] 					=					VGroup(lowercase__   ,		lowercase__       ).arrange(lowercase__   ,		buff=0       )
											snake_case_					:		List[Any] 					=					Text("""CPU"""   ,		font_size=24       )
											snake_case_					:		Tuple 					=					Group(lowercase__   ,		lowercase__       ).arrange(lowercase__   ,		buff=0.5   ,		aligned_edge=lowercase__       )
											cpu.move_to([-2.5, -0.5, 0]       )
											self.add(lowercase__       )

											snake_case_					:		List[Any] 					=					[mem.copy() for i in range(4       )]
											snake_case_					:		Tuple 					=					VGroup(*lowercase__       ).arrange(lowercase__   ,		buff=0       )
											snake_case_					:		List[str] 					=					Text("""GPU"""   ,		font_size=24       )
											snake_case_					:		Any 					=					Group(lowercase__   ,		lowercase__       ).arrange(lowercase__   ,		buff=0.5   ,		aligned_edge=lowercase__       )
											gpu.move_to([-1, -1, 0]       )
											self.add(lowercase__       )

											snake_case_					:		Optional[Any] 					=					[mem.copy() for i in range(6       )]
											snake_case_					:		List[Any] 					=					VGroup(*lowercase__       ).arrange(lowercase__   ,		buff=0       )

											snake_case_					:		Dict 					=					Text("""Model"""   ,		font_size=24       )
											snake_case_					:		int 					=					Group(lowercase__   ,		lowercase__       ).arrange(lowercase__   ,		buff=0.5   ,		aligned_edge=lowercase__       )
											model.move_to([3, -1.0, 0]       )
											self.add(lowercase__       )

											snake_case_					:		Dict 					=					[]
											for i, rect in enumerate(lowercase__       ):
																rect.set_stroke(lowercase__       )
																# target = fill.copy().set_fill(YELLOW, opacity=0.7)
																# target.move_to(rect)
																# self.add(target)

																snake_case_					:		List[str] 					=					Rectangle(height=0.46 / 4   ,		width=0.46 / 3       ).set_stroke(width=0.0       ).set_fill(lowercase__   ,		opacity=0.7       )

																if i == 0:
																					cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT       )   ,		buff=0.02   ,		direction=lowercase__       )
																					cpu_target.set_x(cpu_target.get_x() + 0.1       )
																elif i == 3:
																					cpu_target.next_to(cpu_targs[0]   ,		direction=lowercase__   ,		buff=0.0       )
																else:
																					cpu_target.next_to(cpu_targs[i - 1]   ,		direction=lowercase__   ,		buff=0.0       )
																self.add(lowercase__       )
																cpu_targs.append(lowercase__       )

											snake_case_					:		List[str] 					=					[mem.copy() for i in range(6       )]
											snake_case_					:		List[str] 					=					VGroup(*lowercase__       ).arrange(lowercase__   ,		buff=0       )

											snake_case_					:		str 					=					Text("""Loaded Checkpoint"""   ,		font_size=24       )
											snake_case_					:		Any 					=					Group(lowercase__   ,		lowercase__       ).arrange(lowercase__   ,		aligned_edge=lowercase__   ,		buff=0.4       )
											checkpoint.move_to([3, 0.5, 0]       )

											snake_case_					:		Optional[Any] 					=					Square(side_length=2.2       )
											key.move_to([-5, 2, 0]       )

											snake_case_					:		Union[str, Any] 					=					MarkupText(
											    f'<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'   ,		font_size=18   ,		)

											key_text.move_to([-5, 2.4, 0]       )

											self.add(lowercase__   ,		lowercase__       )

											snake_case_					:		List[Any] 					=					MarkupText(
											    f'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'   ,		font_size=18   ,		)

											blue_text.next_to(lowercase__   ,		DOWN * 2.4   ,		aligned_edge=key_text.get_left()       )

											snake_case_					:		List[Any] 					=					MarkupText(
											    f'Next, a <i><span fgcolor="{BLUE}">second</span></i> model is loaded into memory,\nwith the weights of a <span fgcolor="{BLUE}">single shard</span>.'   ,		font_size=24   ,		)
											step_a.move_to([2, 2, 0]       )
											self.play(Write(lowercase__       )   ,		Write(lowercase__       )       )

											self.play(Write(lowercase__   ,		run_time=1       )   ,		Create(lowercase__   ,		run_time=1       )       )

											snake_case_					:		Optional[int] 					=					[]
											snake_case_					:		List[str] 					=					[]
											for i, rect in enumerate(lowercase__       ):
																snake_case_					:		Optional[Any] 					=					fill.copy().set_fill(lowercase__   ,		opacity=0.7       )
																target.move_to(lowercase__       )
																first_animations.append(GrowFromCenter(lowercase__   ,		run_time=1       )       )

																snake_case_					:		List[Any] 					=					target.copy()
																cpu_target.generate_target()
																if i < 5:
																					cpu_target.target.move_to(cpu_left_col_base[i + 1]       )
																else:
																					cpu_target.target.move_to(cpu_right_col_base[i - 5]       )
																second_animations.append(MoveToTarget(lowercase__   ,		run_time=1.5       )       )

											self.play(*lowercase__       )
											self.play(*lowercase__       )
											self.wait()




