






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


UpperCAmelCase_     :				Union[str, Any]         =				logging.getLogger(__name__)




@dataclass(frozen=UpperCAmelCase__				)
class 							lowerCAmelCase__       :






      '''simple docstring'''

      __UpperCamelCase							=						4_2
      __UpperCamelCase							=						4_2
      __UpperCamelCase							=						None
      __UpperCamelCase							=						None
      __UpperCamelCase							=						None




@dataclass(frozen=UpperCAmelCase__				)
class 							lowerCAmelCase__       :






      '''simple docstring'''

      __UpperCamelCase							=						4_2
      __UpperCamelCase							=						None
      __UpperCamelCase							=						None
      __UpperCamelCase							=						None
      __UpperCamelCase							=						None


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




     class 							lowerCAmelCase__       (    UpperCAmelCase__				):






           '''simple docstring'''

           __UpperCamelCase							=						4_2
           def __init__(		self  :						List[str]		,			lowercase_  :						str		,			lowercase_  :						PreTrainedTokenizer		,			lowercase_  :						str		,			lowercase_  :						Optional[int] = None		,			lowercase_  :						str=False		,			lowercase_  :						bool = False		,			):




               '''simple docstring'''


               SCREAMING_SNAKE_CASE_   :						int							=	hans_processors[task]()

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

               # Make sure only the first process in distributed training processes the dataset,
               # and the others will use the cache.
               SCREAMING_SNAKE_CASE_   :						List[str]							=	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}')
                       SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	torch.load(lowercase_)
                   else:
                       logger.info(F'Creating features from dataset file at {data_dir}')

                       SCREAMING_SNAKE_CASE_   :						Tuple							=	(
                           processor.get_dev_examples(lowercase_) if evaluate else processor.get_train_examples(lowercase_)
                       )

                       logger.info('''Training examples: %s'''		,			len(lowercase_))
                       SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	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  :						Optional[int]):




               '''simple docstring'''


               return len(self.features)
           def __getitem__(		self  :						Optional[Any]		,			lowercase_  :						int):




               '''simple docstring'''


               return self.features[i]



           def _SCREAMING_SNAKE_CASE				(		self  :						str):




               '''simple docstring'''


               return self.label_list


if is_tf_available():
     import tensorflow as tf




     class 							lowerCAmelCase__       :






           '''simple docstring'''

           __UpperCamelCase							=						4_2
           def __init__(		self  :						Any		,			lowercase_  :						str		,			lowercase_  :						PreTrainedTokenizer		,			lowercase_  :						str		,			lowercase_  :						Optional[int] = 128		,			lowercase_  :						Union[str, Any]=False		,			lowercase_  :						bool = False		,			):




               '''simple docstring'''


               SCREAMING_SNAKE_CASE_   :						List[str]							=	hans_processors[task]()
               SCREAMING_SNAKE_CASE_   :						List[str]							=	processor.get_labels()
               if tokenizer.__class__ in (
                   RobertaTokenizer,
                   RobertaTokenizerFast,
                   XLMRobertaTokenizer,
                   BartTokenizer,
                   BartTokenizerFast,
               ):
                   # HACK(label indices are swapped in RoBERTa pretrained model)
                   SCREAMING_SNAKE_CASE_   :						Dict							=	label_list[2], label_list[1]
               SCREAMING_SNAKE_CASE_   :						Any							=	label_list

               SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	processor.get_dev_examples(lowercase_) if evaluate else processor.get_train_examples(lowercase_)
               SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	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 % 10000 == 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,
                       )

               SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	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 _SCREAMING_SNAKE_CASE				(		self  :						Dict):




               '''simple docstring'''


               return self.dataset
           def __len__(		self  :						Tuple):




               '''simple docstring'''


               return len(self.features)
           def __getitem__(		self  :						Dict		,			lowercase_  :						Union[str, Any]):




               '''simple docstring'''


               return self.features[i]



           def _SCREAMING_SNAKE_CASE				(		self  :						str):




               '''simple docstring'''


               return self.label_list




class 							lowerCAmelCase__       (    UpperCAmelCase__				):






      '''simple docstring'''

      def _SCREAMING_SNAKE_CASE				(		self  :						str		,			lowercase_  :						Dict):




          '''simple docstring'''


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




          '''simple docstring'''


          return self._create_examples(self._read_tsv(os.path.join(lowercase_		,			'''heuristics_evaluation_set.txt'''))		,			'''dev''')
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple):




          '''simple docstring'''


          return ["contradiction", "entailment", "neutral"]



      def _SCREAMING_SNAKE_CASE				(		self  :						Dict		,			lowercase_  :						Dict		,			lowercase_  :						int):




          '''simple docstring'''


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



def       _A    (__a					,       __a					,       __a					,       __a					,       )				-> Optional[int]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						Dict							=	{label: i for i, label in enumerate(SCREAMING_SNAKE_CASE_							)}

    SCREAMING_SNAKE_CASE_   :						Dict							=	[]
    for ex_index, example in tqdm.tqdm(enumerate(SCREAMING_SNAKE_CASE_							)					,       desc='''convert examples to features'''							):
        if ex_index % 1_00_00 == 0:
            logger.info('''Writing example %d''' % (ex_index)							)

        SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	tokenizer(
            example.text_a					,       example.text_b					,       add_special_tokens=SCREAMING_SNAKE_CASE_					,       max_length=SCREAMING_SNAKE_CASE_					,       padding='''max_length'''					,       truncation=SCREAMING_SNAKE_CASE_					,       return_overflowing_tokens=SCREAMING_SNAKE_CASE_					,       )

        SCREAMING_SNAKE_CASE_   :						List[str]							=	label_map[example.label] if example.label in label_map else 0

        SCREAMING_SNAKE_CASE_   :						Any							=	int(example.pairID							)

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

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

    return features


UpperCAmelCase_     :				int         =				{
    """hans""": 3,
}





UpperCAmelCase_     :				Tuple         =				{
    """hans""": HansProcessor,
}










"""simple docstring"""






import multiprocessing
import time

from arguments import PretokenizationArguments
from datasets import load_dataset

from transformers import AutoTokenizer, HfArgumentParser



def       _A    (__a							)				-> Optional[int]:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	{}
    SCREAMING_SNAKE_CASE_   :						List[Any]							=	tokenizer(example['''content''']					,       truncation=__a							)['''input_ids''']
    SCREAMING_SNAKE_CASE_   :						List[Any]							=	len(example['''content''']							) / len(output['''input_ids''']							)
    return output


UpperCAmelCase_     :				Tuple         =				HfArgumentParser(PretokenizationArguments)
UpperCAmelCase_     :				Tuple         =				parser.parse_args()
if args.num_workers is None:
  UpperCAmelCase_     :				Tuple         =				multiprocessing.cpu_count()
UpperCAmelCase_     :				str         =				AutoTokenizer.from_pretrained(args.tokenizer_dir)

UpperCAmelCase_     :				Tuple         =				time.time()
UpperCAmelCase_     :				List[Any]         =				load_dataset(args.dataset_name, split="""train""")
print(f'''Dataset loaded in {time.time()-t_start:.2f}s''')

UpperCAmelCase_     :				Dict         =				time.time()
UpperCAmelCase_     :				Dict         =				ds.map(
    tokenize,
    num_proc=args.num_workers,
    remove_columns=[
        """repo_name""",
        """path""",
        """copies""",
        """size""",
        """content""",
        """license""",
        """hash""",
        """line_mean""",
        """line_max""",
        """alpha_frac""",
        """autogenerated""",
    ],
)
print(f'''Dataset tokenized in {time.time()-t_start:.2f}s''')

UpperCAmelCase_     :				str         =				time.time()
ds.push_to_hub(args.tokenized_data_repo)





print(f'''Data pushed to the hub in {time.time()-t_start:.2f}s''')


