"""simple docstring"""







from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
	import tensorflow as tf

	from transformers.models.distilbert.modeling_tf_distilbert import (
	    TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
	    TFDistilBertForMaskedLM,
	    TFDistilBertForMultipleChoice,
	    TFDistilBertForQuestionAnswering,
	    TFDistilBertForSequenceClassification,
	    TFDistilBertForTokenClassification,
	    TFDistilBertModel,
	)

class 		UpperCamelCase			:



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



				def 					_lowercase     (self				:      Tuple)    ->	Union[str, Any]:
					__snake_case			:							Dict      		= ids_tensor([self.batch_size, self.seq_length]	,  self.vocab_size)

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

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

					__snake_case			:							List[Any]      		= 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	,  )

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



				def 					_lowercase     (self				:      Optional[Any]	,  _A				:      Optional[Any]	,  _A				:      List[Any]	,  _A				:      List[Any]	,  _A				:      int	,  _A				:      Optional[int]	,  _A				:      List[str])    ->	Optional[Any]:
					__snake_case			:							List[str]      		= TFDistilBertModel(config=lowercase_)
					__snake_case			:							Any      		= {'''input_ids''': input_ids, '''attention_mask''': input_mask}

					__snake_case			:							List[str]      		= model(lowercase_)

					__snake_case			:							Dict      		= [input_ids, input_mask]

					__snake_case			:							Tuple      		= model(lowercase_)

					self.parent.assertEqual(result.last_hidden_state.shape	,  (self.batch_size, self.seq_length, self.hidden_size))



				def 					_lowercase     (self				:      str	,  _A				:      Optional[int]	,  _A				:      Tuple	,  _A				:      List[str]	,  _A				:      Any	,  _A				:      Dict	,  _A				:      List[Any])    ->	Optional[int]:
					__snake_case			:							int      		= TFDistilBertForMaskedLM(config=lowercase_)
					__snake_case			:							List[str]      		= {'''input_ids''': input_ids, '''attention_mask''': input_mask}
					__snake_case			:							Optional[Any]      		= model(lowercase_)
					self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.seq_length, self.vocab_size))



				def 					_lowercase     (self				:      int	,  _A				:      int	,  _A				:      int	,  _A				:      List[str]	,  _A				:      Union[str, Any]	,  _A				:      Dict	,  _A				:      Optional[int])    ->	str:
					__snake_case			:							Any      		= TFDistilBertForQuestionAnswering(config=lowercase_)
					__snake_case			:							Tuple      		= {
					    '''input_ids''': input_ids,
					    '''attention_mask''': input_mask,
					}
					__snake_case			:							List[Any]      		= model(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 					_lowercase     (self				:      Union[str, Any]	,  _A				:      Dict	,  _A				:      Union[str, Any]	,  _A				:      Any	,  _A				:      List[str]	,  _A				:      Optional[Any]	,  _A				:      Dict)    ->	int:
					__snake_case			:							Tuple      		= self.num_labels
					__snake_case			:							Optional[int]      		= TFDistilBertForSequenceClassification(lowercase_)
					__snake_case			:							Optional[Any]      		= {'''input_ids''': input_ids, '''attention_mask''': input_mask}
					__snake_case			:							str      		= model(lowercase_)
					self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.num_labels))



				def 					_lowercase     (self				:      List[str]	,  _A				:      Tuple	,  _A				:      Any	,  _A				:      Any	,  _A				:      Dict	,  _A				:      Optional[int]	,  _A				:      Any)    ->	Dict:
					__snake_case			:							List[str]      		= self.num_choices
					__snake_case			:							Optional[int]      		= TFDistilBertForMultipleChoice(lowercase_)
					__snake_case			:							List[Any]      		= tf.tile(tf.expand_dims(lowercase_	,  1)	,  (1, self.num_choices, 1))
					__snake_case			:							List[Any]      		= tf.tile(tf.expand_dims(lowercase_	,  1)	,  (1, self.num_choices, 1))
					__snake_case			:							Optional[Any]      		= {
					    '''input_ids''': multiple_choice_inputs_ids,
					    '''attention_mask''': multiple_choice_input_mask,
					}
					__snake_case			:							List[Any]      		= model(lowercase_)
					self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.num_choices))



				def 					_lowercase     (self				:      List[str]	,  _A				:      List[str]	,  _A				:      Optional[int]	,  _A				:      Optional[Any]	,  _A				:      Optional[int]	,  _A				:      Tuple	,  _A				:      Optional[int])    ->	Union[str, Any]:
					__snake_case			:							Optional[Any]      		= self.num_labels
					__snake_case			:							Optional[int]      		= TFDistilBertForTokenClassification(lowercase_)
					__snake_case			:							Any      		= {'''input_ids''': input_ids, '''attention_mask''': input_mask}
					__snake_case			:							List[Any]      		= model(lowercase_)
					self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.seq_length, self.num_labels))






				def 					_lowercase     (self				:      Any)    ->	Tuple:
					__snake_case			:							str      		= self.prepare_config_and_inputs()
					(__snake_case)			:							int      		= config_and_inputs
					__snake_case			:							Dict      		= {'''input_ids''': input_ids, '''attention_mask''': input_mask}
					return config, inputs_dict







@require_tf
class 		UpperCamelCase			(    a_      ,	a_      ,	unittest.TestCase						):
				UpperCAmelCase				:      Union[str, Any]    =		(
				    (
				        TFDistilBertModel,
				        TFDistilBertForMaskedLM,
				        TFDistilBertForQuestionAnswering,
				        TFDistilBertForSequenceClassification,
				        TFDistilBertForTokenClassification,
				        TFDistilBertForMultipleChoice,
				    )
				    if is_tf_available()
				    else None
				)
				UpperCAmelCase				:      Dict    =		(
				    {
				        """feature-extraction""": TFDistilBertModel,
				        """fill-mask""": TFDistilBertForMaskedLM,
				        """question-answering""": TFDistilBertForQuestionAnswering,
				        """text-classification""": TFDistilBertForSequenceClassification,
				        """token-classification""": TFDistilBertForTokenClassification,
				        """zero-shot""": TFDistilBertForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				UpperCAmelCase				:      Any    =		False
				UpperCAmelCase				:      str    =		False



				def 					_lowercase     (self				:      Union[str, Any])    ->	int:
					__snake_case			:							Dict      		= TFDistilBertModelTester(self)
					__snake_case			:							Optional[Any]      		= ConfigTester(self	,  config_class=lowercase_	,  dim=37)



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



				def 					_lowercase     (self				:      Union[str, Any])    ->	str:
					__snake_case			:							Tuple      		= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_distilbert_model(*lowercase_)



				def 					_lowercase     (self				:      str)    ->	Any:
					__snake_case			:							List[str]      		= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_distilbert_for_masked_lm(*lowercase_)



				def 					_lowercase     (self				:      Union[str, Any])    ->	Optional[int]:
					__snake_case			:							Optional[int]      		= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_distilbert_for_question_answering(*lowercase_)



				def 					_lowercase     (self				:      List[str])    ->	Dict:
					__snake_case			:							int      		= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowercase_)



				def 					_lowercase     (self				:      List[str])    ->	Union[str, Any]:
					__snake_case			:							Tuple      		= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowercase_)



				def 					_lowercase     (self				:      Union[str, Any])    ->	int:
					__snake_case			:							Tuple      		= self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_distilbert_for_token_classification(*lowercase_)






				@slow
				def 					_lowercase     (self				:      Tuple)    ->	Dict:
					for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]):
						__snake_case			:							Optional[int]      		= TFDistilBertModel.from_pretrained(lowercase_)
						self.assertIsNotNone(lowercase_)







@require_tf
class 		UpperCamelCase			(    unittest.TestCase						):



				@slow
				def 					_lowercase     (self				:      Optional[Any])    ->	Any:
					__snake_case			:							Tuple      		= TFDistilBertModel.from_pretrained('distilbert-base-uncased')
					__snake_case			:							Optional[Any]      		= tf.constant([[0, 1, 2, 3, 4, 5]])
					__snake_case			:							Union[str, Any]      		= model(lowercase_)[0]

					__snake_case			:							str      		= [1, 6, 7_68]
					self.assertEqual(output.shape	,  lowercase_)

					__snake_case			:							str      		= tf.constant(
					    [
					        [
					            [0.19_261_885, -0.13_732_955, 0.4_119_799],
					            [0.22_150_156, -0.07_422_661, 0.39_037_204],
					            [0.22_756_018, -0.0_896_414, 0.3_701_467],
					        ]
					    ])
					tf.debugging.assert_near(output[:, :3, :3]	,  lowercase_	,  atol=1E-4)








"""simple docstring"""





import argparse
import logging
import pickle
from collections import Counter


logging.basicConfig(
    format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s''', datefmt='''%m/%d/%Y %H:%M:%S''', level=logging.INFO
)
__UpperCamelCase  :      Any							      =		logging.getLogger(__name__)

if __name__ == "__main__":
							__UpperCamelCase  :      Union[str, Any]							      =		argparse.ArgumentParser(
							    description='''Token Counts for smoothing the masking probabilities in MLM (cf XLM/word2vec)'''
							)
							parser.add_argument(
							    '''--data_file''', type=str, default='''data/dump.bert-base-uncased.pickle''', help='''The binarized dataset.'''
							)
							parser.add_argument(
							    '''--token_counts_dump''', type=str, default='''data/token_counts.bert-base-uncased.pickle''', help='''The dump file.'''
							)
							parser.add_argument('''--vocab_size''', default=3_0_5_2_2, type=int)
							__UpperCamelCase  :      Any							      =		parser.parse_args()

							logger.info(F'''Loading data from {args.data_file}''')
							with open(args.data_file, '''rb''') as fp:
														__UpperCamelCase  :      List[Any]							      =		pickle.load(fp)

							logger.info('''Counting occurrences for MLM.''')
							__UpperCamelCase  :      Dict							      =		Counter()
							for tk_ids in data:
														counter.update(tk_ids)
							__UpperCamelCase  :      Optional[Any]							      =		[0] * args.vocab_size
							for k, v in counter.items():
														__UpperCamelCase  :      Dict							      =		v

							logger.info(F'''Dump to {args.token_counts_dump}''')
							with open(args.token_counts_dump, '''wb''') as handle:
														pickle.dump(counts, handle, protocol=pickle.HIGHEST_PROTOCOL)




