



"""simple docstring"""


import logging
import os
from dataclasses import dataclass
from typing import List, Optional, Union

import tqdm
from filelock import FileLock

from transformers import (
    BartTokenizer,
    BartTokenizerFast,
    DataProcessor,
    PreTrainedTokenizer,
    RobertaTokenizer,
    RobertaTokenizerFast,
    XLMRobertaTokenizer,
    is_tf_available,
    is_torch_available,
)


snake_case_    					=	logging.getLogger(__name__)







@dataclass(frozen=lowerCamelCase__    )
class 				A_      :



						"""simple docstring"""
						__UpperCamelCase		    =							42
						__UpperCamelCase		    =							42
						__UpperCamelCase		    =							None
						__UpperCamelCase		    =							None
						__UpperCamelCase		    =							None







@dataclass(frozen=lowerCamelCase__    )
class 				A_      :



						"""simple docstring"""
						__UpperCamelCase		    =							42
						__UpperCamelCase		    =							None
						__UpperCamelCase		    =							None
						__UpperCamelCase		    =							None
						__UpperCamelCase		    =							None


if is_torch_available():
				import torch
				from torch.utils.data import Dataset







				class 				A_      (					lowerCamelCase__    ):



										"""simple docstring"""
										__UpperCamelCase		    =							42






										def __init__( self			:Dict       ,  lowercase_			:str       ,  lowercase_			:PreTrainedTokenizer       ,  lowercase_			:str       ,  lowercase_			:Optional[int] = None       ,  lowercase_			:Any=False       ,  lowercase_			:bool = False       ,  )    ->      str:
																UpperCAmelCase					=			hans_processors[task]()

																UpperCAmelCase					=			os.path.join(
																    lowercase_       ,  'cached_{}_{}_{}_{}'.format(
																        'dev' if evaluate else 'train'       ,  tokenizer.__class__.__name__       ,  str(lowercase_ )       ,  lowercase_       ,  )       ,  )
																UpperCAmelCase					=			processor.get_labels()
																if tokenizer.__class__ in (
																    RobertaTokenizer,
																    RobertaTokenizerFast,
																    XLMRobertaTokenizer,
																    BartTokenizer,
																    BartTokenizerFast,
																):
																						# HACK(label indices are swapped in RoBERTa pretrained model)
																						UpperCAmelCase					=			label_list[2], label_list[1]
																UpperCAmelCase					=			label_list

																# Make sure only the first process in distributed training processes the dataset,
																# and the others will use the cache.
																UpperCAmelCase					=			cached_features_file + ".lock"
																with FileLock(lowercase_ ):
																						if os.path.exists(lowercase_ ) and not overwrite_cache:
																												logger.info(f"""Loading features from cached file {cached_features_file}""" )
																												UpperCAmelCase					=			torch.load(lowercase_ )
																						else:
																												logger.info(f"""Creating features from dataset file at {data_dir}""" )

																												UpperCAmelCase					=			(
																												    processor.get_dev_examples(lowercase_ ) if evaluate else processor.get_train_examples(lowercase_ )
																												)

																												logger.info('Training examples: %s'       ,  len(lowercase_ ) )
																												UpperCAmelCase					=			hans_convert_examples_to_features(lowercase_       ,  lowercase_       ,  lowercase_       ,  lowercase_ )
																												logger.info('Saving features into cached file %s'       ,  lowercase_ )
																												torch.save(self.features       ,  lowercase_ )






										def __len__( self			:str )    ->      Optional[Any]:
																return len(self.features )






										def __getitem__( self			:List[str]       ,  lowercase_			:Optional[Any] )    ->      Optional[Any]:
																return self.features[i]






										def      UpperCAmelCase__					( self			:Optional[int] )    ->      Union[str, Any]:
																return self.label_list


if is_tf_available():
				import tensorflow as tf







				class 				A_      :



										"""simple docstring"""
										__UpperCamelCase		    =							42






										def __init__( self			:int       ,  lowercase_			:str       ,  lowercase_			:PreTrainedTokenizer       ,  lowercase_			:str       ,  lowercase_			:Optional[int] = 1_28       ,  lowercase_			:Dict=False       ,  lowercase_			:bool = False       ,  )    ->      List[str]:
																UpperCAmelCase					=			hans_processors[task]()
																UpperCAmelCase					=			processor.get_labels()
																if tokenizer.__class__ in (
																    RobertaTokenizer,
																    RobertaTokenizerFast,
																    XLMRobertaTokenizer,
																    BartTokenizer,
																    BartTokenizerFast,
																):
																						# HACK(label indices are swapped in RoBERTa pretrained model)
																						UpperCAmelCase					=			label_list[2], label_list[1]
																UpperCAmelCase					=			label_list

																UpperCAmelCase					=			processor.get_dev_examples(lowercase_ ) if evaluate else processor.get_train_examples(lowercase_ )
																UpperCAmelCase					=			hans_convert_examples_to_features(lowercase_       ,  lowercase_       ,  lowercase_       ,  lowercase_ )

																def gen():
																						for ex_index, ex in tqdm.tqdm(enumerate(self.features )       ,  desc='convert examples to features' ):
																												if ex_index % 1_00_00 == 0:
																																		logger.info('Writing example %d of %d' % (ex_index, len(lowercase_ )) )

																												yield (
																												    {
																												        "example_id": 0,
																												        "input_ids": ex.input_ids,
																												        "attention_mask": ex.attention_mask,
																												        "token_type_ids": ex.token_type_ids,
																												    },
																												    ex.label,
																												)

																UpperCAmelCase					=			tf.data.Dataset.from_generator(
																    lowercase_       ,  (
																        {
																            'example_id': tf.intaa,
																            'input_ids': tf.intaa,
																            'attention_mask': tf.intaa,
																            'token_type_ids': tf.intaa,
																        },
																        tf.intaa,
																    )       ,  (
																        {
																            'example_id': tf.TensorShape([] ),
																            'input_ids': tf.TensorShape([None, None] ),
																            'attention_mask': tf.TensorShape([None, None] ),
																            'token_type_ids': tf.TensorShape([None, None] ),
																        },
																        tf.TensorShape([] ),
																    )       ,  )






										def      UpperCAmelCase__					( self			:Optional[Any] )    ->      Optional[int]:
																return self.dataset






										def __len__( self			:int )    ->      Union[str, Any]:
																return len(self.features )






										def __getitem__( self			:Optional[Any]       ,  lowercase_			:List[Any] )    ->      int:
																return self.features[i]






										def      UpperCAmelCase__					( self			:Any )    ->      Union[str, Any]:
																return self.label_list







class 				A_      (					lowerCamelCase__    ):



						"""simple docstring"""






						def      UpperCAmelCase__					( self			:Any       ,  lowercase_			:Dict )    ->      Optional[Any]:
												return self._create_examples(self._read_tsv(os.path.join(lowercase_       ,  'heuristics_train_set.txt' ) )       ,  'train' )






						def      UpperCAmelCase__					( self			:Union[str, Any]       ,  lowercase_			:Union[str, Any] )    ->      Any:
												return self._create_examples(self._read_tsv(os.path.join(lowercase_       ,  'heuristics_evaluation_set.txt' ) )       ,  'dev' )






						def      UpperCAmelCase__					( self			:Union[str, Any] )    ->      Any:
												return ["contradiction", "entailment", "neutral"]






						def      UpperCAmelCase__					( self			:Optional[int]       ,  lowercase_			:Union[str, Any]       ,  lowercase_			:Union[str, Any] )    ->      Optional[Any]:
												UpperCAmelCase					=			[]
												for i, line in enumerate(lowercase_ ):
																		if i == 0:
																								continue
																		UpperCAmelCase					=			"%s-%s" % (set_type, line[0])
																		UpperCAmelCase					=			line[5]
																		UpperCAmelCase					=			line[6]
																		UpperCAmelCase					=			line[7][2:] if line[7].startswith('ex' ) else line[7]
																		UpperCAmelCase					=			line[0]
																		examples.append(InputExample(guid=lowercase_       ,  text_a=lowercase_       ,  text_b=lowercase_       ,  label=lowercase_       ,  pairID=lowercase_ ) )
												return examples


def 			_lowerCAmelCase				(       lowercase_		,     lowercase_		,     lowercase_		,     lowercase_		,     ):
						UpperCAmelCase					=			{label: i for i, label in enumerate(_lowerCAmelCase						)}

						UpperCAmelCase					=			[]
						for ex_index, example in tqdm.tqdm(enumerate(_lowerCAmelCase						)		,     desc='convert examples to features'						):
												if ex_index % 10000 == 0:
																		logger.info('Writing example %d' % (ex_index)						)

												UpperCAmelCase					=			tokenizer(
												    example.text_a		,     example.text_b		,     add_special_tokens=_lowerCAmelCase		,     max_length=_lowerCAmelCase		,     padding='max_length'		,     truncation=_lowerCAmelCase		,     return_overflowing_tokens=_lowerCAmelCase		,     )

												UpperCAmelCase					=			label_map[example.label] if example.label in label_map else 0

												UpperCAmelCase					=			int(example.pairID						)

												features.append(InputFeatures(**_lowerCAmelCase		,     label=_lowerCAmelCase		,     pairID=_lowerCAmelCase						)						)

						for i, example in enumerate(examples[:5]						):
												logger.info('*** Example ***'						)
												logger.info(F"""guid: {example}"""						)
												logger.info(F"""features: {features[i]}"""						)

						return features


snake_case_    					=	{
    '''hans''': 3,
}






snake_case_    					=	{
    '''hans''': HansProcessor,
}










"""simple docstring"""


import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp

def 			_lowerCAmelCase				(       ):
						UpperCAmelCase					=			ArgumentParser(
						    description=(
						        'PyTorch TPU distributed training launch helper utility that will spawn up multiple distributed processes'
						    )						)

						# Optional arguments for the launch helper
						parser.add_argument('--num_cores'		,     type=lowercase_		,     default=1		,     help='Number of TPU cores to use (1 or 8).'						)

						# positional
						parser.add_argument(
						    'training_script'		,     type=lowercase_		,     help=(
						        'The full path to the single TPU training '
						        'program/script to be launched in parallel, '
						        'followed by all the arguments for the '
						        'training script'
						    )		,     )

						# rest from the training program
						parser.add_argument('training_script_args'		,     nargs=lowercase_						)

						return parser.parse_args()

def 			_lowerCAmelCase				(       ):
						UpperCAmelCase					=			parse_args()

						# Import training_script as a module.
						UpperCAmelCase					=			Path(args.training_script						)
						sys.path.append(str(script_fpath.parent.resolve()						)						)
						UpperCAmelCase					=			script_fpath.stem
						UpperCAmelCase					=			importlib.import_module(lowercase_						)

						# Patch sys.argv
						UpperCAmelCase					=			[args.training_script] + args.training_script_args + ['--tpu_num_cores', str(args.num_cores						)]

						xmp.spawn(mod._mp_fn		,     args=()		,     nprocs=args.num_cores						)


if __name__ == "__main__":
				main()





