



import argparse
import json
import re
import string
import sys
from collections import Counter


def 						_UpperCAmelCase							( snake_case     ):





					"""simple docstring"""

					def remove_articles(snake_case     ):
										return re.sub(R"""\b(a|an|the)\b"""      ,		""" """      ,		snake_case     )

					def white_space_fix(snake_case     ):
										return " ".join(text.split()     )

					def remove_punc(snake_case     ):
										_lowerCAmelCase         =       set(string.punctuation     )
										return "".join(ch for ch in text if ch not in exclude     )

					def lower(snake_case     ):
										return text.lower()

					return white_space_fix(remove_articles(remove_punc(lower(snake_case     )     )     )     )


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					_lowerCAmelCase         =       normalize_answer(snake_case     ).split()
					_lowerCAmelCase         =       normalize_answer(snake_case     ).split()
					_lowerCAmelCase         =       Counter(snake_case     ) & Counter(snake_case     )
					_lowerCAmelCase         =       sum(common.values()     )
					if num_same == 0:
										return 0
					_lowerCAmelCase         =       1.0 * num_same / len(snake_case     )
					_lowerCAmelCase         =       1.0 * num_same / len(snake_case     )
					_lowerCAmelCase         =       (2 * precision * recall) / (precision + recall)
					return fa


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					return normalize_answer(snake_case     ) == normalize_answer(snake_case     )


def 						_UpperCAmelCase							( snake_case      ,		snake_case      ,		snake_case     ):





					"""simple docstring"""

					_lowerCAmelCase         =       []
					for ground_truth in ground_truths:
										_lowerCAmelCase         =       metric_fn(snake_case      ,		snake_case     )
										scores_for_ground_truths.append(snake_case     )
					return max(snake_case     )


def 						_UpperCAmelCase							( snake_case      ,		snake_case     ):





					"""simple docstring"""

					_lowerCAmelCase         =       _lowerCAmelCase         =       _lowerCAmelCase         =       0
					for article in dataset:
										for paragraph in article["paragraphs"]:
															for qa in paragraph["qas"]:
																				total += 1
																				if qa["id"] not in predictions:
																									_lowerCAmelCase         =       """Unanswered question """ + qa["""id"""] + """ will receive score 0."""
																									print(snake_case      ,		file=sys.stderr     )
																									continue
																				_lowerCAmelCase         =       [x["""text"""] for x in qa["""answers"""]]
																				_lowerCAmelCase         =       predictions[qa["""id"""]]
																				exact_match += metric_max_over_ground_truths(snake_case      ,		snake_case      ,		snake_case     )
																				fa += metric_max_over_ground_truths(snake_case      ,		snake_case      ,		snake_case     )

					_lowerCAmelCase         =       100.0 * exact_match / total
					_lowerCAmelCase         =       100.0 * fa / total

					return {"exact_match": exact_match, "f1": fa}


if __name__ == "__main__":
	A__       =			"""1.1"""
	A__       =			argparse.ArgumentParser(description="""Evaluation for SQuAD """ + expected_version)
	parser.add_argument("""dataset_file""", help="""Dataset file""")
	parser.add_argument("""prediction_file""", help="""Prediction File""")
	A__       =			parser.parse_args()
	with open(args.dataset_file) as dataset_file:
		A__       =			json.load(dataset_file)
		if dataset_json["version"] != expected_version:
			print(
			    """Evaluation expects v-""" + expected_version + """, but got dataset with v-""" + dataset_json["""version"""],
			    file=sys.stderr,
			)
		A__       =			dataset_json["""data"""]
	with open(args.prediction_file) as prediction_file:
		A__       =			json.load(prediction_file)
	print(json.dumps(evaluate(dataset, predictions)))








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



