






"""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  							= logging.getLogger(__name__)

require_version("pytorch_lightning>=1.0.4")

lowercase  							= {
    '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  							= {
    '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  							= sorted(arg_to_scheduler.keys())
lowercase  							= '{' + ', '.join(arg_to_scheduler_choices) + '}'




class 			UpperCamelCase_				(      pl.LightningModule    ):



	'''simple docstring'''




	def __init__(	self	,      a	,      a=None	,      a="base"	,      a=None	,      a=None	,      a=None	,      **a	,      )	->   int:
								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 )
								snake_case_           =							0
								snake_case_           =							Path(self.hparams.output_dir )
								snake_case_           =							self.hparams.cache_dir if self.hparams.cache_dir else None
								if config is None:
															snake_case_           =							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:
															snake_case_           =							config

								snake_case_           =							('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:
															snake_case_           =							AutoTokenizer.from_pretrained(
															    self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path	,      cache_dir=__UpperCamelCase	,      )
								else:
															snake_case_           =							tokenizer
								snake_case_           =							MODEL_MODES[mode]
								if model is None:
															snake_case_           =							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:
															snake_case_           =							model




	def 							_UpperCamelCase      (	self	,      *a	,      **a )	->   Union[str, Any]:
								snake_case_           =							self.model_type.from_pretrained(*__UpperCamelCase	,      **__UpperCamelCase )




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




	def 							_UpperCamelCase      (	self )	->   Union[str, Any]:
								snake_case_           =							self.model
								snake_case_           =							['bias', 'LayerNorm.weight']
								snake_case_           =							[
								    {
								        '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:
															snake_case_           =							Adafactor(
															    __UpperCamelCase	,      lr=self.hparams.learning_rate	,      scale_parameter=__UpperCamelCase	,      relative_step=__UpperCamelCase )

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

								snake_case_           =							self.get_lr_scheduler()

								return [optimizer], [scheduler]




	def 							_UpperCamelCase      (	self	,      a	,      a )	->   Dict:
								return self.validation_step(__UpperCamelCase	,      __UpperCamelCase )




	def 							_UpperCamelCase      (	self	,      a )	->   List[Any]:
								return self.validation_end(__UpperCamelCase )




	def 							_UpperCamelCase      (	self )	->   List[Any]:
								snake_case_           =							max(1	,      self.hparams.gpus )  # TODO: consider num_tpu_cores
								snake_case_           =							self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
								return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs




	def 							_UpperCamelCase      (	self	,      a )	->   int:
								if mode == "test":
															snake_case_           =							len(self.test_dataloader().dataset )
								else:
															snake_case_           =							self.get_dataloader('train'	,      self.hparams.train_batch_size	,      shuffle=__UpperCamelCase )
															snake_case_           =							len(self.train_dataloader().dataset )




	def 							_UpperCamelCase      (	self	,      a	,      a	,      a = False )	->   Optional[int]:
								raise NotImplementedError('You must implement this for your task' )




	def 							_UpperCamelCase      (	self )	->   List[Any]:
								return self.train_loader




	def 							_UpperCamelCase      (	self )	->   str:
								return self.get_dataloader('dev'	,      self.hparams.eval_batch_size	,      shuffle=__UpperCamelCase )




	def 							_UpperCamelCase      (	self )	->   Tuple:
								return self.get_dataloader('test'	,      self.hparams.eval_batch_size	,      shuffle=__UpperCamelCase )




	def 							_UpperCamelCase      (	self	,      a )	->   int:
								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 							_UpperCamelCase      (	self	,      a )	->   List[Any]:
								snake_case_           =							self.output_dir.joinpath('best_tfmr' )
								snake_case_           =							self.step_count
								self.model.save_pretrained(__UpperCamelCase )
								self.tokenizer.save_pretrained(__UpperCamelCase )






	@staticmethod
	def 							_UpperCamelCase      (	a	,      a )	->   List[Any]:
								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 			UpperCamelCase_				(      pl.Callback    ):



	'''simple docstring'''




	def 							_UpperCamelCase      (	self	,      a	,      a )	->   str:
								snake_case_           =							trainer.lr_schedulers[0]['scheduler']
								snake_case_           =							{F'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr() )}
								pl_module.logger.log_metrics(__UpperCamelCase )




	def 							_UpperCamelCase      (	self	,      a	,      a )	->   List[Any]:
								rank_zero_info('***** Validation results *****' )
								snake_case_           =							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 							_UpperCamelCase      (	self	,      a	,      a )	->   List[str]:
								rank_zero_info('***** Test results *****' )
								snake_case_           =							trainer.callback_metrics
								# Log and save results to file
								snake_case_           =							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      (					a_							,					a_):
							#  To allow all pl args uncomment the following line
							#  parser = pl.Trainer.add_argparse_args(parser)
							parser.add_argument(
							    '--output_dir'							,					default=a_							,					type=a_							,					required=a_							,					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=a_							,					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=a_)
							parser.add_argument('--max_grad_norm'							,					dest='gradient_clip_val'							,					default=1.0							,					type=a_							,					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=a_							,					default=1							,					help='Number of updates steps to accumulate before performing a backward/update pass.'							,					)
							parser.add_argument('--seed'							,					type=a_							,					default=42							,					help='random seed for initialization')
							parser.add_argument(
							    '--data_dir'							,					default=a_							,					type=a_							,					required=a_							,					help='The input data dir. Should contain the training files for the CoNLL-2003 NER task.'							,					)





def 						__UpperCAmelCase      (					a_							,					a_							,					a_=None							,					a_=True							,					a_=[]							,					a_=None							,					a_=None							,					**a_							,					):
							pl.seed_everything(args.seed)

							# init model
							snake_case_           =							Path(model.hparams.output_dir)
							odir.mkdir(exist_ok=a_)

							# add custom checkpoints
							if checkpoint_callback is None:
														snake_case_           =							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(a_)
							if logging_callback is None:
														snake_case_           =							LoggingCallback()

							snake_case_           =							{}

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

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

							snake_case_           =							args.accumulate_grad_batches
							snake_case_           =							extra_train_kwargs.get('accelerator'							,					a_)
							snake_case_           =							extra_train_kwargs.get('profiler'							,					a_)

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

							if args.do_train:
														trainer.fit(a_)

							return trainer













import os
from functools import partial
from glob import glob

import faiss
from datasets import Features, Sequence, Value, concatenate_datasets, load_dataset, load_from_disk

from transformers import DPRContextEncoder, DPRContextEncoderTokenizerFast





def 						__UpperCAmelCase      (					a_							,					a_=1_00							,					a_=" "):
							snake_case_           =							text.split(a_)
							return [character.join(text[i : i + n]).strip() for i in range(0							,					len(a_)							,					a_)]





def 						__UpperCAmelCase      (					a_):
							snake_case_      ,		snake_case_           =							[], []
							for title, text in zip(documents['title']							,					documents['text']):
														if text is not None:
																					for passage in split_text(a_):
																												titles.append(title if title is not None else '')
																												texts.append(a_)
							return {"title": titles, "text": texts}





def 						__UpperCAmelCase      (					a_							,					a_							,					a_							,					a_							,					a_							,					a_):
							snake_case_           =							load_dataset(
							    'csv'							,					data_files=[csv_path]							,					split='train'							,					delimiter='\t'							,					column_names=['title', 'text'])
							snake_case_           =							kb_dataset.map(
							    a_							,					batched=a_							,					num_proc=1)  # if you want you can load already splitted csv.
							snake_case_           =							[kb_dataset.shard(a_							,					a_							,					contiguous=a_) for i in range(a_)]
							snake_case_           =							kb_list[process_num]

							snake_case_           =							'data_' + str(a_)
							snake_case_           =							os.path.join(a_							,					a_)

							snake_case_           =							DPRContextEncoderTokenizerFast.from_pretrained('facebook/dpr-ctx_encoder-multiset-base')
							snake_case_           =							ctx_encoder.to(device=a_)

							def embed(
							    a_							,					a_							,					a_							,					a_) -> dict:
														snake_case_           =							ctx_tokenizer(
														    documents['title']							,					documents['text']							,					truncation=a_							,					padding='longest'							,					return_tensors='pt')['input_ids']
														snake_case_           =							ctx_encoder(input_ids.to(device=a_)							,					return_dict=a_).pooler_output
														return {"embeddings": embeddings.detach().cpu().numpy()}

							snake_case_           =							Features(
							    {'text': Value('string'), 'title': Value('string'), 'embeddings': Sequence(Value('float32'))})  # optional, save as float32 instead of float64 to save space

							snake_case_           =							data_shrad.map(
							    partial(a_							,					ctx_encoder=a_							,					ctx_tokenizer=a_							,					device=a_)							,					batched=a_							,					batch_size=16							,					features=a_							,					)
							dataset.save_to_disk(a_)





def 						__UpperCAmelCase      (					a_							,					a_):
							snake_case_           =							[]

							for shard_address in glob(str(a_) + '/*/'):
														data_shard_list.append(load_from_disk(a_))

							snake_case_           =							concatenate_datasets(a_)
							faiss.omp_set_num_threads(96)

							snake_case_           =							faiss.IndexHNSWFlat(7_68							,					1_28							,					faiss.METRIC_INNER_PRODUCT)
							concat.add_faiss_index('embeddings'							,					custom_index=a_)
							concat.get_index('embeddings').save(
							    a_)  # since we load the index in to memory,we can directly update the index in the disk





