



'''simple docstring'''







import json
from typing import List, Optional, Tuple

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_longformer import LongformerTokenizer


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

lowercase__  :	Union[str, Any]            =      {'vocab_file': 'vocab.json', 'merges_file': 'merges.txt', 'tokenizer_file': 'tokenizer.json'}

lowercase__  :	Any            =      {
    'vocab_file': {
        'allenai/longformer-base-4096': 'https://huggingface.co/allenai/longformer-base-4096/resolve/main/vocab.json',
        'allenai/longformer-large-4096': (
            'https://huggingface.co/allenai/longformer-large-4096/resolve/main/vocab.json'
        ),
        'allenai/longformer-large-4096-finetuned-triviaqa': (
            'https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/vocab.json'
        ),
        'allenai/longformer-base-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/vocab.json'
        ),
        'allenai/longformer-large-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/vocab.json'
        ),
    },
    'merges_file': {
        'allenai/longformer-base-4096': 'https://huggingface.co/allenai/longformer-base-4096/resolve/main/merges.txt',
        'allenai/longformer-large-4096': (
            'https://huggingface.co/allenai/longformer-large-4096/resolve/main/merges.txt'
        ),
        'allenai/longformer-large-4096-finetuned-triviaqa': (
            'https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/merges.txt'
        ),
        'allenai/longformer-base-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/merges.txt'
        ),
        'allenai/longformer-large-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/merges.txt'
        ),
    },
    'tokenizer_file': {
        'allenai/longformer-base-4096': (
            'https://huggingface.co/allenai/longformer-base-4096/resolve/main/tokenizer.json'
        ),
        'allenai/longformer-large-4096': (
            'https://huggingface.co/allenai/longformer-large-4096/resolve/main/tokenizer.json'
        ),
        'allenai/longformer-large-4096-finetuned-triviaqa': (
            'https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/tokenizer.json'
        ),
        'allenai/longformer-base-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/tokenizer.json'
        ),
        'allenai/longformer-large-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/tokenizer.json'
        ),
    },
}

lowercase__  :	List[str]            =      {
    'allenai/longformer-base-4096': 40_96,
    'allenai/longformer-large-4096': 40_96,
    'allenai/longformer-large-4096-finetuned-triviaqa': 40_96,
    'allenai/longformer-base-4096-extra.pos.embd.only': 40_96,
    'allenai/longformer-large-4096-extra.pos.embd.only': 40_96,
}

class    lowerCamelCase_				(						lowerCAmelCase_	):






   '''simple docstring'''

   A			:  Any					=						VOCAB_FILES_NAMES
   A			:  List[Any]					=						PRETRAINED_VOCAB_FILES_MAP
   A			:  str					=						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   A			:  List[str]					=						["input_ids", "attention_mask"]
   A			:  Any					=						LongformerTokenizer


   def __init__(	self  ,    __UpperCamelCase=None  ,    __UpperCamelCase=None  ,    __UpperCamelCase=None  ,    __UpperCamelCase="replace"  ,    __UpperCamelCase="<s>"  ,    __UpperCamelCase="</s>"  ,    __UpperCamelCase="</s>"  ,    __UpperCamelCase="<s>"  ,    __UpperCamelCase="<unk>"  ,    __UpperCamelCase="<pad>"  ,    __UpperCamelCase="<mask>"  ,    __UpperCamelCase=False  ,    __UpperCamelCase=True  ,    **__UpperCamelCase  ,    ):
       super().__init__(
           __UpperCamelCase  ,    __UpperCamelCase  ,    tokenizer_file=__UpperCamelCase  ,    errors=__UpperCamelCase  ,    bos_token=__UpperCamelCase  ,    eos_token=__UpperCamelCase  ,    sep_token=__UpperCamelCase  ,    cls_token=__UpperCamelCase  ,    unk_token=__UpperCamelCase  ,    pad_token=__UpperCamelCase  ,    mask_token=__UpperCamelCase  ,    add_prefix_space=__UpperCamelCase  ,    trim_offsets=__UpperCamelCase  ,    **__UpperCamelCase  ,    )

       UpperCamelCase__	     =      json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()			)
       if pre_tok_state.get("add_prefix_space"  ,    __UpperCamelCase			) != add_prefix_space:
           UpperCamelCase__	     =      getattr(__UpperCamelCase  ,    pre_tok_state.pop("type"			)			)
           UpperCamelCase__	     =      add_prefix_space
           UpperCamelCase__	     =      pre_tok_class(**__UpperCamelCase			)

       UpperCamelCase__	     =      add_prefix_space

       UpperCamelCase__	     =      "post_processor"
       UpperCamelCase__	     =      getattr(self.backend_tokenizer  ,    __UpperCamelCase  ,    __UpperCamelCase			)
       if tokenizer_component_instance:
           UpperCamelCase__	     =      json.loads(tokenizer_component_instance.__getstate__()			)

           # The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
           if "sep" in state:
               UpperCamelCase__	     =      tuple(state["sep"]			)
           if "cls" in state:
               UpperCamelCase__	     =      tuple(state["cls"]			)

           UpperCamelCase__	     =      False

           if state.get("add_prefix_space"  ,    __UpperCamelCase			) != add_prefix_space:
               UpperCamelCase__	     =      add_prefix_space
               UpperCamelCase__	     =      True

           if state.get("trim_offsets"  ,    __UpperCamelCase			) != trim_offsets:
               UpperCamelCase__	     =      trim_offsets
               UpperCamelCase__	     =      True

           if changes_to_apply:
               UpperCamelCase__	     =      getattr(__UpperCamelCase  ,    state.pop("type"			)			)
               UpperCamelCase__	     =      component_class(**__UpperCamelCase			)
               setattr(self.backend_tokenizer  ,    __UpperCamelCase  ,    __UpperCamelCase			)


   @property
   def 						_UpperCamelCase							(	self			):
       if self._mask_token is None:
           if self.verbose:
               logger.error("Using mask_token, but it is not set yet."			)
           return None
       return str(self._mask_token			)


   @mask_token.setter
   def 						_UpperCamelCase							(	self  ,    __UpperCamelCase			):
       UpperCamelCase__	     =      AddedToken(__UpperCamelCase  ,    lstrip=__UpperCamelCase  ,    rstrip=__UpperCamelCase			) if isinstance(__UpperCamelCase  ,    __UpperCamelCase			) else value
       UpperCamelCase__	     =      value


   def 						_UpperCamelCase							(	self  ,    *__UpperCamelCase  ,    **__UpperCamelCase			):
       UpperCamelCase__	     =      kwargs.get("is_split_into_words"  ,    __UpperCamelCase			)
       assert self.add_prefix_space or not is_split_into_words, (
           f'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
           "to use it with pretokenized inputs."
       )

       return super()._batch_encode_plus(*__UpperCamelCase  ,    **__UpperCamelCase			)


   def 						_UpperCamelCase							(	self  ,    *__UpperCamelCase  ,    **__UpperCamelCase			):
       UpperCamelCase__	     =      kwargs.get("is_split_into_words"  ,    __UpperCamelCase			)

       assert self.add_prefix_space or not is_split_into_words, (
           f'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
           "to use it with pretokenized inputs."
       )

       return super()._encode_plus(*__UpperCamelCase  ,    **__UpperCamelCase			)


   def 						_UpperCamelCase							(	self  ,    __UpperCamelCase  ,    __UpperCamelCase = None			):
       UpperCamelCase__	     =      self._tokenizer.model.save(__UpperCamelCase  ,    name=__UpperCamelCase			)
       return tuple(__UpperCamelCase			)


   def 						_UpperCamelCase							(	self  ,    __UpperCamelCase  ,    __UpperCamelCase=None			):
       UpperCamelCase__	     =      [self.bos_token_id] + token_ids_a + [self.eos_token_id]
       if token_ids_a is None:
           return output

       return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]







   def 						_UpperCamelCase							(	self  ,    __UpperCamelCase  ,    __UpperCamelCase = None			):
       UpperCamelCase__	     =      [self.sep_token_id]
       UpperCamelCase__	     =      [self.cls_token_id]

       if token_ids_a is None:
           return len(cls + token_ids_a + sep			) * [0]
       return len(cls + token_ids_a + sep + sep + token_ids_a + sep			) * [0]









'''simple docstring'''







import argparse
import logging
import os
import sys
import time

import tensorflow as tf
from datasets import load_dataset
from tqdm import tqdm

from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
from transformers.utils import is_sagemaker_dp_enabled


if os.environ.get('SDP_ENABLED') or is_sagemaker_dp_enabled():
      lowercase__  :	Any            =      True
      lowercase__  :	Tuple            =      'p3dn.24xlarge'
      import smdistributed.dataparallel.tensorflow as sdp
else:
      lowercase__  :	Dict            =      False





def     snake_case__			(lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_=None						):
    UpperCamelCase__	     =      tqdm(lowerCAmelCase_						)
    for i, batch in enumerate(lowerCAmelCase_						):
        with tf.GradientTape() as tape:
            UpperCamelCase__     ,	UpperCamelCase__	     =      batch
            UpperCamelCase__	     =      model(lowerCAmelCase_						)
            UpperCamelCase__	     =      loss(lowerCAmelCase_       ,	outputs.logits						)

        if SDP_ENABLED:
            UpperCamelCase__	     =      sdp.DistributedGradientTape(lowerCAmelCase_       ,	sparse_as_dense=lowerCAmelCase_						)

        UpperCamelCase__	     =      tape.gradient(lowerCAmelCase_       ,	model.trainable_variables						)
        opt.apply_gradients(zip(lowerCAmelCase_       ,	model.trainable_variables						)						)

        pbar.set_description(F'Loss: {loss_value:.4f}'						)

        if SDP_ENABLED and i == 0:
            sdp.broadcast_variables(model.variables       ,	root_rank=0						)
            sdp.broadcast_variables(opt.variables()       ,	root_rank=0						)

        if max_steps and i >= max_steps:
            break

    UpperCamelCase__	     =      {"loss": loss_value.numpy()}
    return train_results





def     snake_case__			(lowerCAmelCase_       ,	lowerCAmelCase_       ,	lowerCAmelCase_						):
    # Load dataset
    UpperCamelCase__     ,	UpperCamelCase__	     =      load_dataset("imdb"       ,	split=["train", "test"]						)

    # Preprocess train dataset
    UpperCamelCase__	     =      train_dataset.map(
        lambda lowerCAmelCase_						: tokenizer(e["text"]       ,	truncation=lowerCAmelCase_       ,	padding="max_length"						)       ,	batched=lowerCAmelCase_						)
    train_dataset.set_format(type="tensorflow"       ,	columns=["input_ids", "attention_mask", "label"]						)

    UpperCamelCase__	     =      {
        x: train_dataset[x].to_tensor(default_value=0       ,	shape=[None, tokenizer.model_max_length]						)
        for x in ["input_ids", "attention_mask"]
    }
    UpperCamelCase__	     =      tf.data.Dataset.from_tensor_slices((train_features, train_dataset["label"])						)

    # Preprocess test dataset
    UpperCamelCase__	     =      test_dataset.map(
        lambda lowerCAmelCase_						: tokenizer(e["text"]       ,	truncation=lowerCAmelCase_       ,	padding="max_length"						)       ,	batched=lowerCAmelCase_						)
    test_dataset.set_format(type="tensorflow"       ,	columns=["input_ids", "attention_mask", "label"]						)

    UpperCamelCase__	     =      {
        x: test_dataset[x].to_tensor(default_value=0       ,	shape=[None, tokenizer.model_max_length]						)
        for x in ["input_ids", "attention_mask"]
    }
    UpperCamelCase__	     =      tf.data.Dataset.from_tensor_slices((test_features, test_dataset["label"])						)

    if SDP_ENABLED:
        UpperCamelCase__	     =      tf_train_dataset.shard(sdp.size()       ,	sdp.rank()						)
        UpperCamelCase__	     =      tf_test_dataset.shard(sdp.size()       ,	sdp.rank()						)
    UpperCamelCase__	     =      tf_train_dataset.batch(lowerCAmelCase_       ,	drop_remainder=lowerCAmelCase_						)
    UpperCamelCase__	     =      tf_test_dataset.batch(lowerCAmelCase_       ,	drop_remainder=lowerCAmelCase_						)

    return tf_train_dataset, tf_test_dataset


if __name__ == "__main__":
      lowercase__  :	Any            =      argparse.ArgumentParser()

      # Hyperparameters sent by the client are passed as command-line arguments to the script.
      parser.add_argument('--epochs', type=int, default=3)
      parser.add_argument('--per_device_train_batch_size', type=int, default=16)
      parser.add_argument('--per_device_eval_batch_size', type=int, default=8)
      parser.add_argument('--model_name_or_path', type=str)
      parser.add_argument('--learning_rate', type=str, default=5e-5)
      parser.add_argument('--do_train', type=bool, default=True)
      parser.add_argument('--do_eval', type=bool, default=True)
      parser.add_argument('--output_dir', type=str)
      parser.add_argument('--max_steps', type=int, default=None)

      # Data, model, and output directories
      parser.add_argument('--output_data_dir', type=str, default=os.environ['SM_OUTPUT_DATA_DIR'])
      parser.add_argument('--model_dir', type=str, default=os.environ['SM_MODEL_DIR'])
      parser.add_argument('--n_gpus', type=str, default=os.environ['SM_NUM_GPUS'])

      lowercase__       ,	lowercase__  :	Union[str, Any]            =      parser.parse_known_args()

      # Set up logging
      lowercase__  :	int            =      logging.getLogger(__name__)

      logging.basicConfig(
          level=logging.getLevelName('INFO'),
          handlers=[logging.StreamHandler(sys.stdout)],
          format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
      )

      if SDP_ENABLED:
            sdp.init()

            lowercase__  :	Union[str, Any]            =      tf.config.experimental.list_physical_devices('GPU')
            for gpu in gpus:
                  tf.config.experimental.set_memory_growth(gpu, True)
            if gpus:
                  tf.config.experimental.set_visible_devices(gpus[sdp.local_rank()], 'GPU')

    # Load model and tokenizer
      lowercase__  :	int            =      TFAutoModelForSequenceClassification.from_pretrained(args.model_name_or_path)
      lowercase__  :	Any            =      AutoTokenizer.from_pretrained(args.model_name_or_path)

      # get datasets
      lowercase__       ,	lowercase__  :	Dict            =      get_datasets(
          tokenizer=tokenizer,
          train_batch_size=args.per_device_train_batch_size,
          eval_batch_size=args.per_device_eval_batch_size,
      )

      # fine optimizer and loss
      lowercase__  :	Any            =      tf.keras.optimizers.Adam(learning_rate=args.learning_rate)
      lowercase__  :	Tuple            =      tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
      lowercase__  :	Tuple            =      [tf.keras.metrics.SparseCategoricalAccuracy()]
      model.compile(optimizer=optimizer, loss=loss, metrics=metrics)

      # Training
      if args.do_train:
            # train_results = model.fit(tf_train_dataset, epochs=args.epochs, batch_size=args.train_batch_size)
            lowercase__  :	List[str]            =      time.time()
            lowercase__  :	Union[str, Any]            =      fit(
                model,
                loss,
                optimizer,
                tf_train_dataset,
                args.epochs,
                args.per_device_train_batch_size,
                max_steps=args.max_steps,
            )
            lowercase__  :	Dict            =      time.time() - start_train_time
            logger.info('*** Train ***')
            logger.info(F'''train_runtime = {end_train_time}''')

            lowercase__  :	List[str]            =      os.path.join(args.output_dir, 'train_results.txt')

            if not SDP_ENABLED or sdp.rank() == 0:
                  with open(output_eval_file, 'w') as writer:
                        logger.info('***** Train results *****')
                        logger.info(train_results)
                        for key, value in train_results.items():
                              logger.info(F'''  {key} = {value}''')
                              writer.write(F'''{key} = {value}\n''')

    # Evaluation
      if args.do_eval and (not SDP_ENABLED or sdp.rank() == 0):
            lowercase__  :	Tuple            =      model.evaluate(tf_test_dataset, batch_size=args.per_device_eval_batch_size, return_dict=True)
            logger.info('*** Evaluate ***')

            lowercase__  :	Any            =      os.path.join(args.output_dir, 'eval_results.txt')

            with open(output_eval_file, 'w') as writer:
                  logger.info('***** Eval results *****')
                  logger.info(result)
                  for key, value in result.items():
                        logger.info(F'''  {key} = {value}''')
                        writer.write(F'''{key} = {value}\n''')

    # Save result
      if SDP_ENABLED:
            if sdp.rank() == 0:
                  model.save_pretrained(args.output_dir)
                  tokenizer.save_pretrained(args.output_dir)
      else:
            model.save_pretrained(args.output_dir)
            tokenizer.save_pretrained(args.output_dir)




