



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 						_UpperCAmelCase							( snake_case     ):





					"""simple docstring"""

					_lowerCAmelCase         =       filter(lambda snake_case     : p.requires_grad      ,		model.parameters()     )
					_lowerCAmelCase         =       sum([np.prod(p.size()     ) for p in model_parameters]     )
					return params


A__       =			logging.getLogger(__name__)





class    __lowerCAmelCase					(	pl.Callback			):




				def 				snake_case  (				self		,							_snake_case		,							_snake_case						):



									"""simple docstring"""





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




				@rank_zero_only
				def 				snake_case  (				self		,							_snake_case		,							_snake_case		,							_snake_case		,							_snake_case=True						):



									"""simple docstring"""





									logger.info(F'***** {type_path} results at step {trainer.global_step:05d} *****'						)
									_lowerCAmelCase         =       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
									_lowerCAmelCase         =       Path(pl_module.hparams.output_dir						)
									if type_path == "test":
														_lowerCAmelCase         =       od / """test_results.txt"""
														_lowerCAmelCase         =       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.
														_lowerCAmelCase         =       od / F'{type_path}_results/{trainer.global_step:05d}.txt'
														_lowerCAmelCase         =       od / F'{type_path}_generations/{trainer.global_step:05d}.txt'
														results_file.parent.mkdir(exist_ok=_snake_case						)
														generations_file.parent.mkdir(exist_ok=_snake_case						)
									with open(_snake_case		,							"""a+"""						) as writer:
														for key in sorted(_snake_case						):
																			if key in ["log", "progress_bar", "preds"]:
																								continue
																			_lowerCAmelCase         =       metrics[key]
																			if isinstance(_snake_case		,							torch.Tensor						):
																								_lowerCAmelCase         =       val.item()
																			_lowerCAmelCase         =       F'{key}: {val:.6f}\n'
																			writer.write(_snake_case						)

									if not save_generations:
														return

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




				@rank_zero_only
				def 				snake_case  (				self		,							_snake_case		,							_snake_case						):



									"""simple docstring"""





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

									_lowerCAmelCase         =       count_trainable_parameters(_snake_case						)
									# 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 				snake_case  (				self		,							_snake_case		,							_snake_case						):



									"""simple docstring"""





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




				@rank_zero_only
				def 				snake_case  (				self		,							_snake_case		,							_snake_case						):



									"""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 						_UpperCAmelCase							( snake_case      ,		snake_case      ,		snake_case=1      ,		snake_case=False     ):





					"""simple docstring"""

					if metric == "rouge2":
										_lowerCAmelCase         =       """{val_avg_rouge2:.4f}-{step_count}"""
					elif metric == "bleu":
										_lowerCAmelCase         =       """{val_avg_bleu:.4f}-{step_count}"""
					elif metric == "loss":
										_lowerCAmelCase         =       """{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."""     )

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


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					return EarlyStopping(
					    monitor=F'val_{metric}'      ,		mode="""min""" if """loss""" in metric else """max"""      ,		patience=snake_case      ,		verbose=snake_case      ,		)








import os
from typing import Optional, Union

import numpy as np
from huggingface_hub import hf_hub_download

from ... import AutoTokenizer
from ...utils import logging


A__       =			"""block_records.npy"""


A__       =			logging.get_logger(__name__)


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					import tensorflow.compat.va as tf

					_lowerCAmelCase         =       tf.data.TFRecordDataset(snake_case      ,		buffer_size=5_12 * 10_24 * 10_24     )
					_lowerCAmelCase         =       blocks_dataset.batch(snake_case      ,		drop_remainder=snake_case     )
					_lowerCAmelCase         =       next(blocks_dataset.take(1     ).as_numpy_iterator()     )

					return np_record





class    __lowerCAmelCase					:




				def __init__(				self		,							_snake_case		,							_snake_case		,							_snake_case=2		,							_snake_case=1000		,							_snake_case=100		,							_snake_case=100000		,							):



									"""simple docstring"""





									from scann.scann_ops.py.scann_ops_pybind import builder as Builder

									_lowerCAmelCase         =       Builder(db=_snake_case		,							num_neighbors=_snake_case		,							distance_measure="""dot_product"""						)
									_lowerCAmelCase         =       builder.tree(
									    num_leaves=_snake_case		,							num_leaves_to_search=_snake_case		,							training_sample_size=_snake_case						)
									_lowerCAmelCase         =       builder.score_ah(dimensions_per_block=_snake_case						)

									_lowerCAmelCase         =       builder.build()




				def 				snake_case  (				self		,							_snake_case						):



									"""simple docstring"""





									_lowerCAmelCase      ,      _lowerCAmelCase         =       self.searcher.search_batched(question_projection.detach().cpu()						)
									return retrieved_block_ids.astype("""int64"""						)





class    __lowerCAmelCase					:




				def __init__(				self		,							_snake_case		,							_snake_case						):



									"""simple docstring"""





									super().__init__()
									_lowerCAmelCase         =       block_records
									_lowerCAmelCase         =       tokenizer




				def __call__(				self		,							_snake_case		,							_snake_case		,							_snake_case		,							_snake_case=None		,							_snake_case="pt"						):



									"""simple docstring"""





									_lowerCAmelCase         =       np.take(self.block_records		,							indices=_snake_case		,							axis=0						)

									_lowerCAmelCase         =       self.tokenizer.decode(question_input_ids[0]		,							skip_special_tokens=_snake_case						)

									_lowerCAmelCase         =       []
									_lowerCAmelCase         =       []
									for retrieved_block in retrieved_blocks:
														text.append(_snake_case						)
														text_pair.append(retrieved_block.decode()						)

									_lowerCAmelCase         =       self.tokenizer(
									    _snake_case		,							_snake_case		,							padding=_snake_case		,							truncation=_snake_case		,							return_special_tokens_mask=_snake_case		,							max_length=_snake_case						)
									_lowerCAmelCase         =       concat_inputs.convert_to_tensors(_snake_case						)

									if answer_ids is not None:
														return self.block_has_answer(_snake_case		,							_snake_case						) + (concat_inputs_tensors,)
									else:
														return (None, None, None, concat_inputs_tensors)




				@classmethod
				def 				snake_case  (				cls		,							_snake_case		,							*_snake_case		,							**_snake_case						):



									"""simple docstring"""





									if os.path.isdir(_snake_case						):
														_lowerCAmelCase         =       os.path.join(_snake_case		,							_REALM_BLOCK_RECORDS_FILENAME						)
									else:
														_lowerCAmelCase         =       hf_hub_download(
														    repo_id=_snake_case		,							filename=_REALM_BLOCK_RECORDS_FILENAME		,							**_snake_case						)
									_lowerCAmelCase         =       np.load(_snake_case		,							allow_pickle=_snake_case						)

									_lowerCAmelCase         =       AutoTokenizer.from_pretrained(_snake_case		,							*_snake_case		,							**_snake_case						)

									return cls(_snake_case		,							_snake_case						)




				def 				snake_case  (				self		,							_snake_case						):



									"""simple docstring"""





									np.save(os.path.join(_snake_case		,							_REALM_BLOCK_RECORDS_FILENAME						)		,							self.block_records						)
									# save tokenizer
									self.tokenizer.save_pretrained(_snake_case						)




				def 				snake_case  (				self		,							_snake_case		,							_snake_case						):



									"""simple docstring"""





									_lowerCAmelCase         =       []
									_lowerCAmelCase         =       []
									_lowerCAmelCase         =       []
									_lowerCAmelCase         =       0

									for input_id in concat_inputs.input_ids:
														_lowerCAmelCase         =       input_id.tolist()
														# Check answers between two [SEP] tokens
														_lowerCAmelCase         =       input_id_list.index(self.tokenizer.sep_token_id						)
														_lowerCAmelCase         =       first_sep_idx + 1 + input_id_list[first_sep_idx + 1 :].index(self.tokenizer.sep_token_id						)

														start_pos.append([]						)
														end_pos.append([]						)
														for answer in answer_ids:
																			for idx in range(first_sep_idx + 1		,							_snake_case						):
																								if answer[0] == input_id_list[idx]:
																													if input_id_list[idx : idx + len(_snake_case						)] == answer:
																																		start_pos[-1].append(_snake_case						)
																																		end_pos[-1].append(idx + len(_snake_case						) - 1						)

														if len(start_pos[-1]						) == 0:
																			has_answers.append(_snake_case						)
														else:
																			has_answers.append(_snake_case						)
																			if len(start_pos[-1]						) > max_answers:
																								_lowerCAmelCase         =       len(start_pos[-1]						)

        # Pad -1 to max_answers
									for start_pos_, end_pos_ in zip(_snake_case		,							_snake_case						):
														if len(start_pos_						) < max_answers:
																			_lowerCAmelCase         =       [-1] * (max_answers - len(start_pos_						))
																			start_pos_ += padded
																			end_pos_ += padded
									return has_answers, start_pos, end_pos



