






"""simple docstring"""





import unittest

from transformers import AutoTokenizer, NystromformerConfig, is_torch_available
from transformers.testing_utils import require_torch, 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 (
      NystromformerForMaskedLM,
      NystromformerForMultipleChoice,
      NystromformerForQuestionAnswering,
      NystromformerForSequenceClassification,
      NystromformerForTokenClassification,
      NystromformerModel,
  )
  from transformers.models.nystromformer.modeling_nystromformer import NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST






class      a						:






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






    def    UpperCamelCase       (       self     :							List[str]					)     ->		Tuple:
      lowerCamelCase_         =						ids_tensor([self.batch_size, self.seq_length]       ,						self.vocab_size					)

      lowerCamelCase_         =						None
      if self.use_input_mask:
        lowerCamelCase_         =						random_attention_mask([self.batch_size, self.seq_length]					)

      lowerCamelCase_         =						None
      if self.use_token_type_ids:
        lowerCamelCase_         =						ids_tensor([self.batch_size, self.seq_length]       ,						self.type_vocab_size					)

      lowerCamelCase_         =						None
      lowerCamelCase_         =						None
      lowerCamelCase_         =						None
      if self.use_labels:
        lowerCamelCase_         =						ids_tensor([self.batch_size]       ,						self.type_sequence_label_size					)
        lowerCamelCase_         =						ids_tensor([self.batch_size, self.seq_length]       ,						self.num_labels					)
        lowerCamelCase_         =						ids_tensor([self.batch_size]       ,						self.num_choices					)

      lowerCamelCase_         =						self.get_config()

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






    def    UpperCamelCase       (       self     :							List[Any]					)     ->		Union[str, Any]:
      return NystromformerConfig(
          vocab_size=self.vocab_size       ,						hidden_size=self.hidden_size       ,						num_hidden_layers=self.num_hidden_layers       ,						num_attention_heads=self.num_attention_heads       ,						intermediate_size=self.intermediate_size       ,						hidden_act=self.hidden_act       ,						hidden_dropout_prob=self.hidden_dropout_prob       ,						attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,						max_position_embeddings=self.max_position_embeddings       ,						type_vocab_size=self.type_vocab_size       ,						is_decoder=__SCREAMING_SNAKE_CASE       ,						initializer_range=self.initializer_range       ,						)






    def    UpperCamelCase       (       self     :							Any       ,						__SCREAMING_SNAKE_CASE     :							List[Any]       ,						__SCREAMING_SNAKE_CASE     :							str       ,						__SCREAMING_SNAKE_CASE     :							int       ,						__SCREAMING_SNAKE_CASE     :							Tuple       ,						__SCREAMING_SNAKE_CASE     :							List[str]       ,						__SCREAMING_SNAKE_CASE     :							Any       ,						__SCREAMING_SNAKE_CASE     :							Any					)     ->		str:
      lowerCamelCase_         =						NystromformerModel(config=__SCREAMING_SNAKE_CASE					)
      model.to(__SCREAMING_SNAKE_CASE					)
      model.eval()
      lowerCamelCase_         =						model(__SCREAMING_SNAKE_CASE       ,						attention_mask=__SCREAMING_SNAKE_CASE       ,						token_type_ids=__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						model(__SCREAMING_SNAKE_CASE       ,						token_type_ids=__SCREAMING_SNAKE_CASE					)
      lowerCamelCase_         =						model(__SCREAMING_SNAKE_CASE					)
      self.parent.assertEqual(result.last_hidden_state.shape       ,						(self.batch_size, self.seq_length, self.hidden_size)					)






    def    UpperCamelCase       (       self     :							List[Any]       ,						__SCREAMING_SNAKE_CASE     :							Optional[Any]       ,						__SCREAMING_SNAKE_CASE     :							Optional[int]       ,						__SCREAMING_SNAKE_CASE     :							Union[str, Any]       ,						__SCREAMING_SNAKE_CASE     :							List[Any]       ,						__SCREAMING_SNAKE_CASE     :							str       ,						__SCREAMING_SNAKE_CASE     :							Optional[Any]       ,						__SCREAMING_SNAKE_CASE     :							List[str]					)     ->		Dict:
      lowerCamelCase_         =						NystromformerForMaskedLM(config=__SCREAMING_SNAKE_CASE					)
      model.to(__SCREAMING_SNAKE_CASE					)
      model.eval()
      lowerCamelCase_         =						model(__SCREAMING_SNAKE_CASE       ,						attention_mask=__SCREAMING_SNAKE_CASE       ,						token_type_ids=__SCREAMING_SNAKE_CASE       ,						labels=__SCREAMING_SNAKE_CASE					)
      self.parent.assertEqual(result.logits.shape       ,						(self.batch_size, self.seq_length, self.vocab_size)					)






    def    UpperCamelCase       (       self     :							Tuple       ,						__SCREAMING_SNAKE_CASE     :							Any       ,						__SCREAMING_SNAKE_CASE     :							Tuple       ,						__SCREAMING_SNAKE_CASE     :							Union[str, Any]       ,						__SCREAMING_SNAKE_CASE     :							List[str]       ,						__SCREAMING_SNAKE_CASE     :							Optional[int]       ,						__SCREAMING_SNAKE_CASE     :							List[str]       ,						__SCREAMING_SNAKE_CASE     :							List[Any]					)     ->		str:
      lowerCamelCase_         =						NystromformerForQuestionAnswering(config=__SCREAMING_SNAKE_CASE					)
      model.to(__SCREAMING_SNAKE_CASE					)
      model.eval()
      lowerCamelCase_         =						model(
          __SCREAMING_SNAKE_CASE       ,						attention_mask=__SCREAMING_SNAKE_CASE       ,						token_type_ids=__SCREAMING_SNAKE_CASE       ,						start_positions=__SCREAMING_SNAKE_CASE       ,						end_positions=__SCREAMING_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    UpperCamelCase       (       self     :							Any       ,						__SCREAMING_SNAKE_CASE     :							List[str]       ,						__SCREAMING_SNAKE_CASE     :							Dict       ,						__SCREAMING_SNAKE_CASE     :							int       ,						__SCREAMING_SNAKE_CASE     :							Union[str, Any]       ,						__SCREAMING_SNAKE_CASE     :							List[Any]       ,						__SCREAMING_SNAKE_CASE     :							Optional[Any]       ,						__SCREAMING_SNAKE_CASE     :							Union[str, Any]					)     ->		int:
      lowerCamelCase_         =						self.num_labels
      lowerCamelCase_         =						NystromformerForSequenceClassification(__SCREAMING_SNAKE_CASE					)
      model.to(__SCREAMING_SNAKE_CASE					)
      model.eval()
      lowerCamelCase_         =						model(__SCREAMING_SNAKE_CASE       ,						attention_mask=__SCREAMING_SNAKE_CASE       ,						token_type_ids=__SCREAMING_SNAKE_CASE       ,						labels=__SCREAMING_SNAKE_CASE					)
      self.parent.assertEqual(result.logits.shape       ,						(self.batch_size, self.num_labels)					)






    def    UpperCamelCase       (       self     :							Union[str, Any]       ,						__SCREAMING_SNAKE_CASE     :							int       ,						__SCREAMING_SNAKE_CASE     :							Optional[int]       ,						__SCREAMING_SNAKE_CASE     :							Union[str, Any]       ,						__SCREAMING_SNAKE_CASE     :							Any       ,						__SCREAMING_SNAKE_CASE     :							Tuple       ,						__SCREAMING_SNAKE_CASE     :							Union[str, Any]       ,						__SCREAMING_SNAKE_CASE     :							Any					)     ->		int:
      lowerCamelCase_         =						self.num_labels
      lowerCamelCase_         =						NystromformerForTokenClassification(config=__SCREAMING_SNAKE_CASE					)
      model.to(__SCREAMING_SNAKE_CASE					)
      model.eval()
      lowerCamelCase_         =						model(__SCREAMING_SNAKE_CASE       ,						attention_mask=__SCREAMING_SNAKE_CASE       ,						token_type_ids=__SCREAMING_SNAKE_CASE       ,						labels=__SCREAMING_SNAKE_CASE					)
      self.parent.assertEqual(result.logits.shape       ,						(self.batch_size, self.seq_length, self.num_labels)					)






    def    UpperCamelCase       (       self     :							List[str]       ,						__SCREAMING_SNAKE_CASE     :							str       ,						__SCREAMING_SNAKE_CASE     :							Optional[Any]       ,						__SCREAMING_SNAKE_CASE     :							Tuple       ,						__SCREAMING_SNAKE_CASE     :							Optional[Any]       ,						__SCREAMING_SNAKE_CASE     :							int       ,						__SCREAMING_SNAKE_CASE     :							Optional[Any]       ,						__SCREAMING_SNAKE_CASE     :							Optional[int]					)     ->		Union[str, Any]:
      lowerCamelCase_         =						self.num_choices
      lowerCamelCase_         =						NystromformerForMultipleChoice(config=__SCREAMING_SNAKE_CASE					)
      model.to(__SCREAMING_SNAKE_CASE					)
      model.eval()
      lowerCamelCase_         =						input_ids.unsqueeze(1					).expand(-1       ,						self.num_choices       ,						-1					).contiguous()
      lowerCamelCase_         =						token_type_ids.unsqueeze(1					).expand(-1       ,						self.num_choices       ,						-1					).contiguous()
      lowerCamelCase_         =						input_mask.unsqueeze(1					).expand(-1       ,						self.num_choices       ,						-1					).contiguous()
      lowerCamelCase_         =						model(
          __SCREAMING_SNAKE_CASE       ,						attention_mask=__SCREAMING_SNAKE_CASE       ,						token_type_ids=__SCREAMING_SNAKE_CASE       ,						labels=__SCREAMING_SNAKE_CASE       ,						)
      self.parent.assertEqual(result.logits.shape       ,						(self.batch_size, self.num_choices)					)






    def    UpperCamelCase       (       self     :							str					)     ->		Optional[Any]:
      lowerCamelCase_         =						self.prepare_config_and_inputs()
      (
          (
          lowerCamelCase_       
      )     ,	(
          lowerCamelCase_       
      )     ,	(
          lowerCamelCase_       
      )     ,	(
          lowerCamelCase_       
      )     ,	(
          lowerCamelCase_       
      )     ,	(
          lowerCamelCase_       
      )     ,	(
          lowerCamelCase_       
      )     ,	       
      )         =						config_and_inputs
      lowerCamelCase_         =						{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
      return config, inputs_dict






@require_torch
class      a						(					SCREAMING_SNAKE_CASE__  ,      SCREAMING_SNAKE_CASE__  ,      unittest.TestCase		):
    SCREAMING_SNAKE_CASE						:     List[Any]						= (
        (
            NystromformerModel,
            NystromformerForMaskedLM,
            NystromformerForMultipleChoice,
            NystromformerForQuestionAnswering,
            NystromformerForSequenceClassification,
            NystromformerForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    SCREAMING_SNAKE_CASE						:     List[Any]						= (
        {
            """feature-extraction""": NystromformerModel,
            """fill-mask""": NystromformerForMaskedLM,
            """question-answering""": NystromformerForQuestionAnswering,
            """text-classification""": NystromformerForSequenceClassification,
            """token-classification""": NystromformerForTokenClassification,
            """zero-shot""": NystromformerForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    SCREAMING_SNAKE_CASE						:     str						= False
    SCREAMING_SNAKE_CASE						:     Any						= False






    def    UpperCamelCase       (       self     :							List[Any]					)     ->		Optional[Any]:
      lowerCamelCase_         =						NystromformerModelTester(self					)
      lowerCamelCase_         =						ConfigTester(self       ,						config_class=__SCREAMING_SNAKE_CASE       ,						hidden_size=37					)






    def    UpperCamelCase       (       self     :							Optional[Any]					)     ->		List[Any]:
      self.config_tester.run_common_tests()






    def    UpperCamelCase       (       self     :							List[str]					)     ->		str:
      lowerCamelCase_         =						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE					)






    def    UpperCamelCase       (       self     :							Union[str, Any]					)     ->		int:
      lowerCamelCase_         =						self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
        lowerCamelCase_         =						type
        self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE					)






    def    UpperCamelCase       (       self     :							str					)     ->		Dict:
      lowerCamelCase_         =						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*__SCREAMING_SNAKE_CASE					)






    def    UpperCamelCase       (       self     :							List[str]					)     ->		Optional[Any]:
      lowerCamelCase_         =						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_multiple_choice(*__SCREAMING_SNAKE_CASE					)






    def    UpperCamelCase       (       self     :							Any					)     ->		Union[str, Any]:
      lowerCamelCase_         =						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_question_answering(*__SCREAMING_SNAKE_CASE					)






    def    UpperCamelCase       (       self     :							Any					)     ->		Optional[Any]:
      lowerCamelCase_         =						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_sequence_classification(*__SCREAMING_SNAKE_CASE					)






    def    UpperCamelCase       (       self     :							Union[str, Any]					)     ->		List[Any]:
      lowerCamelCase_         =						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*__SCREAMING_SNAKE_CASE					)






    @slow
    def    UpperCamelCase       (       self     :							Tuple					)     ->		Tuple:
      for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        lowerCamelCase_         =						NystromformerModel.from_pretrained(__SCREAMING_SNAKE_CASE					)
        self.assertIsNotNone(__SCREAMING_SNAKE_CASE					)






@require_torch
class      a						(					unittest.TestCase		):






    @slow
    def    UpperCamelCase       (       self     :							Dict					)     ->		Union[str, Any]:
      lowerCamelCase_         =						NystromformerModel.from_pretrained('uw-madison/nystromformer-512'					)
      lowerCamelCase_         =						torch.tensor([[0, 1, 2, 3, 4, 5]]					)

      with torch.no_grad():
        lowerCamelCase_         =						model(__SCREAMING_SNAKE_CASE					)[0]

      lowerCamelCase_         =						torch.Size((1, 6, 768)					)
      self.assertEqual(output.shape       ,						__SCREAMING_SNAKE_CASE					)

      lowerCamelCase_         =						torch.tensor(
          [[[-0.4_532, -0.0_936, 0.5_137], [-0.2_676, 0.0_628, 0.6_186], [-0.3_629, -0.1_726, 0.4_716]]]					)

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






    @slow
    def    UpperCamelCase       (       self     :							Tuple					)     ->		int:
      lowerCamelCase_         =						'the [MASK] of Belgium is Brussels'

      lowerCamelCase_         =						AutoTokenizer.from_pretrained('uw-madison/nystromformer-512'					)
      lowerCamelCase_         =						NystromformerForMaskedLM.from_pretrained('uw-madison/nystromformer-512'					)

      lowerCamelCase_         =						tokenizer(__SCREAMING_SNAKE_CASE       ,						return_tensors='pt'					)

      with torch.no_grad():
        lowerCamelCase_         =						model(encoding.input_ids					).logits

      lowerCamelCase_         =						token_logits[:, 2, :].argmax(-1					)[0]

      self.assertEqual(tokenizer.decode(__SCREAMING_SNAKE_CASE					)       ,						'capital'					)





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


A   :     Any       =				logging.getLogger(__name__)


def    a__       (						__UpperCamelCase   ,					__UpperCamelCase ):
      return (preds == labels).mean()


@dataclass
class      lowerCamelCase    :
  """simple docstring"""







  lowerCamelCase__   =  field(
      metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models'''}    )
  lowerCamelCase__   =  field(
      default=SCREAMING_SNAKE_CASE__     ,					metadata={'''help''': '''Pretrained config name or path if not the same as model_name'''}    )
  lowerCamelCase__   =  field(
      default=SCREAMING_SNAKE_CASE__     ,					metadata={'''help''': '''Pretrained tokenizer name or path if not the same as model_name'''}    )
  lowerCamelCase__   =  field(
      default=SCREAMING_SNAKE_CASE__     ,					metadata={'''help''': '''Where do you want to store the pretrained models downloaded from huggingface.co'''}     ,					)



@dataclass
class      lowerCamelCase    :
  """simple docstring"""







  lowerCamelCase__   =  field(metadata={'''help''': '''The name of the task to train on: ''' + ''', '''.join(processors.keys()    )}    )
  lowerCamelCase__   =  field(metadata={'''help''': '''Should contain the data files for the task.'''}    )
  lowerCamelCase__   =  field(
      default=1_2_8     ,					metadata={
          '''help''': (
              '''The maximum total input sequence length after tokenization. Sequences longer '''
              '''than this will be truncated, sequences shorter will be padded.'''
          )
      }     ,					)
  lowerCamelCase__   =  field(
      default=SCREAMING_SNAKE_CASE__     ,					metadata={'''help''': '''Overwrite the cached training and evaluation sets'''}    )



def    a__       (						):
      # See all possible arguments in src/transformers/training_args.py
      # or by passing the --help flag to this script.
      # We now keep distinct sets of args, for a cleaner separation of concerns.

      SCREAMING_SNAKE_CASE_      =     HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments) )
      SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_      =     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"   ,					__UpperCamelCase )

      # Set seed
      set_seed(training_args.seed )

      try:
            SCREAMING_SNAKE_CASE_      =     processors[data_args.task_name]()
            SCREAMING_SNAKE_CASE_      =     processor.get_labels()
            SCREAMING_SNAKE_CASE_      =     len(__UpperCamelCase )
      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.

      SCREAMING_SNAKE_CASE_      =     AutoConfig.from_pretrained(
          model_args.config_name if model_args.config_name else model_args.model_name_or_path   ,					num_labels=__UpperCamelCase   ,					finetuning_task=data_args.task_name   ,					cache_dir=model_args.cache_dir   ,					)
      SCREAMING_SNAKE_CASE_      =     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   ,					)
      SCREAMING_SNAKE_CASE_      =     AutoModelForMultipleChoice.from_pretrained(
          model_args.model_name_or_path   ,					from_tf=bool(".ckpt" in model_args.model_name_or_path )   ,					config=__UpperCamelCase   ,					cache_dir=model_args.cache_dir   ,					)

      # Get datasets
      SCREAMING_SNAKE_CASE_      =     (
          MultipleChoiceDataset(
              data_dir=data_args.data_dir   ,					tokenizer=__UpperCamelCase   ,					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
      )
      SCREAMING_SNAKE_CASE_      =     (
          MultipleChoiceDataset(
              data_dir=data_args.data_dir   ,					tokenizer=__UpperCamelCase   ,					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:
            SCREAMING_SNAKE_CASE_      =     np.argmax(p.predictions   ,					axis=1 )
            return {"acc": simple_accuracy(__UpperCamelCase   ,					p.label_ids )}

      # Data collator
      SCREAMING_SNAKE_CASE_      =     DataCollatorWithPadding(__UpperCamelCase   ,					pad_to_multiple_of=8 ) if training_args.fpaa else None

      # Initialize our Trainer
      SCREAMING_SNAKE_CASE_      =     Trainer(
          model=__UpperCamelCase   ,					args=__UpperCamelCase   ,					train_dataset=__UpperCamelCase   ,					eval_dataset=__UpperCamelCase   ,					compute_metrics=__UpperCamelCase   ,					data_collator=__UpperCamelCase   ,					)

      # 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
      SCREAMING_SNAKE_CASE_      =     {}
      if training_args.do_eval:
            logger.info("*** Evaluate ***" )

            SCREAMING_SNAKE_CASE_      =     trainer.evaluate()

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

                        results.update(__UpperCamelCase )

      return results


def    a__       (						__UpperCamelCase ):
      # For xla_spawn (TPUs)
      main()


if __name__ == "__main__":
      main()


