

'''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():
      _lowerCAmelCase    			=						True
      _lowerCAmelCase    			=						'''p3dn.24xlarge'''
      import smdistributed.dataparallel.tensorflow as sdp
else:
      _lowerCAmelCase    			=						False




def 						UpperCamelCase       ( a		,						a		,						a		,						a		,						a		,						a		,						a=None )      ->	Dict:






     '''simple docstring'''
     __magic_name__						   =				tqdm(_UpperCAmelCase )
     for i, batch in enumerate(_UpperCAmelCase ):
          with tf.GradientTape() as tape:
               __magic_name__   ,  __magic_name__						   =				batch
               __magic_name__						   =				model(_UpperCAmelCase )
               __magic_name__						   =				loss(_UpperCAmelCase		,						outputs.logits )

          if SDP_ENABLED:
               __magic_name__						   =				sdp.DistributedGradientTape(_UpperCAmelCase		,						sparse_as_dense=_UpperCAmelCase )

          __magic_name__						   =				tape.gradient(_UpperCAmelCase		,						model.trainable_variables )
          opt.apply_gradients(zip(_UpperCAmelCase		,						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

     __magic_name__						   =				{'''loss''': loss_value.numpy()}
     return train_results




def 						UpperCamelCase       ( a		,						a		,						a )      ->	Dict:






     '''simple docstring'''
     # Load dataset
     __magic_name__   ,  __magic_name__						   =				load_dataset('''imdb'''		,						split=['''train''', '''test'''] )

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

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

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

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

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

     return tf_train_dataset, tf_test_dataset


if __name__ == "__main__":
      _lowerCAmelCase    			=						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"])

      _lowerCAmelCase    			=						parser.parse_known_args()

      # Set up logging
      _lowerCAmelCase    			=						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()

            _lowerCAmelCase    			=						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
      _lowerCAmelCase    			=						TFAutoModelForSequenceClassification.from_pretrained(args.model_name_or_path)
      _lowerCAmelCase    			=						AutoTokenizer.from_pretrained(args.model_name_or_path)

      # get datasets
      _lowerCAmelCase    			=						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
      _lowerCAmelCase    			=						tf.keras.optimizers.Adam(learning_rate=args.learning_rate)
      _lowerCAmelCase    			=						tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
      _lowerCAmelCase    			=						[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)
            _lowerCAmelCase    			=						time.time()
            _lowerCAmelCase    			=						fit(
                model,
                loss,
                optimizer,
                tf_train_dataset,
                args.epochs,
                args.per_device_train_batch_size,
                max_steps=args.max_steps,
            )
            _lowerCAmelCase    			=						time.time() - start_train_time
            logger.info("*** Train ***")
            logger.info(F'''train_runtime = {end_train_time}''')

            _lowerCAmelCase    			=						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):
            _lowerCAmelCase    			=						model.evaluate(tf_test_dataset, batch_size=args.per_device_eval_batch_size, return_dict=True)
            logger.info("*** Evaluate ***")

            _lowerCAmelCase    			=						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)













import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
							from .tokenization_xlm_roberta import XLMRobertaTokenizer
else:
							UpperCAmelCase__  :	Union[str, Any]     =None


UpperCAmelCase__  :	str     =logging.get_logger(__name__)

UpperCAmelCase__  :	Optional[int]     ={'''vocab_file''': '''sentencepiece.bpe.model''', '''tokenizer_file''': '''tokenizer.json'''}

UpperCAmelCase__  :	str     ={
    '''vocab_file''': {
        '''xlm-roberta-base''': '''https://huggingface.co/xlm-roberta-base/resolve/main/sentencepiece.bpe.model''',
        '''xlm-roberta-large''': '''https://huggingface.co/xlm-roberta-large/resolve/main/sentencepiece.bpe.model''',
        '''xlm-roberta-large-finetuned-conll02-dutch''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll02-dutch/resolve/main/sentencepiece.bpe.model'''
        ),
        '''xlm-roberta-large-finetuned-conll02-spanish''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll02-spanish/resolve/main/sentencepiece.bpe.model'''
        ),
        '''xlm-roberta-large-finetuned-conll03-english''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll03-english/resolve/main/sentencepiece.bpe.model'''
        ),
        '''xlm-roberta-large-finetuned-conll03-german''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll03-german/resolve/main/sentencepiece.bpe.model'''
        ),
    },
    '''tokenizer_file''': {
        '''xlm-roberta-base''': '''https://huggingface.co/xlm-roberta-base/resolve/main/tokenizer.json''',
        '''xlm-roberta-large''': '''https://huggingface.co/xlm-roberta-large/resolve/main/tokenizer.json''',
        '''xlm-roberta-large-finetuned-conll02-dutch''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll02-dutch/resolve/main/tokenizer.json'''
        ),
        '''xlm-roberta-large-finetuned-conll02-spanish''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll02-spanish/resolve/main/tokenizer.json'''
        ),
        '''xlm-roberta-large-finetuned-conll03-english''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll03-english/resolve/main/tokenizer.json'''
        ),
        '''xlm-roberta-large-finetuned-conll03-german''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll03-german/resolve/main/tokenizer.json'''
        ),
    },
}

UpperCAmelCase__  :	int     ={
    '''xlm-roberta-base''': 5_12,
    '''xlm-roberta-large''': 5_12,
    '''xlm-roberta-large-finetuned-conll02-dutch''': 5_12,
    '''xlm-roberta-large-finetuned-conll02-spanish''': 5_12,
    '''xlm-roberta-large-finetuned-conll03-english''': 5_12,
    '''xlm-roberta-large-finetuned-conll03-german''': 5_12,
}

class 					__A   (  a			):
				__A            = VOCAB_FILES_NAMES
				__A            = PRETRAINED_VOCAB_FILES_MAP
				__A            = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				__A            = ["""input_ids""", """attention_mask"""]
				__A            = XLMRobertaTokenizer


				def __init__(						self      ,       UpperCAmelCase_=None      ,       UpperCAmelCase_=None      ,       UpperCAmelCase_="<s>"      ,       UpperCAmelCase_="</s>"      ,       UpperCAmelCase_="</s>"      ,       UpperCAmelCase_="<s>"      ,       UpperCAmelCase_="<unk>"      ,       UpperCAmelCase_="<pad>"      ,       UpperCAmelCase_="<mask>"      ,       **UpperCAmelCase_      ,       ):
								# Mask token behave like a normal word, i.e. include the space before it
								lowerCamelCase		       =AddedToken(UpperCAmelCase_      ,       lstrip=UpperCAmelCase_      ,       rstrip=UpperCAmelCase_							) if isinstance(UpperCAmelCase_      ,       UpperCAmelCase_							) else mask_token

								super().__init__(
								    UpperCAmelCase_      ,       tokenizer_file=UpperCAmelCase_      ,       bos_token=UpperCAmelCase_      ,       eos_token=UpperCAmelCase_      ,       sep_token=UpperCAmelCase_      ,       cls_token=UpperCAmelCase_      ,       unk_token=UpperCAmelCase_      ,       pad_token=UpperCAmelCase_      ,       mask_token=UpperCAmelCase_      ,       **UpperCAmelCase_      ,       )

								lowerCamelCase		       =vocab_file
								lowerCamelCase		       =False if not self.vocab_file else True


				def _snake_case				(						self      ,       UpperCAmelCase_      ,       UpperCAmelCase_ = None							):

								if token_ids_a is None:
												return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
								lowerCamelCase		       =[self.cls_token_id]
								lowerCamelCase		       =[self.sep_token_id]
								return cls + token_ids_a + sep + sep + token_ids_a + sep


				def _snake_case				(						self      ,       UpperCAmelCase_      ,       UpperCAmelCase_ = None							):
								lowerCamelCase		       =[self.sep_token_id]
								lowerCamelCase		       =[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]


				def _snake_case				(						self      ,       UpperCAmelCase_      ,       UpperCAmelCase_ = None							):
								if not self.can_save_slow_tokenizer:
												raise ValueError(
												    """Your fast tokenizer does not have the necessary information to save the vocabulary for a slow """
												    """tokenizer."""							)

								if not os.path.isdir(UpperCAmelCase_							):
												logger.error(f"""Vocabulary path ({save_directory}) should be a directory."""							)
												return
								lowerCamelCase		       =os.path.join(
								    UpperCAmelCase_      ,       (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]							)

								if os.path.abspath(self.vocab_file							) != os.path.abspath(UpperCAmelCase_							):
												copyfile(self.vocab_file      ,       UpperCAmelCase_							)

								return (out_vocab_file,)





