'''simple docstring'''





import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
SCREAMING_SNAKE_CASE_:	List[Any]    =object()

# For specifying empty leaf dict `{}`
SCREAMING_SNAKE_CASE_:	List[Any]    =object()





def       lowerCAmelCase_							(   snake_case_			:					Dict   ,       snake_case_			:					Any     )       -> Tuple:


     '''simple docstring'''




     UpperCAmelCase_			       =  tuple((re.compile(x + "$"     ) for x in qs)     )
     for i in range(len(snake_case_     ) - len(snake_case_     ) + 1     ):
          UpperCAmelCase_			       =  [x.match(snake_case_     ) for x, y in zip(snake_case_   ,       ks[i:]     )]
          if matches and all(snake_case_     ):
               return True
     return False





def       lowerCAmelCase_							(   snake_case_			:					Optional[Any]     )       -> Tuple:


     '''simple docstring'''




     def replace(snake_case_			:					Optional[int]   ,       snake_case_			:					List[str]     ):
          for rule, replacement in rules:
               if _match(snake_case_   ,       snake_case_     ):
                    return replacement
          return val

     return replace





def       lowerCAmelCase_							(   )       -> Any:


     '''simple docstring'''




     return [
         # embeddings
         (("transformer", "wpe", "embedding"), P("mp"   ,       snake_case_     )),
         (("transformer", "wte", "embedding"), P("mp"   ,       snake_case_     )),
         # atention
         (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(snake_case_   ,       "mp"     )),
         (("attention", "out_proj", "kernel"), P("mp"   ,       snake_case_     )),
         (("attention", "out_proj", "bias"), None),
         # mlp
         (("mlp", "c_fc", "kernel"), P(snake_case_   ,       "mp"     )),
         (("mlp", "c_fc", "bias"), P("mp"     )),
         (("mlp", "c_proj", "kernel"), P("mp"   ,       snake_case_     )),
         (("mlp", "c_proj", "bias"), None),
         # layer norms
         ((r"ln_\d+", "bias"), None),
         ((r"\d+", r"ln_\d+", "scale"), None),
         (("ln_f", "bias"), None),
         (("ln_f", "scale"), None),
     ]






def       lowerCAmelCase_							(   snake_case_			:					int     )       -> Any:


     '''simple docstring'''




     UpperCAmelCase_			       =  _get_partition_rules()
     UpperCAmelCase_			       =  _replacement_rules(snake_case_     )
     UpperCAmelCase_			       =  {k: _unmatched for k in flatten_dict(snake_case_     )}
     UpperCAmelCase_			       =  {k: replace(snake_case_   ,       snake_case_     ) for k, v in initd.items()}
     assert _unmatched not in result.values(), "Incomplete partition spec."
     return freeze(unflatten_dict(snake_case_     )     )






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						(						__snake_case       ):






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






	def       _a(self				: int    ) ->			Dict:
							_lowercase       :	str           =   ids_tensor([self.batch_size, self.seq_length]      ,	self.vocab_size    )

							_lowercase       :	Tuple           =   None
							if self.use_input_mask:
													_lowercase       :	int           =   random_attention_mask([self.batch_size, self.seq_length]    )

							_lowercase       :	Tuple           =   None
							_lowercase       :	Union[str, Any]           =   None
							_lowercase       :	Tuple           =   None
							if self.use_labels:
													_lowercase       :	Dict           =   ids_tensor([self.batch_size]      ,	self.type_sequence_label_size    )
													_lowercase       :	Any           =   ids_tensor([self.batch_size, self.seq_length]      ,	self.num_labels    )
													_lowercase       :	str           =   ids_tensor([self.batch_size]      ,	self.num_choices    )

							_lowercase       :	Dict           =   self.get_config()

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






	def       _a(self				: Any    ) ->			Any:
							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				: int      ,	snake_case				: Optional[Any]      ,	snake_case				: Tuple      ,	snake_case				: List[str]      ,	snake_case				: Tuple      ,	snake_case				: Any      ,	snake_case				: Dict    ) ->			Optional[int]:
							_lowercase       :	Optional[int]           =   DistilBertModel(config=snake_case    )
							model.to(snake_case    )
							model.eval()
							_lowercase       :	List[Any]           =   model(snake_case      ,	snake_case    )
							_lowercase       :	int           =   model(snake_case    )
							self.parent.assertEqual(result.last_hidden_state.shape      ,	(self.batch_size, self.seq_length, self.hidden_size)    )






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






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






	def       _a(self				: int      ,	snake_case				: Union[str, Any]      ,	snake_case				: Optional[Any]      ,	snake_case				: int      ,	snake_case				: Optional[int]      ,	snake_case				: Optional[int]      ,	snake_case				: str    ) ->			str:
							_lowercase       :	str           =   self.num_labels
							_lowercase       :	List[str]           =   DistilBertForTokenClassification(config=snake_case    )
							model.to(snake_case    )
							model.eval()

							_lowercase       :	str           =   model(snake_case      ,	attention_mask=snake_case      ,	labels=snake_case    )
							self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.seq_length, self.num_labels)    )






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






	def       _a(self				: List[str]    ) ->			List[str]:
							_lowercase       :	Union[str, Any]           =   self.prepare_config_and_inputs()
							((_lowercase)      ,    (_lowercase)      ,    (_lowercase)      ,    (_lowercase)      ,    (_lowercase)      ,    (_lowercase))       :	Union[str, Any]           =   config_and_inputs
							_lowercase       :	Union[str, Any]           =   {"input_ids": input_ids, "attention_mask": input_mask}
							return config, inputs_dict




@require_torch
class   __lowercase						(						__snake_case   , __snake_case   , unittest.TestCase       ):
	_A       =   (
	    (
	        DistilBertModel,
	        DistilBertForMaskedLM,
	        DistilBertForMultipleChoice,
	        DistilBertForQuestionAnswering,
	        DistilBertForSequenceClassification,
	        DistilBertForTokenClassification,
	    )
	    if is_torch_available()
	    else None
	)
	_A       =   (
	    {
	        "feature-extraction": DistilBertModel,
	        "fill-mask": DistilBertForMaskedLM,
	        "question-answering": DistilBertForQuestionAnswering,
	        "text-classification": DistilBertForSequenceClassification,
	        "token-classification": DistilBertForTokenClassification,
	        "zero-shot": DistilBertForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	_A       =   True
	_A       =   True
	_A       =   True
	_A       =   True






	def       _a(self				: Dict    ) ->			List[Any]:
							_lowercase       :	Optional[Any]           =   DistilBertModelTester(self    )
							_lowercase       :	str           =   ConfigTester(self      ,	config_class=snake_case      ,	dim=37    )






	def       _a(self				: int    ) ->			List[str]:
							self.config_tester.run_common_tests()






	def       _a(self				: Optional[Any]    ) ->			Optional[int]:
							_lowercase       :	Tuple           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_model(*snake_case    )






	def       _a(self				: Any    ) ->			int:
							_lowercase       :	str           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_masked_lm(*snake_case    )






	def       _a(self				: Dict    ) ->			List[Any]:
							_lowercase       :	List[str]           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_question_answering(*snake_case    )






	def       _a(self				: str    ) ->			Tuple:
							_lowercase       :	Tuple           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_sequence_classification(*snake_case    )






	def       _a(self				: Any    ) ->			List[Any]:
							_lowercase       :	Optional[Any]           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_token_classification(*snake_case    )






	def       _a(self				: Optional[int]    ) ->			Optional[int]:
							_lowercase       :	Any           =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_multiple_choice(*snake_case    )






	@slow
	def       _a(self				: Optional[Any]    ) ->			Dict:
							for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													_lowercase       :	Tuple           =   DistilBertModel.from_pretrained(snake_case    )
													self.assertIsNotNone(snake_case    )






	@slow
	@require_torch_gpu
	def       _a(self				: Optional[int]    ) ->			Optional[int]:
							_lowercase      ,    _lowercase       :	Optional[int]           =   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

													_lowercase       :	str           =   True
													_lowercase       :	Tuple           =   model_class(config=snake_case    )

													_lowercase       :	str           =   self._prepare_for_class(snake_case      ,	snake_case    )
													_lowercase       :	Optional[int]           =   torch.jit.trace(
													    snake_case      ,	(inputs_dict["input_ids"].to("cpu"    ), inputs_dict["attention_mask"].to("cpu"    ))    )

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




@require_torch
class   __lowercase						(						unittest.TestCase       ):






	@slow
	def       _a(self				: int    ) ->			str:
							_lowercase       :	Any           =   DistilBertModel.from_pretrained("distilbert-base-uncased"    )
							_lowercase       :	Dict           =   torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]    )
							_lowercase       :	List[Any]           =   torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]    )
							with torch.no_grad():
													_lowercase       :	Optional[int]           =   model(snake_case      ,	attention_mask=snake_case    )[0]
							_lowercase       :	List[str]           =   torch.Size((1, 11, 768)    )
							self.assertEqual(output.shape      ,	snake_case    )
							_lowercase       :	Any           =   torch.tensor(
							    [[[-0.16_39, 0.32_99, 0.16_48], [-0.17_46, 0.32_89, 0.17_10], [-0.18_84, 0.33_57, 0.18_10]]]    )

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




