





import logging
import os
from dataclasses import dataclass, field
from typing import Dict, Optional

import numpy as np
from utils_multiple_choice import MultipleChoiceDataset, Split, processors

import transformers
from transformers import (
    AutoConfig,
    AutoModelForMultipleChoice,
    AutoTokenizer,
    DataCollatorWithPadding,
    EvalPrediction,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    set_seed,
)
from transformers.trainer_utils import is_main_process


__lowerCAmelCase :List[str]     =    logging.getLogger(__name__)







def    A				(							UpperCAmelCase  ,  UpperCAmelCase				):
					return (preds == labels).mean()






@dataclass
class        _a:
		lowerCamelCase__     :str   				=		field(
		    metadata={'help': 'Path to pretrained model or model identifier from huggingface.co/models'}    )
		lowerCamelCase__     :Tuple   				=		field(
		    default=_UpperCAmelCase	,   metadata={'help': 'Pretrained config name or path if not the same as model_name'}    )
		lowerCamelCase__     :Optional[Any]   				=		field(
		    default=_UpperCAmelCase	,   metadata={'help': 'Pretrained tokenizer name or path if not the same as model_name'}    )
		lowerCamelCase__     :Optional[Any]   				=		field(
		    default=_UpperCAmelCase	,   metadata={'help': 'Where do you want to store the pretrained models downloaded from huggingface.co'}	,   )






@dataclass
class        _a:
		lowerCamelCase__     :Optional[Any]   				=		field(metadata={'help': 'The name of the task to train on: ' + ', '.join(processors.keys()    )}    )
		lowerCamelCase__     :Optional[Any]   				=		field(metadata={'help': 'Should contain the data files for the task.'}    )
		lowerCamelCase__     :Optional[int]   				=		field(
		    default=128	,   metadata={
		        'help': (
		            'The maximum total input sequence length after tokenization. Sequences longer '
		            'than this will be truncated, sequences shorter will be padded.'
		        )
		    }	,   )
		lowerCamelCase__     :List[str]   				=		field(
		    default=_UpperCAmelCase	,   metadata={'help': 'Overwrite the cached training and evaluation sets'}    )







def    A				(							):
					_snake_case       :					str				       =   HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)				)
					_snake_case       :					List[Any]				       =   parser.parse_args_into_dataclasses()

					if (
					    os.path.exists(training_args.output_dir				)
					    and os.listdir(training_args.output_dir				)
					    and training_args.do_train
					    and not training_args.overwrite_output_dir
					):
										raise ValueError(
										    F"""Output directory ({training_args.output_dir}) already exists and is not empty. Use"""
										    " --overwrite_output_dir to overcome."				)

					# Setup logging
					logging.basicConfig(
					    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"  ,  datefmt="%m/%d/%Y %H:%M:%S"  ,  level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN  ,  )
					logger.warning(
					    "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s"  ,  training_args.local_rank  ,  training_args.device  ,  training_args.n_gpu  ,  bool(training_args.local_rank != -1				)  ,  training_args.fpaa  ,  )
					# Set the verbosity to info of the Transformers logger (on main process only):
					if is_main_process(training_args.local_rank				):
										transformers.utils.logging.set_verbosity_info()
										transformers.utils.logging.enable_default_handler()
										transformers.utils.logging.enable_explicit_format()
					logger.info("Training/evaluation parameters %s"  ,  SCREAMING_SNAKE_CASE_				)

					# Set seed
					set_seed(training_args.seed				)

					try:
										_snake_case       :					int				       =   processors[data_args.task_name]()
										_snake_case       :					Union[str, Any]				       =   processor.get_labels()
										_snake_case       :					Union[str, Any]				       =   len(SCREAMING_SNAKE_CASE_				)
					except KeyError:
										raise ValueError("Task not found: %s" % (data_args.task_name)				)

					# Load pretrained model and tokenizer
					#
					# Distributed training:
					# The .from_pretrained methods guarantee that only one local process can concurrently
					# download model & vocab.

					_snake_case       :					List[str]				       =   AutoConfig.from_pretrained(
					    model_args.config_name if model_args.config_name else model_args.model_name_or_path  ,  num_labels=SCREAMING_SNAKE_CASE_  ,  finetuning_task=data_args.task_name  ,  cache_dir=model_args.cache_dir  ,  )
					_snake_case       :					List[str]				       =   AutoTokenizer.from_pretrained(
					    model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path  ,  cache_dir=model_args.cache_dir  ,  )
					_snake_case       :					List[Any]				       =   AutoModelForMultipleChoice.from_pretrained(
					    model_args.model_name_or_path  ,  from_tf=bool(".ckpt" in model_args.model_name_or_path				)  ,  config=SCREAMING_SNAKE_CASE_  ,  cache_dir=model_args.cache_dir  ,  )

					# Get datasets
					_snake_case       :					Dict				       =   (
					    MultipleChoiceDataset(
					        data_dir=data_args.data_dir  ,  tokenizer=SCREAMING_SNAKE_CASE_  ,  task=data_args.task_name  ,  max_seq_length=data_args.max_seq_length  ,  overwrite_cache=data_args.overwrite_cache  ,  mode=Split.train  ,  )
					    if training_args.do_train
					    else None
					)
					_snake_case       :					Union[str, Any]				       =   (
					    MultipleChoiceDataset(
					        data_dir=data_args.data_dir  ,  tokenizer=SCREAMING_SNAKE_CASE_  ,  task=data_args.task_name  ,  max_seq_length=data_args.max_seq_length  ,  overwrite_cache=data_args.overwrite_cache  ,  mode=Split.dev  ,  )
					    if training_args.do_eval
					    else None
					)

					def compute_metrics(UpperCAmelCase				) -> Dict:
										_snake_case       :					int				       =   np.argmax(p.predictions  ,  axis=1				)
										return {"acc": simple_accuracy(SCREAMING_SNAKE_CASE_  ,  p.label_ids				)}

					# Data collator
					_snake_case       :					Optional[Any]				       =   DataCollatorWithPadding(SCREAMING_SNAKE_CASE_  ,  pad_to_multiple_of=8				) if training_args.fpaa else None

					# Initialize our Trainer
					_snake_case       :					str				       =   Trainer(
					    model=SCREAMING_SNAKE_CASE_  ,  args=SCREAMING_SNAKE_CASE_  ,  train_dataset=SCREAMING_SNAKE_CASE_  ,  eval_dataset=SCREAMING_SNAKE_CASE_  ,  compute_metrics=SCREAMING_SNAKE_CASE_  ,  data_collator=SCREAMING_SNAKE_CASE_  ,  )

					# Training
					if training_args.do_train:
										trainer.train(
										    model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path				) else None				)
										trainer.save_model()
										# For convenience, we also re-save the tokenizer to the same directory,
										# so that you can share your model easily on huggingface.co/models =)
										if trainer.is_world_master():
															tokenizer.save_pretrained(training_args.output_dir				)

    # Evaluation
					_snake_case       :					Any				       =   {}
					if training_args.do_eval:
										logger.info("*** Evaluate ***"				)

										_snake_case       :					Any				       =   trainer.evaluate()

										_snake_case       :					Optional[Any]				       =   os.path.join(training_args.output_dir  ,  "eval_results.txt"				)
										if trainer.is_world_master():
															with open(SCREAMING_SNAKE_CASE_  ,  "w"				) as writer:
																				logger.info("***** Eval results *****"				)
																				for key, value in result.items():
																									logger.info("  %s = %s"  ,  SCREAMING_SNAKE_CASE_  ,  SCREAMING_SNAKE_CASE_				)
																									writer.write("%s = %s\n" % (key, value)				)

																				results.update(SCREAMING_SNAKE_CASE_				)

					return results







def    A				(							UpperCAmelCase				):
					main()


if __name__ == "__main__":
							main()






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        _a:







		def __init__(    self ,						__snake_case ,						)      ->      List[str]:


							'''simple docstring'''

							_snake_case       :					Dict				       =   parent
							_snake_case       :					Any				       =   1_3
							_snake_case       :					Union[str, Any]				       =   7
							_snake_case       :					Any				       =   True
							_snake_case       :					Optional[Any]				       =   True
							_snake_case       :					List[str]				       =   False
							_snake_case       :					str				       =   True
							_snake_case       :					int				       =   9_9
							_snake_case       :					Optional[int]				       =   3_2
							_snake_case       :					Dict				       =   2
							_snake_case       :					Optional[Any]				       =   4
							_snake_case       :					Any				       =   3_7
							_snake_case       :					Optional[int]				       =   "gelu"
							_snake_case       :					Optional[Any]				       =   0.1
							_snake_case       :					List[Any]				       =   0.1
							_snake_case       :					Any				       =   5_1_2
							_snake_case       :					Optional[int]				       =   1_6
							_snake_case       :					Dict				       =   2
							_snake_case       :					Optional[int]				       =   0.02
							_snake_case       :					Any				       =   3
							_snake_case       :					Tuple				       =   4
							_snake_case       :					str				       =   None







		def        lowercase       (    self				)      ->      Union[str, Any]:


							'''simple docstring'''

							_snake_case       :					Optional[int]				       =   ids_tensor([self.batch_size, self.seq_length] ,						self.vocab_size				)

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

							_snake_case       :					str				       =   None
							_snake_case       :					Optional[Any]				       =   None
							_snake_case       :					Any				       =   None
							if self.use_labels:
												_snake_case       :					Union[str, 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       :					List[Any]				       =   ids_tensor([self.batch_size] ,						self.num_choices				)

							_snake_case       :					Optional[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 ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case				)      ->      Tuple:


							'''simple docstring'''

							_snake_case       :					str				       =   TFDistilBertModel(config=__snake_case				)
							_snake_case       :					Dict				       =   {"input_ids": input_ids, "attention_mask": input_mask}

							_snake_case       :					Union[str, Any]				       =   model(__snake_case				)

							_snake_case       :					Union[str, Any]				       =   [input_ids, input_mask]

							_snake_case       :					int				       =   model(__snake_case				)

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







		def        lowercase       (    self ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case				)      ->      str:


							'''simple docstring'''

							_snake_case       :					str				       =   TFDistilBertForMaskedLM(config=__snake_case				)
							_snake_case       :					Optional[Any]				       =   {"input_ids": input_ids, "attention_mask": input_mask}
							_snake_case       :					Tuple				       =   model(__snake_case				)
							self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.vocab_size)				)







		def        lowercase       (    self ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case				)      ->      Union[str, Any]:


							'''simple docstring'''

							_snake_case       :					int				       =   TFDistilBertForQuestionAnswering(config=__snake_case				)
							_snake_case       :					Any				       =   {
							    "input_ids": input_ids,
							    "attention_mask": input_mask,
							}
							_snake_case       :					Tuple				       =   model(__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        lowercase       (    self ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case				)      ->      str:


							'''simple docstring'''

							_snake_case       :					Dict				       =   self.num_labels
							_snake_case       :					Dict				       =   TFDistilBertForSequenceClassification(__snake_case				)
							_snake_case       :					Dict				       =   {"input_ids": input_ids, "attention_mask": input_mask}
							_snake_case       :					Optional[Any]				       =   model(__snake_case				)
							self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.num_labels)				)







		def        lowercase       (    self ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case				)      ->      Dict:


							'''simple docstring'''

							_snake_case       :					str				       =   self.num_choices
							_snake_case       :					List[str]				       =   TFDistilBertForMultipleChoice(__snake_case				)
							_snake_case       :					Optional[int]				       =   tf.tile(tf.expand_dims(__snake_case ,						1				) ,						(1, self.num_choices, 1)				)
							_snake_case       :					int				       =   tf.tile(tf.expand_dims(__snake_case ,						1				) ,						(1, self.num_choices, 1)				)
							_snake_case       :					str				       =   {
							    "input_ids": multiple_choice_inputs_ids,
							    "attention_mask": multiple_choice_input_mask,
							}
							_snake_case       :					List[Any]				       =   model(__snake_case				)
							self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.num_choices)				)







		def        lowercase       (    self ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case ,						__snake_case				)      ->      Dict:


							'''simple docstring'''

							_snake_case       :					List[str]				       =   self.num_labels
							_snake_case       :					int				       =   TFDistilBertForTokenClassification(__snake_case				)
							_snake_case       :					Any				       =   {"input_ids": input_ids, "attention_mask": input_mask}
							_snake_case       :					Any				       =   model(__snake_case				)
							self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.num_labels)				)







		def        lowercase       (    self				)      ->      Tuple:


							'''simple docstring'''

							_snake_case       :					List[str]				       =   self.prepare_config_and_inputs()
							((_snake_case)       , (_snake_case)       , (_snake_case)       , (_snake_case)       , (_snake_case)       , (_snake_case))       :					List[str]				       =   config_and_inputs
							_snake_case       :					Optional[int]				       =   {"input_ids": input_ids, "attention_mask": input_mask}
							return config, inputs_dict






@require_tf
class        _a(    __A	,   __A	,   unittest.TestCase    ):
		lowerCamelCase__     :str   				=		(
		    (
		        TFDistilBertModel,
		        TFDistilBertForMaskedLM,
		        TFDistilBertForQuestionAnswering,
		        TFDistilBertForSequenceClassification,
		        TFDistilBertForTokenClassification,
		        TFDistilBertForMultipleChoice,
		    )
		    if is_tf_available()
		    else None
		)
		lowerCamelCase__     :Tuple   				=		(
		    {
		        'feature-extraction': TFDistilBertModel,
		        'fill-mask': TFDistilBertForMaskedLM,
		        'question-answering': TFDistilBertForQuestionAnswering,
		        'text-classification': TFDistilBertForSequenceClassification,
		        'token-classification': TFDistilBertForTokenClassification,
		        'zero-shot': TFDistilBertForSequenceClassification,
		    }
		    if is_tf_available()
		    else {}
		)
		lowerCamelCase__     :Tuple   				=		False
		lowerCamelCase__     :Any   				=		False







		def        lowercase       (    self				)      ->      int:


							'''simple docstring'''

							_snake_case       :					Optional[Any]				       =   TFDistilBertModelTester(self				)
							_snake_case       :					str				       =   ConfigTester(self ,						config_class=__snake_case ,						dim=3_7				)







		def        lowercase       (    self				)      ->      List[Any]:


							'''simple docstring'''

							self.config_tester.run_common_tests()







		def        lowercase       (    self				)      ->      List[Any]:


							'''simple docstring'''

							_snake_case       :					Dict				       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_model(*__snake_case				)







		def        lowercase       (    self				)      ->      Any:


							'''simple docstring'''

							_snake_case       :					Any				       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_masked_lm(*__snake_case				)







		def        lowercase       (    self				)      ->      List[str]:


							'''simple docstring'''

							_snake_case       :					str				       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_question_answering(*__snake_case				)







		def        lowercase       (    self				)      ->      Dict:


							'''simple docstring'''

							_snake_case       :					List[Any]				       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_sequence_classification(*__snake_case				)







		def        lowercase       (    self				)      ->      Dict:


							'''simple docstring'''

							_snake_case       :					List[str]				       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_multiple_choice(*__snake_case				)







		def        lowercase       (    self				)      ->      List[Any]:


							'''simple docstring'''

							_snake_case       :					Dict				       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_distilbert_for_token_classification(*__snake_case				)







		@slow
		def        lowercase       (    self				)      ->      Tuple:


							'''simple docstring'''

							for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]				):
												_snake_case       :					Any				       =   TFDistilBertModel.from_pretrained(__snake_case				)
												self.assertIsNotNone(__snake_case				)






@require_tf
class        _a(    unittest.TestCase    ):







		@slow
		def        lowercase       (    self				)      ->      List[str]:


							'''simple docstring'''

							_snake_case       :					Any				       =   TFDistilBertModel.from_pretrained("distilbert-base-uncased"				)
							_snake_case       :					int				       =   tf.constant([[0, 1, 2, 3, 4, 5]]				)
							_snake_case       :					Dict				       =   model(__snake_case				)[0]

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

							_snake_case       :					Optional[Any]				       =   tf.constant(
							    [
							        [
							            [0.19_26_18_85, -0.13_73_29_55, 0.4_11_97_99],
							            [0.22_15_01_56, -0.07_42_26_61, 0.39_03_72_04],
							            [0.22_75_60_18, -0.0_89_64_14, 0.3_70_14_67],
							        ]
							    ]				)
							tf.debugging.assert_near(output[:, :3, :3] ,						__snake_case ,						atol=1E-4				)