



"""simple docstring"""

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


lowercase__					:						Optional[Any]												=			logging.getLogger(__name__)




def __lowercase    (							_a	,    _a					):
    return (preds == labels).mean()




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





@dataclass
class  _UpperCAmelCase :
 _lowerCAmelCase     :    str	      =	field(metadata={"""help""": """The name of the task to train on: """ + """, """.join(processors.keys())})
 _lowerCAmelCase     :    str	      =	field(metadata={"""help""": """Should contain the data files for the task."""})
 _lowerCAmelCase     :    int	      =	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     :    bool	      =	field(
     default=lowerCAmelCase__					,  metadata={"""help""": """Overwrite the cached training and evaluation sets"""})





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

    snake_case_				:  Tuple       =			HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)					)
    snake_case_				:  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'''	,    _a					)

    # Set seed
    set_seed(training_args.seed					)

    try:
        snake_case_				:  Optional[Any]       =			processors[data_args.task_name]()
        snake_case_				:  str       =			processor.get_labels()
        snake_case_				:  Any       =			len(_a					)
    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_				:  Tuple       =			AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path	,    num_labels=_a	,    finetuning_task=data_args.task_name	,    cache_dir=model_args.cache_dir	,    )
    snake_case_				:  Optional[Any]       =			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_				:  Dict       =			AutoModelForMultipleChoice.from_pretrained(
        model_args.model_name_or_path	,    from_tf=bool('''.ckpt''' in model_args.model_name_or_path					)	,    config=_a	,    cache_dir=model_args.cache_dir	,    )

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

    # Data collator
    snake_case_				:  Any       =			DataCollatorWithPadding(_a	,    pad_to_multiple_of=8					) if training_args.fpaa else None

    # Initialize our Trainer
    snake_case_				:  int       =			Trainer(
        model=_a	,    args=_a	,    train_dataset=_a	,    eval_dataset=_a	,    compute_metrics=_a	,    data_collator=_a	,    )

    # 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_				:  Optional[Any]       =			{}
    if training_args.do_eval:
        logger.info('''*** Evaluate ***'''					)

        snake_case_				:  Optional[Any]       =			trainer.evaluate()

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

                results.update(_a					)

    return results




def __lowercase    (							_a					):
    # For xla_spawn (TPUs)
    main()


if __name__ == "__main__":
     main()










"""simple docstring"""

import tensorflow as tf

from ...tf_utils import shape_list




class  _UpperCAmelCase (    tf.keras.layers.Layer):

 def __init__(  self	: int						,  lowercase_	: Tuple						,  lowercase_	: Union[str, Any]						,  lowercase_	: Tuple						,  lowercase_	: List[Any]						,  lowercase_	: Tuple=1						,  lowercase_	: List[str]=False						,  **lowercase_	: Optional[Any]       ):
     super().__init__(**lowercase_       )

     snake_case_				:  int       =			vocab_size
     snake_case_				:  Union[str, Any]       =			d_embed
     snake_case_				:  Optional[int]       =			d_proj

     snake_case_				:  int       =			cutoffs + [vocab_size]
     snake_case_				:  Optional[int]       =			[0] + self.cutoffs
     snake_case_				:  List[str]       =			div_val

     snake_case_				:  int       =			self.cutoffs[0]
     snake_case_				:  Optional[Any]       =			len(self.cutoffs       ) - 1
     snake_case_				:  Any       =			self.shortlist_size + self.n_clusters
     snake_case_				:  Dict       =			keep_order

     snake_case_				:  Tuple       =			[]
     snake_case_				:  Optional[Any]       =			[]

 def 		_snake_case      (  self	: Dict						,  lowercase_	: int       ):
     if self.n_clusters > 0:
         snake_case_				:  Tuple       =			self.add_weight(
             shape=(self.n_clusters, self.d_embed)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name='''cluster_weight'''       )
         snake_case_				:  List[str]       =			self.add_weight(
             shape=(self.n_clusters,)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name='''cluster_bias'''       )

     if self.div_val == 1:
         for i in range(len(self.cutoffs       )       ):
             if self.d_proj != self.d_embed:
                 snake_case_				:  List[Any]       =			self.add_weight(
                     shape=(self.d_embed, self.d_proj)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name=f"out_projs_._{i}"						,  )
                 self.out_projs.append(lowercase_       )
             else:
                 self.out_projs.append(lowercase_       )
             snake_case_				:  Any       =			self.add_weight(
                 shape=(self.vocab_size, self.d_embed)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name=f"out_layers_._{i}_._weight"						,  )
             snake_case_				:  Any       =			self.add_weight(
                 shape=(self.vocab_size,)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name=f"out_layers_._{i}_._bias"						,  )
             self.out_layers.append((weight, bias)       )
     else:
         for i in range(len(self.cutoffs       )       ):
             snake_case_,			snake_case_				:  str       =			self.cutoff_ends[i], self.cutoff_ends[i + 1]
             snake_case_				:  List[str]       =			self.d_embed // (self.div_val**i)

             snake_case_				:  Optional[Any]       =			self.add_weight(
                 shape=(d_emb_i, self.d_proj)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name=f"out_projs_._{i}"       )
             self.out_projs.append(lowercase_       )
             snake_case_				:  Any       =			self.add_weight(
                 shape=(r_idx - l_idx, d_emb_i)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name=f"out_layers_._{i}_._weight"						,  )
             snake_case_				:  List[Any]       =			self.add_weight(
                 shape=(r_idx - l_idx,)						,  initializer='''zeros'''						,  trainable=lowercase_						,  name=f"out_layers_._{i}_._bias"						,  )
             self.out_layers.append((weight, bias)       )
     super().build(lowercase_       )

 @staticmethod
 def 		_snake_case      (  lowercase_	: Dict						,  lowercase_	: Union[str, Any]						,  lowercase_	: Optional[int]						,  lowercase_	: str=None       ):
     snake_case_				:  str       =			x
     if proj is not None:
         snake_case_				:  str       =			tf.einsum('''ibd,ed->ibe'''						,  lowercase_						,  lowercase_       )
     return tf.einsum('''ibd,nd->ibn'''						,  lowercase_						,  lowercase_       ) + b

 @staticmethod
 def 		_snake_case      (  lowercase_	: Tuple						,  lowercase_	: Union[str, Any]       ):
     snake_case_				:  Optional[int]       =			shape_list(lowercase_       )
     snake_case_				:  Any       =			tf.range(lp_size[0]						,  dtype=target.dtype       )
     snake_case_				:  Optional[int]       =			tf.stack([r, target]						,  1       )
     return tf.gather_nd(lowercase_						,  lowercase_       )






 def 		_snake_case      (  self	: Tuple						,  lowercase_	: Union[str, Any]						,  lowercase_	: int						,  lowercase_	: Tuple=True						,  lowercase_	: Tuple=False       ):
     snake_case_				:  Optional[int]       =			0
     if self.n_clusters == 0:
         snake_case_				:  List[str]       =			self._logit(lowercase_						,  self.out_layers[0][0]						,  self.out_layers[0][1]						,  self.out_projs[0]       )
         if target is not None:
             snake_case_				:  List[Any]       =			tf.nn.sparse_softmax_cross_entropy_with_logits(labels=lowercase_						,  logits=lowercase_       )
         snake_case_				:  List[Any]       =			tf.nn.log_softmax(lowercase_						,  axis=-1       )
     else:
         snake_case_				:  Optional[int]       =			shape_list(lowercase_       )
         snake_case_				:  Any       =			[]
         snake_case_				:  List[str]       =			tf.zeros(hidden_sizes[:2]       )
         for i in range(len(self.cutoffs       )       ):
             snake_case_,			snake_case_				:  Union[str, Any]       =			self.cutoff_ends[i], self.cutoff_ends[i + 1]
             if target is not None:
                 snake_case_				:  Optional[Any]       =			(target >= l_idx) & (target < r_idx)
                 snake_case_				:  Any       =			tf.where(lowercase_       )
                 snake_case_				:  List[Any]       =			tf.boolean_mask(lowercase_						,  lowercase_       ) - l_idx

             if self.div_val == 1:
                 snake_case_				:  Optional[int]       =			self.out_layers[0][0][l_idx:r_idx]
                 snake_case_				:  List[Any]       =			self.out_layers[0][1][l_idx:r_idx]
             else:
                 snake_case_				:  Dict       =			self.out_layers[i][0]
                 snake_case_				:  str       =			self.out_layers[i][1]

             if i == 0:
                 snake_case_				:  Tuple       =			tf.concat([cur_W, self.cluster_weight]						,  0       )
                 snake_case_				:  List[str]       =			tf.concat([cur_b, self.cluster_bias]						,  0       )

                 snake_case_				:  List[Any]       =			self._logit(lowercase_						,  lowercase_						,  lowercase_						,  self.out_projs[0]       )
                 snake_case_				:  Union[str, Any]       =			tf.nn.log_softmax(lowercase_       )
                 out.append(head_logprob[..., : self.cutoffs[0]]       )
                 if target is not None:
                     snake_case_				:  int       =			tf.boolean_mask(lowercase_						,  lowercase_       )
                     snake_case_				:  str       =			self._gather_logprob(lowercase_						,  lowercase_       )
             else:
                 snake_case_				:  Union[str, Any]       =			self._logit(lowercase_						,  lowercase_						,  lowercase_						,  self.out_projs[i]       )
                 snake_case_				:  List[str]       =			tf.nn.log_softmax(lowercase_       )
                 snake_case_				:  Dict       =			self.cutoffs[0] + i - 1  # No probability for the head cluster
                 snake_case_				:  Union[str, Any]       =			head_logprob[..., cluster_prob_idx, None] + tail_logprob
                 out.append(lowercase_       )
                 if target is not None:
                     snake_case_				:  Optional[int]       =			tf.boolean_mask(lowercase_						,  lowercase_       )
                     snake_case_				:  Dict       =			tf.boolean_mask(lowercase_						,  lowercase_       )
                     snake_case_				:  Union[str, Any]       =			self._gather_logprob(lowercase_						,  lowercase_       )
                     cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
             if target is not None:
                 loss += tf.scatter_nd(lowercase_						,  -cur_logprob						,  shape_list(lowercase_       )       )
         snake_case_				:  Union[str, Any]       =			tf.concat(lowercase_						,  axis=-1       )

     if target is not None:
         if return_mean:
             snake_case_				:  Dict       =			tf.reduce_mean(lowercase_       )
         # Add the training-time loss value to the layer using `self.add_loss()`.
         self.add_loss(lowercase_       )

         # Log the loss as a metric (we could log arbitrary metrics,
         # including different metrics for training and inference.
         self.add_metric(lowercase_						,  name=self.name						,  aggregation='''mean''' if return_mean else ''''''       )

     return out





