


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     =			logging.getLogger(__name__)
def   _UpperCamelCase				(						lowerCAmelCase_  ,						lowerCAmelCase_   )						->str:
   return (preds == labels).mean()






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






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

def   _UpperCamelCase				(						)						->Tuple:
   # 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.

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

   # Set seed
   set_seed(training_args.seed   )

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

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

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

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

   # Initialize our Trainer
   UpperCAmelCase							=     Trainer(
       model=lowerCAmelCase_  ,						args=lowerCAmelCase_  ,						train_dataset=lowerCAmelCase_  ,						eval_dataset=lowerCAmelCase_  ,						compute_metrics=lowerCAmelCase_  ,						data_collator=lowerCAmelCase_  ,						)

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

      UpperCAmelCase							=     trainer.evaluate()

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

            results.update(lowerCAmelCase_   )

   return results
def   _UpperCamelCase				(						lowerCAmelCase_   )						->str:
   # For xla_spawn (TPUs)
   main()


if __name__ == "__main__":
 main()







from pickle import UnpicklingError

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict

from ..utils import logging


__a     =			logging.get_logger(__name__)
def   _UpperCamelCase				(						lowerCAmelCase_  ,						lowerCAmelCase_   )						->Tuple:
   try:
      with open(lowerCAmelCase_  ,						"""rb"""   ) as flax_state_f:
         UpperCAmelCase							=     from_bytes(lowerCAmelCase_  ,						flax_state_f.read()   )
   except UnpicklingError as e:
      try:
         with open(lowerCAmelCase_   ) as f:
            if f.read().startswith("""version"""   ):
               raise OSError(
                   """You seem to have cloned a repository without having git-lfs installed. Please"""
                   """ install git-lfs and run `git lfs install` followed by `git lfs pull` in the"""
                   """ folder you cloned."""   )
            else:
               raise ValueError from e
      except (UnicodeDecodeError, ValueError):
         raise EnvironmentError(F"""Unable to convert {model_file} to Flax deserializable object. """   )

   return load_flax_weights_in_pytorch_model(lowerCAmelCase_  ,						lowerCAmelCase_   )







def   _UpperCamelCase				(						lowerCAmelCase_  ,						lowerCAmelCase_   )						->Dict:

   try:
      import torch  # noqa: F401
   except ImportError:
      logger.error(
          """Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see"""
          """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
          """ instructions."""   )
      raise

   # check if we have bf16 weights
   UpperCAmelCase							=     flatten_dict(jax.tree_util.tree_map(lambda lowerCAmelCase_   : x.dtype == jnp.bfloataa  ,						lowerCAmelCase_   )   ).values()
   if any(lowerCAmelCase_   ):
      # convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16

      # and bf16 is not fully supported in PT yet.
      logger.warning(
          """Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` """
          """before loading those in PyTorch model."""   )
      UpperCAmelCase							=     jax.tree_util.tree_map(
          lambda lowerCAmelCase_   : params.astype(np.floataa   ) if params.dtype == jnp.bfloataa else params  ,						lowerCAmelCase_   )

   UpperCAmelCase							=     """"""

   UpperCAmelCase							=     flatten_dict(lowerCAmelCase_  ,						sep="""."""   )
   UpperCAmelCase							=     pt_model.state_dict()

   # keep track of unexpected & missing keys
   UpperCAmelCase							=     []
   UpperCAmelCase							=     set(pt_model_dict.keys()   )

   for flax_key_tuple, flax_tensor in flax_state_dict.items():
      UpperCAmelCase							=     flax_key_tuple.split("""."""   )

      if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
         UpperCAmelCase							=     flax_key_tuple_array[:-1] + ["""weight"""]
         UpperCAmelCase							=     jnp.transpose(lowerCAmelCase_  ,						(3, 2, 0, 1)   )
      elif flax_key_tuple_array[-1] == "kernel":
         UpperCAmelCase							=     flax_key_tuple_array[:-1] + ["""weight"""]
         UpperCAmelCase							=     flax_tensor.T
      elif flax_key_tuple_array[-1] == "scale":
         UpperCAmelCase							=     flax_key_tuple_array[:-1] + ["""weight"""]

      if "time_embedding" not in flax_key_tuple_array:
         for i, flax_key_tuple_string in enumerate(lowerCAmelCase_   ):
            UpperCAmelCase							=     (
                flax_key_tuple_string.replace("""_0"""  ,						""".0"""   )
                .replace("""_1"""  ,						""".1"""   )
                .replace("""_2"""  ,						""".2"""   )
                .replace("""_3"""  ,						""".3"""   )
                .replace("""_4"""  ,						""".4"""   )
                .replace("""_5"""  ,						""".5"""   )
                .replace("""_6"""  ,						""".6"""   )
                .replace("""_7"""  ,						""".7"""   )
                .replace("""_8"""  ,						""".8"""   )
                .replace("""_9"""  ,						""".9"""   )
            )

      UpperCAmelCase							=     """.""".join(lowerCAmelCase_   )

      if flax_key in pt_model_dict:
         if flax_tensor.shape != pt_model_dict[flax_key].shape:
            raise ValueError(
                F"""Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected """
                F"""to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}."""   )
         else:
            # add weight to pytorch dict
            UpperCAmelCase							=     np.asarray(lowerCAmelCase_   ) if not isinstance(lowerCAmelCase_  ,						np.ndarray   ) else flax_tensor
            UpperCAmelCase							=     torch.from_numpy(lowerCAmelCase_   )
            # remove from missing keys
            missing_keys.remove(lowerCAmelCase_   )
      else:
         # weight is not expected by PyTorch model
         unexpected_keys.append(lowerCAmelCase_   )

   pt_model.load_state_dict(lowerCAmelCase_   )

   # re-transform missing_keys to list
   UpperCAmelCase							=     list(lowerCAmelCase_   )

   if len(lowerCAmelCase_   ) > 0:
      logger.warning(
          """Some weights of the Flax model were not used when initializing the PyTorch model"""
          F""" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing"""
          F""" {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture"""
          """ (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"""
          F""" IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect"""
          """ to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"""
          """ FlaxBertForSequenceClassification model)."""   )
   if len(lowerCAmelCase_   ) > 0:
      logger.warning(
          F"""Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly"""
          F""" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to"""
          """ use it for predictions and inference."""   )

   return pt_model



