




import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    MobileViTConfig,
    MobileViTForImageClassification,
    MobileViTForSemanticSegmentation,
    MobileViTImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
_lowerCamelCase			:	Dict							 =     logging.get_logger(__name__)







def SCREAMING_SNAKE_CASE       ( lowercase_		)     -> Tuple:



       """simple docstring"""

       A__   				=					MobileViTConfig()

       # size of the architecture
       if "mobilevit_s" in mobilevit_name:
              A__   				=					[144, 192, 240]
              A__   				=					[16, 32, 64, 96, 128, 160, 640]
       elif "mobilevit_xs" in mobilevit_name:
              A__   				=					[96, 120, 144]
              A__   				=					[16, 32, 48, 64, 80, 96, 384]
       elif "mobilevit_xxs" in mobilevit_name:
              A__   				=					[64, 80, 96]
              A__   				=					[16, 16, 24, 48, 64, 80, 320]
              A__   				=					0.05
              A__   				=					2.0

       if mobilevit_name.startswith('''deeplabv3_'''		):
              A__   				=					512
              A__   				=					16
              A__   				=					21
              A__   				=					'''pascal-voc-id2label.json'''
       else:
              A__   				=					1_000
              A__   				=					'''imagenet-1k-id2label.json'''

       A__   				=					'''huggingface/label-files'''
       A__   				=					json.load(open(hf_hub_download(lowercase_				,       lowercase_				,       repo_type='''dataset'''		)				,       '''r'''		)		)
       A__   				=					{int(lowercase_		): v for k, v in idalabel.items()}
       A__   				=					idalabel
       A__   				=					{v: k for k, v in idalabel.items()}

       return config







def SCREAMING_SNAKE_CASE       ( lowercase_				,       lowercase_=False		)     -> List[Any]:



       """simple docstring"""

       for i in range(1				,       6		):
              if f"""layer_{i}.""" in name:
                     A__   				=					name.replace(f"""layer_{i}."""				,       f"""encoder.layer.{i - 1}."""		)

       if "conv_1." in name:
              A__   				=					name.replace('''conv_1.'''				,       '''conv_stem.'''		)
       if ".block." in name:
              A__   				=					name.replace('''.block.'''				,       '''.'''		)

       if "exp_1x1" in name:
              A__   				=					name.replace('''exp_1x1'''				,       '''expand_1x1'''		)
       if "red_1x1" in name:
              A__   				=					name.replace('''red_1x1'''				,       '''reduce_1x1'''		)
       if ".local_rep.conv_3x3." in name:
              A__   				=					name.replace('''.local_rep.conv_3x3.'''				,       '''.conv_kxk.'''		)
       if ".local_rep.conv_1x1." in name:
              A__   				=					name.replace('''.local_rep.conv_1x1.'''				,       '''.conv_1x1.'''		)
       if ".norm." in name:
              A__   				=					name.replace('''.norm.'''				,       '''.normalization.'''		)
       if ".conv." in name:
              A__   				=					name.replace('''.conv.'''				,       '''.convolution.'''		)
       if ".conv_proj." in name:
              A__   				=					name.replace('''.conv_proj.'''				,       '''.conv_projection.'''		)

       for i in range(0				,       2		):
              for j in range(0				,       4		):
                     if f""".{i}.{j}.""" in name:
                            A__   				=					name.replace(f""".{i}.{j}."""				,       f""".{i}.layer.{j}."""		)

       for i in range(2				,       6		):
              for j in range(0				,       4		):
                     if f""".{i}.{j}.""" in name:
                            A__   				=					name.replace(f""".{i}.{j}."""				,       f""".{i}."""		)
                            if "expand_1x1" in name:
                                   A__   				=					name.replace('''expand_1x1'''				,       '''downsampling_layer.expand_1x1'''		)
                            if "conv_3x3" in name:
                                   A__   				=					name.replace('''conv_3x3'''				,       '''downsampling_layer.conv_3x3'''		)
                            if "reduce_1x1" in name:
                                   A__   				=					name.replace('''reduce_1x1'''				,       '''downsampling_layer.reduce_1x1'''		)

       for i in range(2				,       5		):
              if f""".global_rep.{i}.weight""" in name:
                     A__   				=					name.replace(f""".global_rep.{i}.weight"""				,       '''.layernorm.weight'''		)
              if f""".global_rep.{i}.bias""" in name:
                     A__   				=					name.replace(f""".global_rep.{i}.bias"""				,       '''.layernorm.bias'''		)

       if ".global_rep." in name:
              A__   				=					name.replace('''.global_rep.'''				,       '''.transformer.'''		)
       if ".pre_norm_mha.0." in name:
              A__   				=					name.replace('''.pre_norm_mha.0.'''				,       '''.layernorm_before.'''		)
       if ".pre_norm_mha.1.out_proj." in name:
              A__   				=					name.replace('''.pre_norm_mha.1.out_proj.'''				,       '''.attention.output.dense.'''		)
       if ".pre_norm_ffn.0." in name:
              A__   				=					name.replace('''.pre_norm_ffn.0.'''				,       '''.layernorm_after.'''		)
       if ".pre_norm_ffn.1." in name:
              A__   				=					name.replace('''.pre_norm_ffn.1.'''				,       '''.intermediate.dense.'''		)
       if ".pre_norm_ffn.4." in name:
              A__   				=					name.replace('''.pre_norm_ffn.4.'''				,       '''.output.dense.'''		)
       if ".transformer." in name:
              A__   				=					name.replace('''.transformer.'''				,       '''.transformer.layer.'''		)

       if ".aspp_layer." in name:
              A__   				=					name.replace('''.aspp_layer.'''				,       '''.'''		)
       if ".aspp_pool." in name:
              A__   				=					name.replace('''.aspp_pool.'''				,       '''.'''		)
       if "seg_head." in name:
              A__   				=					name.replace('''seg_head.'''				,       '''segmentation_head.'''		)
       if "segmentation_head.classifier.classifier." in name:
              A__   				=					name.replace('''segmentation_head.classifier.classifier.'''				,       '''segmentation_head.classifier.'''		)

       if "classifier.fc." in name:
              A__   				=					name.replace('''classifier.fc.'''				,       '''classifier.'''		)
       elif (not base_model) and ("segmentation_head." not in name):
              A__   				=					'''mobilevit.''' + name

       return name







def SCREAMING_SNAKE_CASE       ( lowercase_				,       lowercase_				,       lowercase_=False		)     -> List[Any]:



       """simple docstring"""

       if base_model:
              A__   				=					''''''
       else:
              A__   				=					'''mobilevit.'''

       for key in orig_state_dict.copy().keys():
              A__   				=					orig_state_dict.pop(lowercase_		)

              if key[:8] == "encoder.":
                     A__   				=					key[8:]

              if "qkv" in key:
                     A__   				=					key.split('''.'''		)
                     A__   				=					int(key_split[0][6:]		) - 1
                     A__   				=					int(key_split[3]		)
                     A__   				=					model.get_submodule(f"""{model_prefix}encoder.layer.{layer_num}"""		)
                     A__   				=					layer.transformer.layer[transformer_num].attention.attention.all_head_size
                     A__   				=					(
                         f"""{model_prefix}encoder.layer.{layer_num}.transformer.layer.{transformer_num}.attention.attention."""
                     )
                     if "weight" in key:
                            A__   				=					val[:dim, :]
                            A__   				=					val[dim : dim * 2, :]
                            A__   				=					val[-dim:, :]
                     else:
                            A__   				=					val[:dim]
                            A__   				=					val[dim : dim * 2]
                            A__   				=					val[-dim:]
              else:
                     A__   				=					val

       return orig_state_dict







def SCREAMING_SNAKE_CASE       ( )     -> Dict:



       """simple docstring"""

       A__   				=					'''http://images.cocodataset.org/val2017/000000039769.jpg'''
       A__   				=					Image.open(requests.get(lowercase_				,       stream=lowercase_		).raw		)
       return im







@torch.no_grad()
def SCREAMING_SNAKE_CASE       ( lowercase_				,       lowercase_				,       lowercase_				,       lowercase_=False		)     -> Optional[int]:



       """simple docstring"""

       A__   				=					get_mobilevit_config(lowercase_		)

       # load original state_dict
       A__   				=					torch.load(lowercase_				,       map_location='''cpu'''		)

       # load 🤗 model
       if mobilevit_name.startswith('''deeplabv3_'''		):
              A__   				=					MobileViTForSemanticSegmentation(lowercase_		).eval()
       else:
              A__   				=					MobileViTForImageClassification(lowercase_		).eval()

       A__   				=					convert_state_dict(lowercase_				,       lowercase_		)
       model.load_state_dict(lowercase_		)

       # Check outputs on an image, prepared by MobileViTImageProcessor
       A__   				=					MobileViTImageProcessor(crop_size=config.image_size				,       size=config.image_size + 32		)
       A__   				=					image_processor(images=prepare_img()				,       return_tensors='''pt'''		)
       A__   				=					model(**lowercase_		)
       A__   				=					outputs.logits

       if mobilevit_name.startswith('''deeplabv3_'''		):
              assert logits.shape == (1, 21, 32, 32)

              if mobilevit_name == "deeplabv3_mobilevit_s":
                     A__   				=					torch.tensor(
                         [
                             [[6.20_65, 6.12_92, 6.20_70], [6.10_79, 6.12_54, 6.17_47], [6.00_42, 6.10_71, 6.10_34]],
                             [[-6.92_53, -6.86_53, -7.03_98], [-7.32_18, -7.39_83, -7.36_70], [-7.19_61, -7.24_82, -7.15_69]],
                             [[-4.47_23, -4.43_48, -4.37_69], [-5.36_29, -5.46_32, -5.45_98], [-5.15_87, -5.34_02, -5.50_59]],
                         ]		)
              elif mobilevit_name == "deeplabv3_mobilevit_xs":
                     A__   				=					torch.tensor(
                         [
                             [[5.44_49, 5.57_33, 5.63_14], [5.18_15, 5.39_30, 5.59_63], [5.16_56, 5.43_33, 5.48_53]],
                             [[-9.44_23, -9.77_66, -9.67_14], [-9.15_81, -9.57_20, -9.55_19], [-9.10_06, -9.64_58, -9.57_03]],
                             [[-7.77_21, -7.37_16, -7.15_83], [-8.45_99, -8.06_24, -7.79_44], [-8.41_72, -7.83_66, -7.50_25]],
                         ]		)
              elif mobilevit_name == "deeplabv3_mobilevit_xxs":
                     A__   				=					torch.tensor(
                         [
                             [[6.98_11, 6.97_43, 7.31_23], [7.17_77, 7.19_31, 7.39_38], [7.56_33, 7.80_50, 7.89_01]],
                             [[-10.55_36, -10.23_32, -10.29_24], [-10.23_36, -9.86_24, -9.59_64], [-10.88_40, -10.81_58, -10.66_59]],
                             [[-3.49_38, -3.06_31, -2.86_20], [-3.42_05, -2.81_35, -2.68_75], [-3.41_79, -2.79_45, -2.87_50]],
                         ]		)
              else:
                     raise ValueError(f"""Unknown mobilevit_name: {mobilevit_name}"""		)

              assert torch.allclose(logits[0, :3, :3, :3]				,       lowercase_				,       atol=1E-4		)
       else:
              assert logits.shape == (1, 1_000)

              if mobilevit_name == "mobilevit_s":
                     A__   				=					torch.tensor([-0.98_66, 0.23_92, -1.12_41]		)
              elif mobilevit_name == "mobilevit_xs":
                     A__   				=					torch.tensor([-2.47_61, -0.93_99, -1.95_87]		)
              elif mobilevit_name == "mobilevit_xxs":
                     A__   				=					torch.tensor([-1.93_64, -1.23_27, -0.46_53]		)
              else:
                     raise ValueError(f"""Unknown mobilevit_name: {mobilevit_name}"""		)

              assert torch.allclose(logits[0, :3]				,       lowercase_				,       atol=1E-4		)

       Path(lowercase_		).mkdir(exist_ok=lowercase_		)
       print(f"""Saving model {mobilevit_name} to {pytorch_dump_folder_path}"""		)
       model.save_pretrained(lowercase_		)
       print(f"""Saving image processor to {pytorch_dump_folder_path}"""		)
       image_processor.save_pretrained(lowercase_		)

       if push_to_hub:
              A__   				=					{
                  '''mobilevit_s''': '''mobilevit-small''',
                  '''mobilevit_xs''': '''mobilevit-x-small''',
                  '''mobilevit_xxs''': '''mobilevit-xx-small''',
                  '''deeplabv3_mobilevit_s''': '''deeplabv3-mobilevit-small''',
                  '''deeplabv3_mobilevit_xs''': '''deeplabv3-mobilevit-x-small''',
                  '''deeplabv3_mobilevit_xxs''': '''deeplabv3-mobilevit-xx-small''',
              }

              print('''Pushing to the hub...'''		)
              A__   				=					model_mapping[mobilevit_name]
              image_processor.push_to_hub(lowercase_				,       organization='''apple'''		)
              model.push_to_hub(lowercase_				,       organization='''apple'''		)


if __name__ == "__main__":
  _lowerCamelCase			:	List[str]							 =     argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--mobilevit_name""",
      default="""mobilevit_s""",
      type=str,
      help=(
          """Name of the MobileViT model you'd like to convert. Should be one of 'mobilevit_s', 'mobilevit_xs',"""
          """ 'mobilevit_xxs', 'deeplabv3_mobilevit_s', 'deeplabv3_mobilevit_xs', 'deeplabv3_mobilevit_xxs'."""
      ),
  )
  parser.add_argument(
      """--checkpoint_path""", required=True, type=str, help="""Path to the original state dict (.pt file)."""
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""", required=True, type=str, help="""Path to the output PyTorch model directory."""
  )
  parser.add_argument(
      """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
  )

  _lowerCamelCase			:	List[str]							 =     parser.parse_args()
  convert_movilevit_checkpoint(
      args.mobilevit_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
  )








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""")

_lowerCamelCase			:	Any							 =     logging.getLogger(__name__)


@dataclass
class        UpperCamelCase_					:

    '''simple docstring'''






    UpperCAmelCase__													=   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__													=   field(
        default=UpperCAmelCase__	,      metadata={'''help''': '''Overwrite the cached preprocessed datasets or not.'''}  )
    UpperCAmelCase__													=   field(
        default=UpperCAmelCase__	,      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__													=   field(
        default=UpperCAmelCase__	,      metadata={
            '''help''': (
                '''For debugging purposes or quicker training, truncate the number of training examples to this '''
                '''value if set.'''
            )
        }	,      )
    UpperCAmelCase__													=   field(
        default=UpperCAmelCase__	,      metadata={
            '''help''': (
                '''For debugging purposes or quicker training, truncate the number of evaluation examples to this '''
                '''value if set.'''
            )
        }	,      )
    UpperCAmelCase__													=   field(
        default=UpperCAmelCase__	,      metadata={
            '''help''': (
                '''For debugging purposes or quicker training, truncate the number of prediction examples to this '''
                '''value if set.'''
            )
        }	,      )


@dataclass
class        UpperCamelCase_					:

    '''simple docstring'''






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







def SCREAMING_SNAKE_CASE       ( )     -> str:



       """simple docstring"""

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

       A__   				=					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.
       A__   				=					None
       if os.path.isdir(training_args.output_dir		) and training_args.do_train and not training_args.overwrite_output_dir:
              A__   				=					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:
                     A__   				=					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:
                     A__   				=					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				,       )
              A__   				=					train_dataset.features['''label'''].names

       if training_args.do_eval:
              A__   				=					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				,       )
              A__   				=					eval_dataset.features['''label'''].names

       if training_args.do_predict:
              A__   				=					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				,       )
              A__   				=					predict_dataset.features['''label'''].names

       # Labels
       A__   				=					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.
       A__   				=					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				,       )
       A__   				=					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				,       )
       A__   				=					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:
              A__   				=					'''max_length'''
       else:
              # We will pad later, dynamically at batch creation, to the max sequence length in each batch
              A__   				=					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:
                     A__   				=					min(len(lowercase_		)				,       data_args.max_train_samples		)
                     A__   				=					train_dataset.select(range(lowercase_		)		)
              with training_args.main_process_first(desc='''train dataset map pre-processing'''		):
                     A__   				=					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:
                     A__   				=					min(len(lowercase_		)				,       data_args.max_eval_samples		)
                     A__   				=					eval_dataset.select(range(lowercase_		)		)
              with training_args.main_process_first(desc='''validation dataset map pre-processing'''		):
                     A__   				=					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:
                     A__   				=					min(len(lowercase_		)				,       data_args.max_predict_samples		)
                     A__   				=					predict_dataset.select(range(lowercase_		)		)
              with training_args.main_process_first(desc='''prediction dataset map pre-processing'''		):
                     A__   				=					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
       A__   				=					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_		):
              A__   				=					p.predictions[0] if isinstance(p.predictions				,       lowercase_		) else p.predictions
              A__   				=					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:
              A__   				=					default_data_collator
       elif training_args.fpaa:
              A__   				=					DataCollatorWithPadding(lowercase_				,       pad_to_multiple_of=8		)
       else:
              A__   				=					None

       # Initialize our Trainer
       A__   				=					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:
              A__   				=					None
              if training_args.resume_from_checkpoint is not None:
                     A__   				=					training_args.resume_from_checkpoint
              elif last_checkpoint is not None:
                     A__   				=					last_checkpoint
              A__   				=					trainer.train(resume_from_checkpoint=lowercase_		)
              A__   				=					train_result.metrics
              A__   				=					(
                  data_args.max_train_samples if data_args.max_train_samples is not None else len(lowercase_		)
              )
              A__   				=					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 ***'''		)
              A__   				=					trainer.evaluate(eval_dataset=lowercase_		)

              A__   				=					data_args.max_eval_samples if data_args.max_eval_samples is not None else len(lowercase_		)
              A__   				=					min(lowercase_				,       len(lowercase_		)		)

              trainer.log_metrics('''eval'''				,       lowercase_		)
              trainer.save_metrics('''eval'''				,       lowercase_		)

       # Prediction
       if training_args.do_predict:
              logger.info('''*** Predict ***'''		)
              A__				,							A__				,							A__   				=					trainer.predict(lowercase_				,       metric_key_prefix='''predict'''		)

              A__   				=					(
                  data_args.max_predict_samples if data_args.max_predict_samples is not None else len(lowercase_		)
              )
              A__   				=					min(lowercase_				,       len(lowercase_		)		)

              trainer.log_metrics('''predict'''				,       lowercase_		)
              trainer.save_metrics('''predict'''				,       lowercase_		)

              A__   				=					np.argmax(lowercase_				,       axis=1		)
              A__   				=					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_		):
                                   A__   				=					label_list[item]
                                   writer.write(f"""{index}\t{item}\n"""		)


if __name__ == "__main__":
  main()


