


'''simple docstring'''



import argparse
import logging

import datasets
from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors
from tokenizers.models import Unigram
from tokenizers.trainers import UnigramTrainer

from transformers import AlbertTokenizerFast


__UpperCAmelCase			:int 				=   logging.getLogger(__name__)




def 	_a	(						):





     '''simple docstring'''






     __UpperCAmelCase	: Optional[Any]  							=			argparse.ArgumentParser(description='''Train a unigram tokenizer on the wikitext dataset.'''  )
     parser.add_argument(
         '''--dataset_name'''			,					type=__UpperCamelCase			,					default='''wikitext'''			,					help='''Name of the training. Explore datasets at: hf.co/datasets.'''			,					)
     parser.add_argument(
         '''--dataset_config'''			,					type=__UpperCamelCase			,					default='''wikitext-103-raw-v1'''			,					help='''Configuration name of the dataset.'''  )
     parser.add_argument(
         '''--batch_size'''			,					type=__UpperCamelCase			,					default=1000			,					help='''Batch size during training.'''			,					)
     parser.add_argument(
         '''--vocab_size'''			,					type=__UpperCamelCase			,					default=10048			,					help='''Size of the desired vocabulary.'''			,					)
     parser.add_argument(
         '''--limit'''			,					default=__UpperCamelCase			,					type=__UpperCamelCase			,					help='''Limit the number of shards (used for debugging).'''			,					)
     parser.add_argument(
         '''--export_to_hub'''			,					action='''store_true'''			,					)

     __UpperCAmelCase	: int  							=			parser.parse_args()
     return args




def 	_a	(						_lowercase					:	Union[str, Any]  ):





     '''simple docstring'''






     __UpperCAmelCase	: Dict  							=			datasets.load_dataset(args.dataset_name			,					args.dataset_config			,					split='''train'''  )

     if args.limit is not None:
          __UpperCAmelCase	: int  							=			min(len(__UpperCamelCase  )			,					args.limit  )
          __UpperCAmelCase	: Union[str, Any]  							=			dataset.select(range(__UpperCamelCase  )  )
          logger.info(F'Limiting the dataset to {args.limit} entries.'  )

     def batch_iterator():
          for i in range(0			,					len(__UpperCamelCase  )			,					args.batch_size  ):
               yield dataset[i : i + args.batch_size]["text"]

    # Prepare the tokenizer.
     __UpperCAmelCase	: Tuple  							=			Tokenizer(Unigram()  )
     __UpperCAmelCase	: Dict  							=			normalizers.Sequence([normalizers.Replace('''``'''			,					'''\"'''  ), normalizers.Replace('''\'\''''			,					'''\"'''  )]  )
     __UpperCAmelCase	: Optional[Any]  							=			pre_tokenizers.Metaspace()

     # Prepare the trainer.
     __UpperCAmelCase	: Optional[Any]  							=			UnigramTrainer(
         unk_token='''<unk>'''			,					special_tokens=['''[CLS]''', '''[SEP]''', '''<unk>''', '''<pad>''', '''[MASK]''']			,					vocab_size=args.vocab_size			,					)

     logger.info('''Training the tokenizer.'''  )
     tokenizer.train_from_iterator(batch_iterator()			,					trainer=__UpperCamelCase  )
     logger.info('''Tokenizer training complete!'''  )

     __UpperCAmelCase	: List[Any]  							=			tokenizer.token_to_id('''[CLS]'''  )
     __UpperCAmelCase	: List[str]  							=			tokenizer.token_to_id('''[SEP]'''  )
     __UpperCAmelCase	: int  							=			processors.TemplateProcessing(
         single='''[CLS]:0 $A:0 [SEP]:0'''			,					pair='''[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1'''			,					special_tokens=[
             ('''[CLS]''', cls_token_id),
             ('''[SEP]''', sep_token_id),
         ]			,					)
     __UpperCAmelCase	: Union[str, Any]  							=			decoders.Metaspace()

     if args.export_to_hub:
          logger.info('''Exporting the trained tokenzier to Hub.'''  )
          __UpperCAmelCase	: Union[str, Any]  							=			AlbertTokenizerFast(tokenizer_object=__UpperCamelCase  )
          new_tokenizer.push_to_hub('''unigram-tokenizer-dataset'''  )


if __name__ == "__main__":
      __UpperCAmelCase			:Tuple 				=   parse_args()
      main(args)





import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only

from utils import save_json

def 					SCREAMING_SNAKE_CASE							(   __UpperCamelCase)			->     int:
     a      =      filter(lambda __UpperCamelCase: p.requires_grad		,					model.parameters())
     a      =      sum([np.prod(p.size()) for p in model_parameters])
     return params


lowercase__      :			Optional[int]           =					logging.getLogger(__name__)




class 			a__       (					pl.Callback      ):







     def 			lowerCAmelCase_				(	self  ,		A  ,		A )  ->    str:



          '''simple docstring'''



          a      =      {F'''lr_group_{i}''': param["lr"] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups )}
          pl_module.logger.log_metrics(A )







     @rank_zero_only
     def 			lowerCAmelCase_				(	self  ,		A  ,		A  ,		A  ,		A=True )  ->    None:



          '''simple docstring'''



          logger.info(F'''***** {type_path} results at step {trainer.global_step:05d} *****''' )
          a      =      trainer.callback_metrics
          trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ["log", "progress_bar", "preds"]} )
          # Log results
          a      =      Path(pl_module.hparams.output_dir )
          if type_path == "test":
               a      =      od / "test_results.txt"
               a      =      od / "test_generations.txt"
          else:
               # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
               # If people want this it will be easy enough to add back.
               a      =      od / F'''{type_path}_results/{trainer.global_step:05d}.txt'''
               a      =      od / F'''{type_path}_generations/{trainer.global_step:05d}.txt'''
               results_file.parent.mkdir(exist_ok=A )
               generations_file.parent.mkdir(exist_ok=A )
          with open(A  ,		"a+" ) as writer:
               for key in sorted(A ):
                    if key in ["log", "progress_bar", "preds"]:
                         continue
                    a      =      metrics[key]
                    if isinstance(A  ,		torch.Tensor ):
                         a      =      val.item()
                    a      =      F'''{key}: {val:.6f}\n'''
                    writer.write(A )

          if not save_generations:
               return

          if "preds" in metrics:
               a      =      "\n".join(metrics["preds"] )
               generations_file.open("w+" ).write(A )







     @rank_zero_only
     def 			lowerCAmelCase_				(	self  ,		A  ,		A )  ->    Any:



          '''simple docstring'''



          try:
               a      =      pl_module.model.model.num_parameters()
          except AttributeError:
               a      =      pl_module.model.num_parameters()

          a      =      count_trainable_parameters(A )
          # mp stands for million parameters
          trainer.logger.log_metrics({"n_params": npars, "mp": npars / 1e6, "grad_mp": n_trainable_pars / 1e6} )







     @rank_zero_only
     def 			lowerCAmelCase_				(	self  ,		A  ,		A )  ->    Optional[int]:



          '''simple docstring'''



          save_json(pl_module.metrics  ,		pl_module.metrics_save_path )
          return self._write_logs(A  ,		A  ,		"test" )







     @rank_zero_only
     def 			lowerCAmelCase_				(	self  ,		A  ,		A )  ->    Dict:



          '''simple docstring'''



          save_json(pl_module.metrics  ,		pl_module.metrics_save_path )
          # Uncommenting this will save val generations
          # return self._write_logs(trainer, pl_module, "valid")







def 					SCREAMING_SNAKE_CASE							(   __UpperCamelCase		,					__UpperCamelCase		,					__UpperCamelCase=1		,					__UpperCamelCase=False)			->     Optional[int]:
     if metric == "rouge2":
          a      =      "{val_avg_rouge2:.4f}-{step_count}"
     elif metric == "bleu":
          a      =      "{val_avg_bleu:.4f}-{step_count}"
     elif metric == "loss":
          a      =      "{val_avg_loss:.4f}-{step_count}"
     else:
          raise NotImplementedError(
              f'''seq2seq callbacks only support rouge2, bleu and loss, got {metric}, You can make your own by adding to'''
              " this function.")

     a      =      ModelCheckpoint(
         dirpath=__UpperCamelCase		,					filename=__UpperCamelCase		,					monitor=f'''val_{metric}'''		,					mode="min" if "loss" in metric else "max"		,					save_top_k=__UpperCamelCase		,					)
     return checkpoint_callback





def 					SCREAMING_SNAKE_CASE							(   __UpperCamelCase		,					__UpperCamelCase)			->     Union[str, Any]:
     return EarlyStopping(
         monitor=f'''val_{metric}'''		,					mode="min" if "loss" in metric else "max"		,					patience=__UpperCamelCase		,					verbose=__UpperCamelCase		,					)

