import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


UpperCAmelCase_    		=				logging.getLogger(__name__)

require_version('''pytorch_lightning>=1.0.4''')

UpperCAmelCase_    		=				{
    'base': AutoModel,
    'sequence-classification': AutoModelForSequenceClassification,
    'question-answering': AutoModelForQuestionAnswering,
    'pretraining': AutoModelForPreTraining,
    'token-classification': AutoModelForTokenClassification,
    'language-modeling': AutoModelWithLMHead,
    'summarization': AutoModelForSeqaSeqLM,
    'translation': AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
UpperCAmelCase_    		=				{
    'linear': get_linear_schedule_with_warmup,
    'cosine': get_cosine_schedule_with_warmup,
    'cosine_w_restarts': get_cosine_with_hard_restarts_schedule_with_warmup,
    'polynomial': get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
UpperCAmelCase_    		=				sorted(arg_to_scheduler.keys())
UpperCAmelCase_    		=				'{' + ', '.join(arg_to_scheduler_choices) + '}'





class   lowerCAmelCase    (       pl.LightningModule      ):



     def __init__(       self					,   lowerCAmelCase__					,   lowerCAmelCase__=None					,   lowerCAmelCase__="base"					,   lowerCAmelCase__=None					,   lowerCAmelCase__=None					,   lowerCAmelCase__=None					,   **lowerCAmelCase__					,   ):
            super().__init__()
            # TODO: move to self.save_hyperparameters()
            # self.save_hyperparameters()
            # can also expand arguments into trainer signature for easier reading

            self.save_hyperparameters(__UpperCamelCase							)
            _A=						0
            _A=						Path(self.hparams.output_dir							)
            _A=						self.hparams.cache_dir if self.hparams.cache_dir else None
            if config is None:
                   _A=						AutoConfig.from_pretrained(
                       self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path					,   **({'num_labels': num_labels} if num_labels is not None else {})					,   cache_dir=__UpperCamelCase					,   **__UpperCamelCase					,   )
            else:
                   _A=						config

            _A=						('encoder_layerdrop', 'decoder_layerdrop', 'dropout', 'attention_dropout')
            for p in extra_model_params:
                   if getattr(self.hparams					,   __UpperCamelCase					,   __UpperCamelCase							):
                          assert hasattr(self.config					,   __UpperCamelCase							), f"model config doesn\'t have a `{p}` attribute"
                          setattr(self.config					,   __UpperCamelCase					,   getattr(self.hparams					,   __UpperCamelCase							)							)

            if tokenizer is None:
                   _A=						AutoTokenizer.from_pretrained(
                       self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path					,   cache_dir=__UpperCamelCase					,   )
            else:
                   _A=						tokenizer
            _A=						MODEL_MODES[mode]
            if model is None:
                   _A=						self.model_type.from_pretrained(
                       self.hparams.model_name_or_path					,   from_tf=bool('.ckpt' in self.hparams.model_name_or_path							)					,   config=self.config					,   cache_dir=__UpperCamelCase					,   )
            else:
                   _A=						model



     def       a__	(       self					,   *lowerCAmelCase__					,   **lowerCAmelCase__							):
            _A=						self.model_type.from_pretrained(*__UpperCamelCase					,   **__UpperCamelCase							)



     def       a__	(       self							):
            _A=						arg_to_scheduler[self.hparams.lr_scheduler]
            _A=						get_schedule_func(
                self.opt					,   num_warmup_steps=self.hparams.warmup_steps					,   num_training_steps=self.total_steps()							)
            _A=						{'scheduler': scheduler, 'interval': 'step', 'frequency': 1}
            return scheduler



     def       a__	(       self							):
            _A=						self.model
            _A=						['bias', 'LayerNorm.weight']
            _A=						[
                {
                    'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay							)],
                    'weight_decay': self.hparams.weight_decay,
                },
                {
                    'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay							)],
                    'weight_decay': 0.0,
                },
            ]
            if self.hparams.adafactor:
                   _A=						Adafactor(
                       __UpperCamelCase					,   lr=self.hparams.learning_rate					,   scale_parameter=__UpperCamelCase					,   relative_step=__UpperCamelCase							)

            else:
                   _A=						AdamW(
                       __UpperCamelCase					,   lr=self.hparams.learning_rate					,   eps=self.hparams.adam_epsilon							)
            _A=						optimizer

            _A=						self.get_lr_scheduler()

            return [optimizer], [scheduler]



     def       a__	(       self					,   lowerCAmelCase__					,   lowerCAmelCase__							):
            return self.validation_step(__UpperCamelCase					,   __UpperCamelCase							)



     def       a__	(       self					,   lowerCAmelCase__							):
            return self.validation_end(__UpperCamelCase							)



     def       a__	(       self							):
            _A=						max(1					,   self.hparams.gpus							)  # TODO: consider num_tpu_cores
            _A=						self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
            return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs



     def       a__	(       self					,   lowerCAmelCase__							):
            if mode == "test":
                   _A=						len(self.test_dataloader().dataset							)
            else:
                   _A=						self.get_dataloader('train'					,   self.hparams.train_batch_size					,   shuffle=__UpperCamelCase							)
                   _A=						len(self.train_dataloader().dataset							)



     def       a__	(       self					,   lowerCAmelCase__					,   lowerCAmelCase__					,   lowerCAmelCase__ = False							):
            raise NotImplementedError('You must implement this for your task'							)



     def       a__	(       self							):
            return self.train_loader



     def       a__	(       self							):
            return self.get_dataloader('dev'					,   self.hparams.eval_batch_size					,   shuffle=__UpperCamelCase							)



     def       a__	(       self							):
            return self.get_dataloader('test'					,   self.hparams.eval_batch_size					,   shuffle=__UpperCamelCase							)



     def       a__	(       self					,   lowerCAmelCase__							):
            return os.path.join(
                self.hparams.data_dir					,   'cached_{}_{}_{}'.format(
                    __UpperCamelCase					,   list(filter(__UpperCamelCase					,   self.hparams.model_name_or_path.split('/'							)							)							).pop()					,   str(self.hparams.max_seq_length							)					,   )					,   )



     @pl.utilities.rank_zero_only
     def       a__	(       self					,   lowerCAmelCase__							):
            _A=						self.output_dir.joinpath('best_tfmr'							)
            _A=						self.step_count
            self.model.save_pretrained(__UpperCamelCase							)
            self.tokenizer.save_pretrained(__UpperCamelCase							)





     @staticmethod
     def       a__	(       lowerCAmelCase__					,   lowerCAmelCase__							):
            parser.add_argument(
                '--model_name_or_path'					,   default=__UpperCamelCase					,   type=__UpperCamelCase					,   required=__UpperCamelCase					,   help='Path to pretrained model or model identifier from huggingface.co/models'					,   )
            parser.add_argument(
                '--config_name'					,   default=''					,   type=__UpperCamelCase					,   help='Pretrained config name or path if not the same as model_name'							)
            parser.add_argument(
                '--tokenizer_name'					,   default=__UpperCamelCase					,   type=__UpperCamelCase					,   help='Pretrained tokenizer name or path if not the same as model_name'					,   )
            parser.add_argument(
                '--cache_dir'					,   default=''					,   type=__UpperCamelCase					,   help='Where do you want to store the pre-trained models downloaded from huggingface.co'					,   )
            parser.add_argument(
                '--encoder_layerdrop'					,   type=__UpperCamelCase					,   help='Encoder layer dropout probability (Optional). Goes into model.config'					,   )
            parser.add_argument(
                '--decoder_layerdrop'					,   type=__UpperCamelCase					,   help='Decoder layer dropout probability (Optional). Goes into model.config'					,   )
            parser.add_argument(
                '--dropout'					,   type=__UpperCamelCase					,   help='Dropout probability (Optional). Goes into model.config'					,   )
            parser.add_argument(
                '--attention_dropout'					,   type=__UpperCamelCase					,   help='Attention dropout probability (Optional). Goes into model.config'					,   )
            parser.add_argument('--learning_rate'					,   default=5E-5					,   type=__UpperCamelCase					,   help='The initial learning rate for Adam.'							)
            parser.add_argument(
                '--lr_scheduler'					,   default='linear'					,   choices=__UpperCamelCase					,   metavar=__UpperCamelCase					,   type=__UpperCamelCase					,   help='Learning rate scheduler'					,   )
            parser.add_argument('--weight_decay'					,   default=0.0					,   type=__UpperCamelCase					,   help='Weight decay if we apply some.'							)
            parser.add_argument('--adam_epsilon'					,   default=1E-8					,   type=__UpperCamelCase					,   help='Epsilon for Adam optimizer.'							)
            parser.add_argument('--warmup_steps'					,   default=0					,   type=__UpperCamelCase					,   help='Linear warmup over warmup_steps.'							)
            parser.add_argument('--num_workers'					,   default=4					,   type=__UpperCamelCase					,   help='kwarg passed to DataLoader'							)
            parser.add_argument('--num_train_epochs'					,   dest='max_epochs'					,   default=3					,   type=__UpperCamelCase							)
            parser.add_argument('--train_batch_size'					,   default=32					,   type=__UpperCamelCase							)
            parser.add_argument('--eval_batch_size'					,   default=32					,   type=__UpperCamelCase							)
            parser.add_argument('--adafactor'					,   action='store_true'							)





class   lowerCAmelCase    (       pl.Callback      ):



     def       a__	(       self					,   lowerCAmelCase__					,   lowerCAmelCase__							):
            _A=						trainer.lr_schedulers[0]['scheduler']
            _A=						{f"lr_group_{i}": lr for i, lr in enumerate(lr_scheduler.get_lr()							)}
            pl_module.logger.log_metrics(__UpperCamelCase							)



     def       a__	(       self					,   lowerCAmelCase__					,   lowerCAmelCase__							):
            rank_zero_info('***** Validation results *****'							)
            _A=						trainer.callback_metrics
            # Log results
            for key in sorted(__UpperCamelCase							):
                   if key not in ["log", "progress_bar"]:
                          rank_zero_info('{} = {}\n'.format(__UpperCamelCase					,   str(metrics[key]							)							)							)





     def       a__	(       self					,   lowerCAmelCase__					,   lowerCAmelCase__							):
            rank_zero_info('***** Test results *****'							)
            _A=						trainer.callback_metrics
            # Log and save results to file
            _A=						os.path.join(pl_module.hparams.output_dir					,   'test_results.txt'							)
            with open(__UpperCamelCase					,   'w'							) as writer:
                   for key in sorted(__UpperCamelCase							):
                          if key not in ["log", "progress_bar"]:
                                 rank_zero_info('{} = {}\n'.format(__UpperCamelCase					,   str(metrics[key]							)							)							)
                                 writer.write('{} = {}\n'.format(__UpperCamelCase					,   str(metrics[key]							)							)							)


def 				UpperCamelCase					(     lowerCAmelCase_							,       lowerCAmelCase_ )		->    List[Any]:




       '''simple docstring'''






       parser.add_argument(
           '--output_dir'							,       default=lowerCAmelCase_							,       type=lowerCAmelCase_							,       required=lowerCAmelCase_							,       help='The output directory where the model predictions and checkpoints will be written.'							,       )
       parser.add_argument(
           '--fp16'							,       action='store_true'							,       help='Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit'							,       )

       parser.add_argument(
           '--fp16_opt_level'							,       type=lowerCAmelCase_							,       default='O2'							,       help=(
               'For fp16: Apex AMP optimization level selected in [\'O0\', \'O1\', \'O2\', and \'O3\'].'
               'See details at https://nvidia.github.io/apex/amp.html'
           )							,       )
       parser.add_argument('--n_tpu_cores'							,       dest='tpu_cores'							,       type=lowerCAmelCase_ )
       parser.add_argument('--max_grad_norm'							,       dest='gradient_clip_val'							,       default=1.0							,       type=lowerCAmelCase_							,       help='Max gradient norm' )
       parser.add_argument('--do_train'							,       action='store_true'							,       help='Whether to run training.' )
       parser.add_argument('--do_predict'							,       action='store_true'							,       help='Whether to run predictions on the test set.' )
       parser.add_argument(
           '--gradient_accumulation_steps'							,       dest='accumulate_grad_batches'							,       type=lowerCAmelCase_							,       default=1							,       help='Number of updates steps to accumulate before performing a backward/update pass.'							,       )
       parser.add_argument('--seed'							,       type=lowerCAmelCase_							,       default=42							,       help='random seed for initialization' )
       parser.add_argument(
           '--data_dir'							,       default=lowerCAmelCase_							,       type=lowerCAmelCase_							,       required=lowerCAmelCase_							,       help='The input data dir. Should contain the training files for the CoNLL-2003 NER task.'							,       )





def 				UpperCamelCase					(     lowerCAmelCase_							,       lowerCAmelCase_							,       lowerCAmelCase_=None							,       lowerCAmelCase_=True							,       lowerCAmelCase_=[]							,       lowerCAmelCase_=None							,       lowerCAmelCase_=None							,       **lowerCAmelCase_							,       )		->    Optional[int]:




       '''simple docstring'''






       pl.seed_everything(args.seed )

       # init model
       _A=						Path(model.hparams.output_dir )
       odir.mkdir(exist_ok=lowerCAmelCase_ )

       # add custom checkpoints
       if checkpoint_callback is None:
              _A=						pl.callbacks.ModelCheckpoint(
                  filepath=args.output_dir							,       prefix='checkpoint'							,       monitor='val_loss'							,       mode='min'							,       save_top_k=1 )
       if early_stopping_callback:
              extra_callbacks.append(lowerCAmelCase_ )
       if logging_callback is None:
              _A=						LoggingCallback()

       _A=						{}

       # TODO: remove with PyTorch 1.6 since pl uses native amp
       if args.fpaa:
              _A=						16
              _A=						args.fpaa_opt_level

       if args.gpus > 1:
              _A=						'ddp'

       _A=						args.accumulate_grad_batches
       _A=						extra_train_kwargs.get('accelerator'							,       lowerCAmelCase_ )
       _A=						extra_train_kwargs.get('profiler'							,       lowerCAmelCase_ )

       _A=						pl.Trainer.from_argparse_args(
           lowerCAmelCase_							,       weights_summary=lowerCAmelCase_							,       callbacks=[logging_callback] + extra_callbacks							,       logger=lowerCAmelCase_							,       checkpoint_callback=lowerCAmelCase_							,       **lowerCAmelCase_							,       )

       if args.do_train:
              trainer.fit(lowerCAmelCase_ )

       return trainer




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




