





import collections
from typing import List, Optional, Union

from ...tokenization_utils_base import BatchEncoding
from ...utils import TensorType, add_end_docstrings, add_start_docstrings, logging
from ..bert.tokenization_bert import BertTokenizer


UpperCAmelCase_           =	logging.get_logger(__name__)

UpperCAmelCase_           =	{"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}

UpperCAmelCase_           =	{
    "vocab_file": {
        "facebook/dpr-ctx_encoder-single-nq-base": (
            "https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/vocab.txt"
        ),
        "facebook/dpr-ctx_encoder-multiset-base": (
            "https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/vocab.txt"
        ),
    },
    "tokenizer_file": {
        "facebook/dpr-ctx_encoder-single-nq-base": (
            "https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/tokenizer.json"
        ),
        "facebook/dpr-ctx_encoder-multiset-base": (
            "https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/tokenizer.json"
        ),
    },
}
UpperCAmelCase_           =	{
    "vocab_file": {
        "facebook/dpr-question_encoder-single-nq-base": (
            "https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/vocab.txt"
        ),
        "facebook/dpr-question_encoder-multiset-base": (
            "https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/vocab.txt"
        ),
    },
    "tokenizer_file": {
        "facebook/dpr-question_encoder-single-nq-base": (
            "https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/tokenizer.json"
        ),
        "facebook/dpr-question_encoder-multiset-base": (
            "https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/tokenizer.json"
        ),
    },
}
UpperCAmelCase_           =	{
    "vocab_file": {
        "facebook/dpr-reader-single-nq-base": (
            "https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/vocab.txt"
        ),
        "facebook/dpr-reader-multiset-base": (
            "https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/vocab.txt"
        ),
    },
    "tokenizer_file": {
        "facebook/dpr-reader-single-nq-base": (
            "https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/tokenizer.json"
        ),
        "facebook/dpr-reader-multiset-base": (
            "https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/tokenizer.json"
        ),
    },
}

UpperCAmelCase_           =	{
    "facebook/dpr-ctx_encoder-single-nq-base": 5_12,
    "facebook/dpr-ctx_encoder-multiset-base": 5_12,
}
UpperCAmelCase_           =	{
    "facebook/dpr-question_encoder-single-nq-base": 5_12,
    "facebook/dpr-question_encoder-multiset-base": 5_12,
}
UpperCAmelCase_           =	{
    "facebook/dpr-reader-single-nq-base": 5_12,
    "facebook/dpr-reader-multiset-base": 5_12,
}


UpperCAmelCase_           =	{
    "facebook/dpr-ctx_encoder-single-nq-base": {"do_lower_case": True},
    "facebook/dpr-ctx_encoder-multiset-base": {"do_lower_case": True},
}
UpperCAmelCase_           =	{
    "facebook/dpr-question_encoder-single-nq-base": {"do_lower_case": True},
    "facebook/dpr-question_encoder-multiset-base": {"do_lower_case": True},
}
UpperCAmelCase_           =	{
    "facebook/dpr-reader-single-nq-base": {"do_lower_case": True},
    "facebook/dpr-reader-multiset-base": {"do_lower_case": True},
}




class 			__UpperCamelCase    (							A__    ):
     __A	:  Optional[int]							  =   VOCAB_FILES_NAMES
     __A	:  int							  =   CONTEXT_ENCODER_PRETRAINED_VOCAB_FILES_MAP
     __A	:  List[str]							  =   CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     __A	:  Optional[int]							  =   CONTEXT_ENCODER_PRETRAINED_INIT_CONFIGURATION




class 			__UpperCamelCase    (							A__    ):
     __A	:  Dict							  =   VOCAB_FILES_NAMES
     __A	:  List[str]							  =   QUESTION_ENCODER_PRETRAINED_VOCAB_FILES_MAP
     __A	:  Any							  =   QUESTION_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     __A	:  str							  =   QUESTION_ENCODER_PRETRAINED_INIT_CONFIGURATION




UpperCAmelCase_           =	collections.namedtuple(
    "DPRSpanPrediction", ["span_score", "relevance_score", "doc_id", "start_index", "end_index", "text"]
)

UpperCAmelCase_           =	collections.namedtuple("DPRReaderOutput", ["start_logits", "end_logits", "relevance_logits"])


UpperCAmelCase_           =	r"\n    Return a dictionary with the token ids of the input strings and other information to give to `.decode_best_spans`.\n    It converts the strings of a question and different passages (title and text) in a sequence of IDs (integers),\n    using the tokenizer and vocabulary. The resulting `input_ids` is a matrix of size `(n_passages, sequence_length)`\n    with the format:\n\n    ```\n    [CLS] <question token ids> [SEP] <titles ids> [SEP] <texts ids>\n    ```\n\n    Args:\n        questions (`str` or `List[str]`):\n            The questions to be encoded. You can specify one question for many passages. In this case, the question\n            will be duplicated like `[questions] * n_passages`. Otherwise you have to specify as many questions as in\n            `titles` or `texts`.\n        titles (`str` or `List[str]`):\n            The passages titles to be encoded. This can be a string or a list of strings if there are several passages.\n        texts (`str` or `List[str]`):\n            The passages texts to be encoded. This can be a string or a list of strings if there are several passages.\n        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`):\n            Activates and controls padding. Accepts the following values:\n\n            - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single sequence\n              if provided).\n            - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum\n              acceptable input length for the model if that argument is not provided.\n            - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different\n              lengths).\n        truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`):\n            Activates and controls truncation. Accepts the following values:\n\n            - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or to\n              the maximum acceptable input length for the model if that argument is not provided. This will truncate\n              token by token, removing a token from the longest sequence in the pair if a pair of sequences (or a batch\n              of pairs) is provided.\n            - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum\n              acceptable input length for the model if that argument is not provided. This will only truncate the first\n              sequence of a pair if a pair of sequences (or a batch of pairs) is provided.\n            - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum\n              acceptable input length for the model if that argument is not provided. This will only truncate the\n              second sequence of a pair if a pair of sequences (or a batch of pairs) is provided.\n            - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths\n              greater than the model maximum admissible input size).\n        max_length (`int`, *optional*):\n                Controls the maximum length to use by one of the truncation/padding parameters.\n\n                If left unset or set to `None`, this will use the predefined model maximum length if a maximum length\n                is required by one of the truncation/padding parameters. If the model has no specific maximum input\n                length (like XLNet) truncation/padding to a maximum length will be deactivated.\n        return_tensors (`str` or [`~utils.TensorType`], *optional*):\n                If set, will return tensors instead of list of python integers. Acceptable values are:\n\n                - `'tf'`: Return TensorFlow `tf.constant` objects.\n                - `'pt'`: Return PyTorch `torch.Tensor` objects.\n                - `'np'`: Return Numpy `np.ndarray` objects.\n        return_attention_mask (`bool`, *optional*):\n            Whether or not to return the attention mask. If not set, will return the attention mask according to the\n            specific tokenizer's default, defined by the `return_outputs` attribute.\n\n            [What are attention masks?](../glossary#attention-mask)\n\n    Returns:\n        `Dict[str, List[List[int]]]`: A dictionary with the following keys:\n\n        - `input_ids`: List of token ids to be fed to a model.\n        - `attention_mask`: List of indices specifying which tokens should be attended to by the model.\n    "




@add_start_docstrings(A__    )
class 			__UpperCamelCase    :





     def __call__(     self				,							_UpperCamelCase				,							_UpperCamelCase = None				,							_UpperCamelCase = None				,							_UpperCamelCase = False				,							_UpperCamelCase = False				,							_UpperCamelCase = None				,							_UpperCamelCase = None				,							_UpperCamelCase = None				,							**_UpperCamelCase				,							):
        if titles is None and texts is None:
           return super().__call__(
               _UpperCamelCase				,							padding=_UpperCamelCase				,							truncation=_UpperCamelCase				,							max_length=_UpperCamelCase				,							return_tensors=_UpperCamelCase				,							return_attention_mask=_UpperCamelCase				,							**_UpperCamelCase				,							)
        elif titles is None or texts is None:
           _UpperCAmelCase           =   titles if texts is None else texts
           return super().__call__(
               _UpperCamelCase				,							_UpperCamelCase				,							padding=_UpperCamelCase				,							truncation=_UpperCamelCase				,							max_length=_UpperCamelCase				,							return_tensors=_UpperCamelCase				,							return_attention_mask=_UpperCamelCase				,							**_UpperCamelCase				,							)
        _UpperCAmelCase           =   titles if not isinstance(_UpperCamelCase				,							_UpperCamelCase					) else [titles]
        _UpperCAmelCase           =   texts if not isinstance(_UpperCamelCase				,							_UpperCamelCase					) else [texts]
        _UpperCAmelCase           =   len(_UpperCamelCase					)
        _UpperCAmelCase           =   questions if not isinstance(_UpperCamelCase				,							_UpperCamelCase					) else [questions] * n_passages
        if len(_UpperCamelCase					) != len(_UpperCamelCase					):
           raise ValueError(
               f'''There should be as many titles than texts but got {len(_UpperCamelCase					)} titles and {len(_UpperCamelCase					)} texts.'''					)
        _UpperCAmelCase           =   super().__call__(_UpperCamelCase				,							_UpperCamelCase				,							padding=_UpperCamelCase				,							truncation=_UpperCamelCase					)['''input_ids''']
        _UpperCAmelCase           =   super().__call__(_UpperCamelCase				,							add_special_tokens=_UpperCamelCase				,							padding=_UpperCamelCase				,							truncation=_UpperCamelCase					)['''input_ids''']
        _UpperCAmelCase           =   {
            '''input_ids''': [
                (encoded_question_and_title + encoded_text)[:max_length]
                if max_length is not None and truncation
                else encoded_question_and_title + encoded_text
                for encoded_question_and_title, encoded_text in zip(_UpperCamelCase				,							_UpperCamelCase					)
            ]
        }
        if return_attention_mask is not False:
           _UpperCAmelCase           =   []
           for input_ids in encoded_inputs["input_ids"]:
              attention_mask.append([int(input_id != self.pad_token_id					) for input_id in input_ids]					)
           _UpperCAmelCase           =   attention_mask
        return self.pad(_UpperCamelCase				,							padding=_UpperCamelCase				,							max_length=_UpperCamelCase				,							return_tensors=_UpperCamelCase					)





     def 		UpperCamelCase(     self				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase = 16				,							_UpperCamelCase = 64				,							_UpperCamelCase = 4				,							):
        _UpperCAmelCase           =   reader_input['''input_ids''']
        _UpperCAmelCase       ,      _UpperCAmelCase       ,      _UpperCAmelCase           =   reader_output[:3]
        _UpperCAmelCase           =   len(_UpperCamelCase					)
        _UpperCAmelCase           =   sorted(range(_UpperCamelCase					)				,							reverse=_UpperCamelCase				,							key=relevance_logits.__getitem__					)
        _UpperCAmelCase           =   []
        for doc_id in sorted_docs:
           _UpperCAmelCase           =   list(input_ids[doc_id]					)
           # assuming question & title information is at the beginning of the sequence
           _UpperCAmelCase           =   sequence_ids.index(self.sep_token_id				,							2					) + 1  # second sep id
           if sequence_ids[-1] == self.pad_token_id:
              _UpperCAmelCase           =   sequence_ids.index(self.pad_token_id					)
           else:
              _UpperCAmelCase           =   len(_UpperCamelCase					)

           _UpperCAmelCase           =   self._get_best_spans(
               start_logits=start_logits[doc_id][passage_offset:sequence_len]				,							end_logits=end_logits[doc_id][passage_offset:sequence_len]				,							max_answer_length=_UpperCamelCase				,							top_spans=_UpperCamelCase				,							)
           for start_index, end_index in best_spans:
              start_index += passage_offset
              end_index += passage_offset
              nbest_spans_predictions.append(
                  DPRSpanPrediction(
                      span_score=start_logits[doc_id][start_index] + end_logits[doc_id][end_index]				,							relevance_score=relevance_logits[doc_id]				,							doc_id=_UpperCamelCase				,							start_index=_UpperCamelCase				,							end_index=_UpperCamelCase				,							text=self.decode(sequence_ids[start_index : end_index + 1]					)				,							)					)
           if len(_UpperCamelCase					) >= num_spans:
              break
        return nbest_spans_predictions[:num_spans]





     def 		UpperCamelCase(     self				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							):
        _UpperCAmelCase           =   []
        for start_index, start_score in enumerate(_UpperCamelCase					):
           for answer_length, end_score in enumerate(end_logits[start_index : start_index + max_answer_length]					):
              scores.append(((start_index, start_index + answer_length), start_score + end_score)					)
        _UpperCAmelCase           =   sorted(_UpperCamelCase				,							key=lambda _UpperCamelCase					: x[1]				,							reverse=_UpperCamelCase					)
        _UpperCAmelCase           =   []
        for (start_index, end_index), score in scores:
           if start_index > end_index:
              raise ValueError(f'''Wrong span indices: [{start_index}:{end_index}]'''					)
           _UpperCAmelCase           =   end_index - start_index + 1
           if length > max_answer_length:
              raise ValueError(f'''Span is too long: {length} > {max_answer_length}'''					)
           if any(
               start_index <= prev_start_index <= prev_end_index <= end_index
               or prev_start_index <= start_index <= end_index <= prev_end_index
               for (prev_start_index, prev_end_index) in chosen_span_intervals					):
              continue
           chosen_span_intervals.append((start_index, end_index)					)

           if len(_UpperCamelCase					) == top_spans:
              break
        return chosen_span_intervals




@add_end_docstrings(A__    )
class 			__UpperCamelCase    (							A__	, A__    ):
     __A	:  Dict							  =   VOCAB_FILES_NAMES
     __A	:  Dict							  =   READER_PRETRAINED_VOCAB_FILES_MAP
     __A	:  int							  =   READER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     __A	:  Optional[Any]							  =   READER_PRETRAINED_INIT_CONFIGURATION
     __A	:  Tuple							  =   ["""input_ids""", """attention_mask"""]


"""simple docstring"""

import logging
import os
import random
import sys
from dataclasses import dataclass, field
from typing import Optional

import datasets
import evaluate
import numpy as np
from datasets import load_dataset

import transformers
from transformers import (
    AutoConfig,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    EvalPrediction,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    default_data_collator,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version


# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
check_min_version("""4.31.0""")

require_version("""datasets>=1.8.0""", """To fix: pip install -r examples/pytorch/text-classification/requirements.txt""")

_UpperCamelCase		     =    logging.getLogger(__name__)






@dataclass
class 			__a					:





       """simple docstring"""


       __UpperCamelCase					:      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.'
               )
           }       ,	)
       __UpperCamelCase					:      bool    =						field(
           default=__magic_name__       ,	metadata={'help': 'Overwrite the cached preprocessed datasets or not.'}       )
       __UpperCamelCase					:      bool    =						field(
           default=__magic_name__       ,	metadata={
               'help': (
                   'Whether to pad all samples to `max_seq_length`. '
                   'If False, will pad the samples dynamically when batching to the maximum length in the batch.'
               )
           }       ,	)
       __UpperCamelCase					:      Optional[int]    =						field(
           default=__magic_name__       ,	metadata={
               'help': (
                   'For debugging purposes or quicker training, truncate the number of training examples to this '
                   'value if set.'
               )
           }       ,	)
       __UpperCamelCase					:      Optional[int]    =						field(
           default=__magic_name__       ,	metadata={
               'help': (
                   'For debugging purposes or quicker training, truncate the number of evaluation examples to this '
                   'value if set.'
               )
           }       ,	)
       __UpperCamelCase					:      Optional[int]    =						field(
           default=__magic_name__       ,	metadata={
               'help': (
                   'For debugging purposes or quicker training, truncate the number of prediction examples to this '
                   'value if set.'
               )
           }       ,	)






@dataclass
class 			__a					:





       """simple docstring"""


       __UpperCamelCase					:      str    =						field(
           default=__magic_name__       ,	metadata={'help': 'Path to pretrained model or model identifier from huggingface.co/models'}       )
       __UpperCamelCase					:      str    =						field(
           default=__magic_name__       ,	metadata={'help': 'Evaluation language. Also train language if `train_language` is set to None.'}       )
       __UpperCamelCase					:      Optional[str]    =						field(
           default=__magic_name__       ,	metadata={'help': 'Train language if it is different from the evaluation language.'}       )
       __UpperCamelCase					:      Optional[str]    =						field(
           default=__magic_name__       ,	metadata={'help': 'Pretrained config name or path if not the same as model_name'}       )
       __UpperCamelCase					:      Optional[str]    =						field(
           default=__magic_name__       ,	metadata={'help': 'Pretrained tokenizer name or path if not the same as model_name'}       )
       __UpperCamelCase					:      Optional[str]    =						field(
           default=__magic_name__       ,	metadata={'help': 'Where do you want to store the pretrained models downloaded from huggingface.co'}       ,	)
       __UpperCamelCase					:      Optional[bool]    =						field(
           default=__magic_name__       ,	metadata={'help': 'arg to indicate if tokenizer should do lower case in AutoTokenizer.from_pretrained()'}       ,	)
       __UpperCamelCase					:      bool    =						field(
           default=__magic_name__       ,	metadata={'help': 'Whether to use one of the fast tokenizer (backed by the tokenizers library) or not.'}       ,	)
       __UpperCamelCase					:      str    =						field(
           default='main'       ,	metadata={'help': 'The specific model version to use (can be a branch name, tag name or commit id).'}       ,	)
       __UpperCamelCase					:      bool    =						field(
           default=__magic_name__       ,	metadata={
               'help': (
                   'Will use the token generated when running `huggingface-cli login` (necessary to use this script '
                   'with private models).'
               )
           }       ,	)
       __UpperCamelCase					:      bool    =						field(
           default=__magic_name__       ,	metadata={'help': 'Will enable to load a pretrained model whose head dimensions are different.'}       ,	)




def       SCREAMING_SNAKE_CASE							( )			->						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.

 lowerCAmelCase__   :				List[Any] 			=				HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)   )
 lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__   :				Optional[Any] 			=				parser.parse_args_into_dataclasses()

 # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
 # information sent is the one passed as arguments along with your Python/PyTorch versions.
 send_example_telemetry("run_xnli"					,					lowercase__   )

 # Setup logging
 logging.basicConfig(
     format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"					,					datefmt="%m/%d/%Y %H:%M:%S"					,					handlers=[logging.StreamHandler(sys.stdout   )]					,					)

 if training_args.should_log:
  # The default of training_args.log_level is passive, so we set log level at info here to have that default.
  transformers.utils.logging.set_verbosity_info()

 lowerCAmelCase__   :				Optional[int] 			=				training_args.get_process_log_level()
 logger.setLevel(lowercase__   )
 datasets.utils.logging.set_verbosity(lowercase__   )
 transformers.utils.logging.set_verbosity(lowercase__   )
 transformers.utils.logging.enable_default_handler()
 transformers.utils.logging.enable_explicit_format()

 # Log on each process the small summary:
 logger.warning(
     F"""Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"""
     + F"""distributed training: {bool(training_args.local_rank != -1   )}, 16-bits training: {training_args.fpaa}"""   )
 logger.info(F"""Training/evaluation parameters {training_args}"""   )

 # Detecting last checkpoint.
 lowerCAmelCase__   :				str 			=				None
 if os.path.isdir(training_args.output_dir   ) and training_args.do_train and not training_args.overwrite_output_dir:
  lowerCAmelCase__   :				List[Any] 			=				get_last_checkpoint(training_args.output_dir   )
  if last_checkpoint is None and len(os.listdir(training_args.output_dir   )   ) > 0:
   raise ValueError(
       F"""Output directory ({training_args.output_dir}) already exists and is not empty. """
       "Use --overwrite_output_dir to overcome."   )
  elif last_checkpoint is not None:
   logger.info(
       F"""Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change """
       "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."   )

    # Set seed before initializing model.
 set_seed(training_args.seed   )

 # In distributed training, the load_dataset function guarantees that only one local process can concurrently
 # download the dataset.
 # Downloading and loading xnli dataset from the hub.
 if training_args.do_train:
  if model_args.train_language is None:
   lowerCAmelCase__   :				Any 			=				load_dataset(
       "xnli"					,					model_args.language					,					split="train"					,					cache_dir=model_args.cache_dir					,					use_auth_token=True if model_args.use_auth_token else None					,					)
  else:
   lowerCAmelCase__   :				Optional[Any] 			=				load_dataset(
       "xnli"					,					model_args.train_language					,					split="train"					,					cache_dir=model_args.cache_dir					,					use_auth_token=True if model_args.use_auth_token else None					,					)
  lowerCAmelCase__   :				Optional[int] 			=				train_dataset.features["label"].names

 if training_args.do_eval:
  lowerCAmelCase__   :				Tuple 			=				load_dataset(
      "xnli"					,					model_args.language					,					split="validation"					,					cache_dir=model_args.cache_dir					,					use_auth_token=True if model_args.use_auth_token else None					,					)
  lowerCAmelCase__   :				Dict 			=				eval_dataset.features["label"].names

 if training_args.do_predict:
  lowerCAmelCase__   :				Any 			=				load_dataset(
      "xnli"					,					model_args.language					,					split="test"					,					cache_dir=model_args.cache_dir					,					use_auth_token=True if model_args.use_auth_token else None					,					)
  lowerCAmelCase__   :				Union[str, Any] 			=				predict_dataset.features["label"].names

 # Labels
 lowerCAmelCase__   :				Optional[Any] 			=				len(lowercase__   )

 # Load pretrained model and tokenizer
 # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently
 # download model & vocab.
 lowerCAmelCase__   :				Tuple 			=				AutoConfig.from_pretrained(
     model_args.config_name if model_args.config_name else model_args.model_name_or_path					,					num_labels=lowercase__					,					idalabel={str(lowercase__   ): label for i, label in enumerate(lowercase__   )}					,					labelaid={label: i for i, label in enumerate(lowercase__   )}					,					finetuning_task="xnli"					,					cache_dir=model_args.cache_dir					,					revision=model_args.model_revision					,					use_auth_token=True if model_args.use_auth_token else None					,					)
 lowerCAmelCase__   :				Dict 			=				AutoTokenizer.from_pretrained(
     model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path					,					do_lower_case=model_args.do_lower_case					,					cache_dir=model_args.cache_dir					,					use_fast=model_args.use_fast_tokenizer					,					revision=model_args.model_revision					,					use_auth_token=True if model_args.use_auth_token else None					,					)
 lowerCAmelCase__   :				List[str] 			=				AutoModelForSequenceClassification.from_pretrained(
     model_args.model_name_or_path					,					from_tf=bool(".ckpt" in model_args.model_name_or_path   )					,					config=lowercase__					,					cache_dir=model_args.cache_dir					,					revision=model_args.model_revision					,					use_auth_token=True if model_args.use_auth_token else None					,					ignore_mismatched_sizes=model_args.ignore_mismatched_sizes					,					)

 # Preprocessing the datasets
 # Padding strategy
 if data_args.pad_to_max_length:
  lowerCAmelCase__   :				List[str] 			=				"max_length"
 else:
  # We will pad later, dynamically at batch creation, to the max sequence length in each batch
  lowerCAmelCase__   :				List[Any] 			=				False

 def preprocess_function(lowercase__   ):
  # Tokenize the texts
  return tokenizer(
      examples["premise"]					,					examples["hypothesis"]					,					padding=lowercase__					,					max_length=data_args.max_seq_length					,					truncation=lowercase__					,					)

 if training_args.do_train:
  if data_args.max_train_samples is not None:
   lowerCAmelCase__   :				Optional[Any] 			=				min(len(lowercase__   )					,					data_args.max_train_samples   )
   lowerCAmelCase__   :				Optional[int] 			=				train_dataset.select(range(lowercase__   )   )
  with training_args.main_process_first(desc="train dataset map pre-processing"   ):
   lowerCAmelCase__   :				int 			=				train_dataset.map(
       lowercase__					,					batched=lowercase__					,					load_from_cache_file=not data_args.overwrite_cache					,					desc="Running tokenizer on train dataset"					,					)
  # Log a few random samples from the training set:
  for index in random.sample(range(len(lowercase__   )   )					,					3   ):
   logger.info(F"""Sample {index} of the training set: {train_dataset[index]}."""   )

 if training_args.do_eval:
  if data_args.max_eval_samples is not None:
   lowerCAmelCase__   :				Optional[int] 			=				min(len(lowercase__   )					,					data_args.max_eval_samples   )
   lowerCAmelCase__   :				Union[str, Any] 			=				eval_dataset.select(range(lowercase__   )   )
  with training_args.main_process_first(desc="validation dataset map pre-processing"   ):
   lowerCAmelCase__   :				List[str] 			=				eval_dataset.map(
       lowercase__					,					batched=lowercase__					,					load_from_cache_file=not data_args.overwrite_cache					,					desc="Running tokenizer on validation dataset"					,					)

 if training_args.do_predict:
  if data_args.max_predict_samples is not None:
   lowerCAmelCase__   :				str 			=				min(len(lowercase__   )					,					data_args.max_predict_samples   )
   lowerCAmelCase__   :				Union[str, Any] 			=				predict_dataset.select(range(lowercase__   )   )
  with training_args.main_process_first(desc="prediction dataset map pre-processing"   ):
   lowerCAmelCase__   :				Dict 			=				predict_dataset.map(
       lowercase__					,					batched=lowercase__					,					load_from_cache_file=not data_args.overwrite_cache					,					desc="Running tokenizer on prediction dataset"					,					)

    # Get the metric function
 lowerCAmelCase__   :				int 			=				evaluate.load("xnli"   )

 # You can define your custom compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with a
 # predictions and label_ids field) and has to return a dictionary string to float.
 def compute_metrics(lowercase__   ):
  lowerCAmelCase__   :				Union[str, Any] 			=				p.predictions[0] if isinstance(p.predictions					,					lowercase__   ) else p.predictions
  lowerCAmelCase__   :				Union[str, Any] 			=				np.argmax(lowercase__					,					axis=1   )
  return metric.compute(predictions=lowercase__					,					references=p.label_ids   )

 # Data collator will default to DataCollatorWithPadding, so we change it if we already did the padding.
 if data_args.pad_to_max_length:
  lowerCAmelCase__   :				Tuple 			=				default_data_collator
 elif training_args.fpaa:
  lowerCAmelCase__   :				Tuple 			=				DataCollatorWithPadding(lowercase__					,					pad_to_multiple_of=8   )
 else:
  lowerCAmelCase__   :				List[Any] 			=				None

 # Initialize our Trainer
 lowerCAmelCase__   :				int 			=				Trainer(
     model=lowercase__					,					args=lowercase__					,					train_dataset=train_dataset if training_args.do_train else None					,					eval_dataset=eval_dataset if training_args.do_eval else None					,					compute_metrics=lowercase__					,					tokenizer=lowercase__					,					data_collator=lowercase__					,					)

 # Training
 if training_args.do_train:
  lowerCAmelCase__   :				str 			=				None
  if training_args.resume_from_checkpoint is not None:
   lowerCAmelCase__   :				Optional[Any] 			=				training_args.resume_from_checkpoint
  elif last_checkpoint is not None:
   lowerCAmelCase__   :				List[str] 			=				last_checkpoint
  lowerCAmelCase__   :				Optional[Any] 			=				trainer.train(resume_from_checkpoint=lowercase__   )
  lowerCAmelCase__   :				str 			=				train_result.metrics
  lowerCAmelCase__   :				Union[str, Any] 			=				(
      data_args.max_train_samples if data_args.max_train_samples is not None else len(lowercase__   )
  )
  lowerCAmelCase__   :				List[Any] 			=				min(lowercase__					,					len(lowercase__   )   )

  trainer.save_model()  # Saves the tokenizer too for easy upload

  trainer.log_metrics("train"					,					lowercase__   )
  trainer.save_metrics("train"					,					lowercase__   )
  trainer.save_state()

 # Evaluation
 if training_args.do_eval:
  logger.info("*** Evaluate ***"   )
  lowerCAmelCase__   :				str 			=				trainer.evaluate(eval_dataset=lowercase__   )

  lowerCAmelCase__   :				Tuple 			=				data_args.max_eval_samples if data_args.max_eval_samples is not None else len(lowercase__   )
  lowerCAmelCase__   :				Any 			=				min(lowercase__					,					len(lowercase__   )   )

  trainer.log_metrics("eval"					,					lowercase__   )
  trainer.save_metrics("eval"					,					lowercase__   )

 # Prediction
 if training_args.do_predict:
  logger.info("*** Predict ***"   )
  lowerCAmelCase__			,						lowerCAmelCase__			,						lowerCAmelCase__   :				List[Any] 			=				trainer.predict(lowercase__					,					metric_key_prefix="predict"   )

  lowerCAmelCase__   :				List[str] 			=				(
      data_args.max_predict_samples if data_args.max_predict_samples is not None else len(lowercase__   )
  )
  lowerCAmelCase__   :				str 			=				min(lowercase__					,					len(lowercase__   )   )

  trainer.log_metrics("predict"					,					lowercase__   )
  trainer.save_metrics("predict"					,					lowercase__   )

  lowerCAmelCase__   :				Union[str, Any] 			=				np.argmax(lowercase__					,					axis=1   )
  lowerCAmelCase__   :				List[str] 			=				os.path.join(training_args.output_dir					,					"predictions.txt"   )
  if trainer.is_world_process_zero():
   with open(lowercase__					,					"w"   ) as writer:
    writer.write("index\tprediction\n"   )
    for index, item in enumerate(lowercase__   ):
     lowerCAmelCase__   :				List[Any] 			=				label_list[item]
     writer.write(F"""{index}\t{item}\n"""   )


if __name__ == "__main__":
    main()






