


'''simple docstring'''




#
# Sort large text files in a minimum amount of memory
#
import argparse
import os





class _snake_case :
							"""simple docstring"""





							__a				       =  '''block_{0}.dat'''







							def __init__( self      ,					__A						)  ->int:



											"""simple docstring"""







											__lowerCamelCase		:   Dict           =       filename
											__lowerCamelCase		:   Union[str, Any]           =       []







							def 		__SCREAMING_SNAKE_CASE( self      ,					__A      ,					__A						)  ->int:



											"""simple docstring"""







											__lowerCamelCase		:   Optional[int]           =       self.BLOCK_FILENAME_FORMAT.format(__A						)
											with open(__A      ,					"w"						) as file:
															file.write(__A						)
											self.block_filenames.append(__A						)







							def 		__SCREAMING_SNAKE_CASE( self						)  ->int:



											"""simple docstring"""







											return self.block_filenames







							def 		__SCREAMING_SNAKE_CASE( self      ,					__A      ,					__A=None						)  ->Dict:



											"""simple docstring"""







											__lowerCamelCase		:   Optional[int]           =       0
											with open(self.filename						) as file:
															while True:
																			__lowerCamelCase		:   Union[str, Any]           =       file.readlines(__A						)

																			if lines == []:
																							break

																			if sort_key is None:
																							lines.sort()
																			else:
																							lines.sort(key=__A						)

																			self.write_block("".join(__A						)      ,					__A						)
																			i += 1







							def 		__SCREAMING_SNAKE_CASE( self						)  ->List[Any]:



											"""simple docstring"""







											map(os.remove      ,					self.block_filenames						)





class _snake_case :
							"""simple docstring"""





							def 		__SCREAMING_SNAKE_CASE( self      ,					__A						)  ->Optional[int]:



											"""simple docstring"""







											__lowerCamelCase		:   Union[str, Any]           =       -1
											__lowerCamelCase		:   Optional[Any]           =       None

											for i in range(len(__A						)						):
															if min_str is None or choices[i] < min_str:
																			__lowerCamelCase		:   Tuple           =       i

											return min_index





class _snake_case :
							"""simple docstring"""





							def __init__( self      ,					__A						)  ->List[Any]:



											"""simple docstring"""







											__lowerCamelCase		:   Tuple           =       files
											__lowerCamelCase		:   Tuple           =       set()
											__lowerCamelCase		:   Optional[Any]           =       len(__A						)
											__lowerCamelCase		:   List[str]           =       {i: None for i in range(self.num_buffers						)}







							def 		__SCREAMING_SNAKE_CASE( self						)  ->Any:



											"""simple docstring"""







											return {
											    i: self.buffers[i] for i in range(self.num_buffers						) if i not in self.empty
											}







							def 		__SCREAMING_SNAKE_CASE( self						)  ->Dict:



											"""simple docstring"""







											for i in range(self.num_buffers						):
															if self.buffers[i] is None and i not in self.empty:
																			__lowerCamelCase		:   Tuple           =       self.files[i].readline()

																			if self.buffers[i] == "":
																							self.empty.add(__A						)
																							self.files[i].close()

											if len(self.empty						) == self.num_buffers:
															return False

											return True







							def 		__SCREAMING_SNAKE_CASE( self      ,					__A						)  ->Any:



											"""simple docstring"""







											__lowerCamelCase		:   Tuple           =       self.buffers[index]
											__lowerCamelCase		:   List[Any]           =       None

											return value





class _snake_case :
							"""simple docstring"""





							def __init__( self      ,					__A						)  ->List[str]:



											"""simple docstring"""







											__lowerCamelCase		:   List[Any]           =       merge_strategy







							def 		__SCREAMING_SNAKE_CASE( self      ,					__A      ,					__A      ,					__A						)  ->Optional[int]:



											"""simple docstring"""







											__lowerCamelCase		:   List[str]           =       FilesArray(self.get_file_handles(__A      ,					__A						)						)
											with open(__A      ,					"w"      ,					__A						) as outfile:
															while buffers.refresh():
																			__lowerCamelCase		:   Optional[Any]           =       self.merge_strategy.select(buffers.get_dict()						)
																			outfile.write(buffers.unshift(__A						)						)







							def 		__SCREAMING_SNAKE_CASE( self      ,					__A      ,					__A						)  ->str:



											"""simple docstring"""







											__lowerCamelCase		:   Union[str, Any]           =       {}

											for i in range(len(__A						)						):
															__lowerCamelCase		:   Tuple           =       open(filenames[i]      ,					"r"      ,					__A						)  # noqa: UP015

											return files





class _snake_case :
							"""simple docstring"""





							def __init__( self      ,					__A						)  ->str:



											"""simple docstring"""







											__lowerCamelCase		:   Tuple           =       block_size







							def 		__SCREAMING_SNAKE_CASE( self      ,					__A      ,					__A=None						)  ->List[str]:



											"""simple docstring"""







											__lowerCamelCase		:   Optional[int]           =       self.get_number_blocks(__A      ,					self.block_size						)
											__lowerCamelCase		:   Optional[int]           =       FileSplitter(__A						)
											splitter.split(self.block_size      ,					__A						)

											__lowerCamelCase		:   List[str]           =       FileMerger(NWayMerge()						)
											__lowerCamelCase		:   str           =       self.block_size / (num_blocks + 1)
											merger.merge(splitter.get_block_filenames()      ,					filename + ".out"      ,					__A						)

											splitter.cleanup()







							def 		__SCREAMING_SNAKE_CASE( self      ,					__A      ,					__A						)  ->List[str]:



											"""simple docstring"""







											return (os.stat(__A						).st_size / block_size) + 1





def 				a_							(    lowercase     )      ->List[Any]:






				'''simple docstring'''

				if string[-1].lower() == "k":
								return int(string[:-1]     ) * 1024
				elif string[-1].lower() == "m":
								return int(string[:-1]     ) * 1024 * 1024
				elif string[-1].lower() == "g":
								return int(string[:-1]     ) * 1024 * 1024 * 1024
				else:
								return int(lowercase     )





def 				a_							(    )      ->Optional[Any]:






				'''simple docstring'''

				__lowerCamelCase		:   List[str]           =       argparse.ArgumentParser()
				parser.add_argument(
				    "-m"				,    "--mem"				,    help="amount of memory to use for sorting"				,    default="100M"     )
				parser.add_argument(
				    "filename"				,    metavar="<filename>"				,    nargs=1				,    help="name of file to sort"     )
				__lowerCamelCase		:   List[str]           =       parser.parse_args()

				__lowerCamelCase		:   Tuple           =       ExternalSort(parse_memory(args.mem     )     )
				sorter.sort(args.filename[0]     )


if __name__ == "__main__":
				main()









'''simple docstring'''




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


__lowercase	:   str   		=					logging.getLogger(__name__)

require_version("""pytorch_lightning>=1.0.4""")

__lowercase	:   List[str]   		=					{
    """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
__lowercase	:   List[str]   		=					{
    """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
}
__lowercase	:   Optional[Any]   		=					sorted(arg_to_scheduler.keys())
__lowercase	:   Union[str, Any]   		=					"""{""" + """, """.join(arg_to_scheduler_choices) + """}"""





class _snake_case (			pl.LightningModule ):
							"""simple docstring"""





							def __init__( self      ,					__A      ,					__A=None      ,					__A="base"      ,					__A=None      ,					__A=None      ,					__A=None      ,					**__A      ,					)  ->Any:



											"""simple docstring"""







											super().__init__()
											# TODO: move to self.save_hyperparameters()
											# self.save_hyperparameters()
											# can also expand arguments into trainer signature for easier reading

											self.save_hyperparameters(__A						)
											__lowerCamelCase		:   str           =       0
											__lowerCamelCase		:   int           =       Path(self.hparams.output_dir						)
											__lowerCamelCase		:   Tuple           =       self.hparams.cache_dir if self.hparams.cache_dir else None
											if config is None:
															__lowerCamelCase		:   str           =       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=__A      ,					**__A      ,					)
											else:
															__lowerCamelCase		:   PretrainedConfig           =       config

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

											if tokenizer is None:
															__lowerCamelCase		:   Optional[int]           =       AutoTokenizer.from_pretrained(
															    self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path      ,					cache_dir=__A      ,					)
											else:
															__lowerCamelCase		:   PreTrainedTokenizer           =       tokenizer
											__lowerCamelCase		:   List[str]           =       MODEL_MODES[mode]
											if model is None:
															__lowerCamelCase		:   int           =       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=__A      ,					)
											else:
															__lowerCamelCase		:   Tuple           =       model







							def 		__SCREAMING_SNAKE_CASE( self      ,					*__A      ,					**__A						)  ->Any:



											"""simple docstring"""







											__lowerCamelCase		:   str           =       self.model_type.from_pretrained(*__A      ,					**__A						)







							def 		__SCREAMING_SNAKE_CASE( self						)  ->int:



											"""simple docstring"""







											__lowerCamelCase		:   Tuple           =       arg_to_scheduler[self.hparams.lr_scheduler]
											__lowerCamelCase		:   Optional[Any]           =       get_schedule_func(
											    self.opt      ,					num_warmup_steps=self.hparams.warmup_steps      ,					num_training_steps=self.total_steps()						)
											__lowerCamelCase		:   List[Any]           =       {"scheduler": scheduler, "interval": "step", "frequency": 1}
											return scheduler







							def 		__SCREAMING_SNAKE_CASE( self						)  ->Optional[Any]:



											"""simple docstring"""







											__lowerCamelCase		:   Dict           =       self.model
											__lowerCamelCase		:   List[Any]           =       ["bias", "LayerNorm.weight"]
											__lowerCamelCase		:   int           =       [
											    {
											        "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:
															__lowerCamelCase		:   Any           =       Adafactor(
															    __A      ,					lr=self.hparams.learning_rate      ,					scale_parameter=__A      ,					relative_step=__A						)

											else:
															__lowerCamelCase		:   str           =       AdamW(
															    __A      ,					lr=self.hparams.learning_rate      ,					eps=self.hparams.adam_epsilon						)
											__lowerCamelCase		:   Optional[int]           =       optimizer

											__lowerCamelCase		:   List[str]           =       self.get_lr_scheduler()

											return [optimizer], [scheduler]







							def 		__SCREAMING_SNAKE_CASE( self      ,					__A      ,					__A						)  ->Optional[int]:



											"""simple docstring"""







											return self.validation_step(__A      ,					__A						)







							def 		__SCREAMING_SNAKE_CASE( self      ,					__A						)  ->Optional[int]:



											"""simple docstring"""







											return self.validation_end(__A						)







							def 		__SCREAMING_SNAKE_CASE( self						)  ->int:



											"""simple docstring"""







											__lowerCamelCase		:   Union[str, Any]           =       max(1      ,					self.hparams.gpus						)  # TODO: consider num_tpu_cores
											__lowerCamelCase		:   Optional[Any]           =       self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
											return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs







							def 		__SCREAMING_SNAKE_CASE( self      ,					__A						)  ->Union[str, Any]:



											"""simple docstring"""







											if mode == "test":
															__lowerCamelCase		:   str           =       len(self.test_dataloader().dataset						)
											else:
															__lowerCamelCase		:   List[str]           =       self.get_dataloader("train"      ,					self.hparams.train_batch_size      ,					shuffle=__A						)
															__lowerCamelCase		:   List[str]           =       len(self.train_dataloader().dataset						)







							def 		__SCREAMING_SNAKE_CASE( self      ,					__A      ,					__A      ,					__A = False						)  ->Any:



											"""simple docstring"""







											raise NotImplementedError("You must implement this for your task"						)







							def 		__SCREAMING_SNAKE_CASE( self						)  ->Optional[Any]:



											"""simple docstring"""







											return self.train_loader







							def 		__SCREAMING_SNAKE_CASE( self						)  ->Optional[Any]:



											"""simple docstring"""







											return self.get_dataloader("dev"      ,					self.hparams.eval_batch_size      ,					shuffle=__A						)







							def 		__SCREAMING_SNAKE_CASE( self						)  ->Dict:



											"""simple docstring"""







											return self.get_dataloader("test"      ,					self.hparams.eval_batch_size      ,					shuffle=__A						)







							def 		__SCREAMING_SNAKE_CASE( self      ,					__A						)  ->Dict:



											"""simple docstring"""







											return os.path.join(
											    self.hparams.data_dir      ,					"cached_{}_{}_{}".format(
											        __A      ,					list(filter(__A      ,					self.hparams.model_name_or_path.split("/"						)						)						).pop()      ,					str(self.hparams.max_seq_length						)      ,					)      ,					)







							@pl.utilities.rank_zero_only
							def 		__SCREAMING_SNAKE_CASE( self      ,					__A						)  ->None:



											"""simple docstring"""







											__lowerCamelCase		:   int           =       self.output_dir.joinpath("best_tfmr"						)
											__lowerCamelCase		:   str           =       self.step_count
											self.model.save_pretrained(__A						)
											self.tokenizer.save_pretrained(__A						)







							@staticmethod
							def 		__SCREAMING_SNAKE_CASE( __A      ,					__A						)  ->str:



											"""simple docstring"""







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





class _snake_case (			pl.Callback ):
							"""simple docstring"""





							def 		__SCREAMING_SNAKE_CASE( self      ,					__A      ,					__A						)  ->Union[str, Any]:



											"""simple docstring"""







											__lowerCamelCase		:   List[Any]           =       trainer.lr_schedulers[0]["scheduler"]
											__lowerCamelCase		:   str           =       {f'lr_group_{i}': lr for i, lr in enumerate(lr_scheduler.get_lr()						)}
											pl_module.logger.log_metrics(__A						)







							def 		__SCREAMING_SNAKE_CASE( self      ,					__A      ,					__A						)  ->List[Any]:



											"""simple docstring"""







											rank_zero_info("***** Validation results *****"						)
											__lowerCamelCase		:   Dict           =       trainer.callback_metrics
											# Log results
											for key in sorted(__A						):
															if key not in ["log", "progress_bar"]:
																			rank_zero_info("{} = {}\n".format(__A      ,					str(metrics[key]						)						)						)







							def 		__SCREAMING_SNAKE_CASE( self      ,					__A      ,					__A						)  ->Optional[Any]:



											"""simple docstring"""







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





def 				a_							(    lowercase				,    lowercase     )      ->None:






				'''simple docstring'''

				parser.add_argument(
				    "--output_dir"				,    default=lowercase				,    type=lowercase				,    required=lowercase				,    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=lowercase				,    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=lowercase     )
				parser.add_argument("--max_grad_norm"				,    dest="gradient_clip_val"				,    default=1.0				,    type=lowercase				,    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=lowercase				,    default=1				,    help="Number of updates steps to accumulate before performing a backward/update pass."				,    )
				parser.add_argument("--seed"				,    type=lowercase				,    default=42				,    help="random seed for initialization"     )
				parser.add_argument(
				    "--data_dir"				,    default=lowercase				,    type=lowercase				,    required=lowercase				,    help="The input data dir. Should contain the training files for the CoNLL-2003 NER task."				,    )





def 				a_							(    lowercase				,    lowercase				,    lowercase=None				,    lowercase=True				,    lowercase=[]				,    lowercase=None				,    lowercase=None				,    **lowercase				,    )      ->List[Any]:






				'''simple docstring'''

				pl.seed_everything(args.seed     )

				# init model
				__lowerCamelCase		:   str           =       Path(model.hparams.output_dir     )
				odir.mkdir(exist_ok=lowercase     )

				# add custom checkpoints
				if checkpoint_callback is None:
								__lowerCamelCase		:   Optional[Any]           =       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(lowercase     )
				if logging_callback is None:
								__lowerCamelCase		:   List[Any]           =       LoggingCallback()

				__lowerCamelCase		:   int           =       {}

				# TODO: remove with PyTorch 1.6 since pl uses native amp
				if args.fpaa:
								__lowerCamelCase		:   List[Any]           =       16
								__lowerCamelCase		:   Optional[int]           =       args.fpaa_opt_level

				if args.gpus > 1:
								__lowerCamelCase		:   Tuple           =       "ddp"

				__lowerCamelCase		:   str           =       args.accumulate_grad_batches
				__lowerCamelCase		:   Any           =       extra_train_kwargs.get("accelerator"				,    lowercase     )
				__lowerCamelCase		:   Tuple           =       extra_train_kwargs.get("profiler"				,    lowercase     )

				__lowerCamelCase		:   Any           =       pl.Trainer.from_argparse_args(
				    lowercase				,    weights_summary=lowercase				,    callbacks=[logging_callback] + extra_callbacks				,    logger=lowercase				,    checkpoint_callback=lowercase				,    **lowercase				,    )

				if args.do_train:
								trainer.fit(lowercase     )

				return trainer





