


import argparse
import os
import re


lowerCamelCase__							:						Optional[int]							     =   """src/transformers/models/auto"""


# re pattern that matches mapping introductions:
#    SUPER_MODEL_MAPPING_NAMES = OrderedDict or SUPER_MODEL_MAPPING = OrderedDict
lowerCamelCase__							:						List[Any]							     =   re.compile(r"""[A-Z_]+_MAPPING(\s+|_[A-Z_]+\s+)=\s+OrderedDict""")
# re pattern that matches identifiers in mappings
lowerCamelCase__							:						Dict							     =   re.compile(r"""\s*\(\s*\"(\S[^\"]+)\"""")







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase = False				)    ->				Union[str, Any]:
  with open(__lowerCAmelCase							,   '''r'''							,   encoding='''utf-8'''				) as f:
    snake_case__						= f.read()

  snake_case__						= content.split('''\n'''				)
  snake_case__						= []
  snake_case__						= 0
  while line_idx < len(__lowerCAmelCase				):
    if _re_intro_mapping.search(lines[line_idx]				) is not None:
      snake_case__						= len(re.search(r'''^(\s*)\S'''							,   lines[line_idx]				).groups()[0]				) + 8
      # Start of a new mapping!
      while not lines[line_idx].startswith(''' ''' * indent + '''('''				):
        new_lines.append(lines[line_idx]				)
        line_idx += 1

      snake_case__						= []
      while lines[line_idx].strip() != "]":
        # Blocks either fit in one line or not
        if lines[line_idx].strip() == "(":
          snake_case__						= line_idx
          while not lines[line_idx].startswith(''' ''' * indent + ''')'''				):
            line_idx += 1
          blocks.append('''\n'''.join(lines[start_idx : line_idx + 1]				)				)
        else:
          blocks.append(lines[line_idx]				)
        line_idx += 1

      # Sort blocks by their identifiers
      snake_case__						= sorted(__lowerCAmelCase							,   key=lambda __lowerCAmelCase				: _re_identifier.search(__lowerCAmelCase				).groups()[0]				)
      new_lines += blocks
    else:
      new_lines.append(lines[line_idx]				)
      line_idx += 1

  if overwrite:
    with open(__lowerCAmelCase							,   '''w'''							,   encoding='''utf-8'''				) as f:
      f.write('''\n'''.join(__lowerCAmelCase				)				)
  elif "\n".join(__lowerCAmelCase				) != content:
    return True







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase = False				)    ->				Tuple:
  snake_case__						= [os.path.join(__lowerCAmelCase							,   __lowerCAmelCase				) for f in os.listdir(__lowerCAmelCase				) if f.endswith('''.py'''				)]
  snake_case__						= [sort_auto_mapping(__lowerCAmelCase							,   overwrite=__lowerCAmelCase				) for fname in fnames]

  if not overwrite and any(__lowerCAmelCase				):
    snake_case__						= [f for f, d in zip(__lowerCAmelCase							,   __lowerCAmelCase				) if d]
    raise ValueError(
        F"""The following files have auto mappings that need sorting: {', '.join(__lowerCAmelCase				)}. Run `make style` to fix"""
        ''' this.'''				)


if __name__ == "__main__":
   lowerCamelCase__							:						Any							     =   argparse.ArgumentParser()
   parser.add_argument("""--check_only""", action="""store_true""", help="""Whether to only check or fix style.""")
   lowerCamelCase__							:						Union[str, Any]							     =   parser.parse_args()

   sort_all_auto_mappings(not args.check_only)




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						)						)