




import jax
import jax.numpy as jnp
from bigbird_flax import FlaxBigBirdForNaturalQuestions
from datasets import load_from_disk

from transformers import BigBirdTokenizerFast


lowercase_		=							{0: '''null''', 1: '''short''', 2: '''long''', 3: '''yes''', 4: '''no'''}
lowercase_		=							set(''''''.join(['''‘''', '''’''', '''´''', '''`''', '''.''', ''',''', '''-''', '''"''']))







def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase=0,	UpperCAmelCase=None       )  ->Union[str, Any]:






					"""simple docstring"""






					return [" ".join(x.split(''' '''       )[begin:end]       ) for x in answers if len(x.split(''' '''       )       ) > 1]







def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase=False       )  ->List[str]:






					"""simple docstring"""






					if make_sub_answers:
										# if answers are longer than one word, make sure a predictions is correct if it coresponds to the complete 1: or :-1 sub word
										# *e.g.* if the correct answer contains a prefix such as "the", or "a"
										__magic_name__					:							str		   =							(
										    given_answers + get_sub_answers(UpperCAmelCase,	begin=1       ) + get_sub_answers(UpperCAmelCase,	end=-1       )
										)
					__magic_name__					:							str		   =							[]
					for answer in given_answers:
										__magic_name__					:							Optional[Any]		   =							answer.replace('''_''',	''' '''       ).lower()
										__magic_name__					:							int		   =							''''''.join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else ''' ''' for c in alias       )
										answers.append(''' '''.join(alias.split()       ).strip()       )
					return set(UpperCAmelCase       )







def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase,	UpperCAmelCase=1,	UpperCAmelCase=100       )  ->int:






					"""simple docstring"""






					__magic_name__   , __magic_name__					:							Tuple		   =							jax.lax.top_k(UpperCAmelCase,	UpperCAmelCase       )
					__magic_name__   , __magic_name__					:							Dict		   =							jax.lax.top_k(UpperCAmelCase,	UpperCAmelCase       )

					__magic_name__					:							Union[str, Any]		   =							best_end_idx[:, None] - best_start_idx[None, :]
					__magic_name__					:							Union[str, Any]		   =							jnp.logical_or(widths < 0,	widths > max_size       )
					__magic_name__					:							List[str]		   =							(best_end_scores[:, None] + best_start_scores[None, :]) - (1E8 * mask)
					__magic_name__					:							List[Any]		   =							jnp.argmax(UpperCAmelCase       ).item()

					return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]







def  lowerCAmelCase  (	UpperCAmelCase       )  ->int:






					"""simple docstring"""






					__magic_name__					:							Dict		   =							sample['''question''']['''text''']
					__magic_name__					:							int		   =							sample['''document''']['''tokens''']['''token''']
					__magic_name__					:							Union[str, Any]		   =							sample['''document''']['''tokens''']['''is_html''']
					__magic_name__					:							int		   =							sample['''annotations''']['''long_answer''']
					__magic_name__					:							Optional[Any]		   =							sample['''annotations''']['''short_answers''']

					__magic_name__					:							Optional[int]		   =							''' '''.join([context[i] for i in range(len(UpperCAmelCase       )       ) if not is_html[i]]       )

					# 0 - No ; 1 - Yes
					for answer in sample["annotations"]["yes_no_answer"]:
										if answer == 0 or answer == 1:
															return {
															    "question": question,
															    "context": context_string,
															    "short": [],
															    "long": [],
															    "category": "no" if answer == 0 else "yes",
															}

					__magic_name__					:							List[Any]		   =							[]
					for s in short_answers:
										short_targets.extend(s['''text''']       )
					__magic_name__					:							Dict		   =							list(set(UpperCAmelCase       )       )

					__magic_name__					:							Tuple		   =							[]
					for s in long_answers:
										if s["start_token"] == -1:
															continue
										__magic_name__					:							Union[str, Any]		   =							context[s['''start_token'''] : s['''end_token''']]
										__magic_name__					:							int		   =							is_html[s['''start_token'''] : s['''end_token''']]
										__magic_name__					:							Any		   =							''' '''.join([answer[i] for i in range(len(UpperCAmelCase       )       ) if not html[i]]       )
										if new_answer not in long_targets:
															long_targets.append(UpperCAmelCase       )

					__magic_name__					:							Union[str, Any]		   =							'''long_short''' if len(short_targets + long_targets       ) > 0 else '''null'''

					return {
					    "question": question,
					    "context": context_string,
					    "short": short_targets,
					    "long": long_targets,
					    "category": category,
					}







def  lowerCAmelCase  (	)  ->Any:






					"""simple docstring"""






					__magic_name__					:							List[str]		   =							load_from_disk('''natural-questions-validation'''       )
					__magic_name__					:							Tuple		   =							dataset.map(UpperCAmelCase       ).remove_columns(['''annotations''', '''document''', '''id''']       )
					print(UpperCAmelCase       )

					__magic_name__					:							Tuple		   =							dataset.filter(lambda UpperCAmelCase       : (len(x['''question''']       ) + len(x['''context''']       )) < 4 * 4096       )
					__magic_name__					:							Dict		   =							short_validation_dataset.filter(lambda UpperCAmelCase       : x["category"] != "null"       )
					short_validation_dataset

					__magic_name__					:							Any		   =							'''vasudevgupta/flax-bigbird-natural-questions'''
					__magic_name__					:							str		   =							FlaxBigBirdForNaturalQuestions.from_pretrained(UpperCAmelCase       )
					__magic_name__					:							Optional[Any]		   =							BigBirdTokenizerFast.from_pretrained(UpperCAmelCase       )

					@jax.jit
					def forward(*UpperCAmelCase,	**UpperCAmelCase       ):
										__magic_name__   , __magic_name__   , __magic_name__					:							Dict		   =							model(*UpperCAmelCase,	**UpperCAmelCase       )
										return start_logits, end_logits, jnp.argmax(UpperCAmelCase,	axis=-1       )

					def evaluate(UpperCAmelCase       ):
										# encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
										__magic_name__					:							Optional[int]		   =							tokenizer(
										    example['''question'''],	example['''context'''],	return_tensors='''np''',	max_length=4096,	padding='''max_length''',	truncation=UpperCAmelCase,	)

										__magic_name__   , __magic_name__   , __magic_name__					:							Tuple		   =							forward(**UpperCAmelCase       )

										__magic_name__					:							Dict		   =							CATEGORY_MAPPING[category.item()]

										__magic_name__					:							Optional[Any]		   =							example['''long'''] + example['''short''']
										if example["category"] in ["yes", "no", "null"]:
															__magic_name__					:							Dict		   =							[example['''category''']]
										__magic_name__					:							List[str]		   =							example['''category'''] != '''null'''
										# Now target can be: "yes", "no", "null", "list of long & short answers"

										if predicted_category in ["yes", "no", "null"]:
															__magic_name__					:							Optional[Any]		   =							[predicted_category]
															__magic_name__					:							Dict		   =							example['''output'''] == example['''targets''']
															__magic_name__					:							Union[str, Any]		   =							predicted_category != '''null'''
															return example

										__magic_name__					:							List[str]		   =							38 if predicted_category == '''short''' else 1024
										__magic_name__   , __magic_name__					:							Tuple		   =							get_best_valid_start_end_idx(
										    start_scores[0],	end_scores[0],	top_k=8,	max_size=UpperCAmelCase       )

										__magic_name__					:							str		   =							inputs['''input_ids'''][0].tolist()
										__magic_name__					:							Dict		   =							[tokenizer.decode(input_ids[start_score : end_score + 1]       )]

										__magic_name__					:							int		   =							expand_to_aliases(example['''targets'''],	make_sub_answers=UpperCAmelCase       )
										__magic_name__					:							List[Any]		   =							expand_to_aliases(example['''output''']       )

										# some preprocessing to both prediction and answer
										__magic_name__					:							int		   =							{''''''.join(a.split()       ) for a in answers}
										__magic_name__					:							Optional[int]		   =							{''''''.join(p.split()       ) for p in predictions}
										__magic_name__					:							int		   =							{s for s in predictions if s not in ['''``''', '''\'\'''', '''`''', '''\'''']}

										# if there is a common element, it's a exact match
										__magic_name__					:							str		   =							len(list(answers & predictions       )       ) > 0
										__magic_name__					:							Dict		   =							predicted_category != '''null''' and len(UpperCAmelCase       ) > 0

										return example

					__magic_name__					:							int		   =							short_validation_dataset.map(UpperCAmelCase       )

					__magic_name__					:							Dict		   =							len(UpperCAmelCase       )
					__magic_name__					:							Dict		   =							len(short_validation_dataset.filter(lambda UpperCAmelCase       : x["match"] == 1       )       )
					print('''EM score:''',	(matched / total) * 100,	'''%'''       )


if __name__ == "__main__":
				main()









import csv
import os

import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.csv import CsvDatasetReader, CsvDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases







def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase       )  ->str:






					"""simple docstring"""






					assert isinstance(UpperCAmelCase,	UpperCAmelCase       )
					assert dataset.num_rows == 4
					assert dataset.num_columns == 3
					assert dataset.column_names == ["col_1", "col_2", "col_3"]
					for feature, expected_dtype in expected_features.items():
										assert dataset.features[feature].dtype == expected_dtype







@pytest.mark.parametrize('''keep_in_memory''',	[False, True]       )
def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase,	UpperCAmelCase       )  ->Any:






					"""simple docstring"""






					__magic_name__					:							Dict		   =							tmp_path / '''cache'''
					__magic_name__					:							Tuple		   =							{'''col_1''': '''int64''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
					with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
										__magic_name__					:							str		   =							CsvDatasetReader(UpperCAmelCase,	cache_dir=UpperCAmelCase,	keep_in_memory=UpperCAmelCase       ).read()
					_check_csv_dataset(UpperCAmelCase,	UpperCAmelCase       )







@pytest.mark.parametrize(
    '''features''',	[
        None,
        {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''},
        {'''col_1''': '''string''', '''col_2''': '''string''', '''col_3''': '''string'''},
        {'''col_1''': '''int32''', '''col_2''': '''int32''', '''col_3''': '''int32'''},
        {'''col_1''': '''float32''', '''col_2''': '''float32''', '''col_3''': '''float32'''},
    ],	)
def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase,	UpperCAmelCase       )  ->Tuple:






					"""simple docstring"""






					__magic_name__					:							Tuple		   =							tmp_path / '''cache'''
					# CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
					__magic_name__					:							List[Any]		   =							{'''col_1''': '''int64''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
					__magic_name__					:							Tuple		   =							features.copy() if features else default_expected_features
					__magic_name__					:							str		   =							(
					    Features({feature: Value(UpperCAmelCase       ) for feature, dtype in features.items()}       ) if features is not None else None
					)
					__magic_name__					:							Dict		   =							CsvDatasetReader(UpperCAmelCase,	features=UpperCAmelCase,	cache_dir=UpperCAmelCase       ).read()
					_check_csv_dataset(UpperCAmelCase,	UpperCAmelCase       )







@pytest.mark.parametrize('''split''',	[None, NamedSplit('''train'''       ), '''train''', '''test''']       )
def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase,	UpperCAmelCase       )  ->List[Any]:






					"""simple docstring"""






					__magic_name__					:							str		   =							tmp_path / '''cache'''
					__magic_name__					:							List[str]		   =							{'''col_1''': '''int64''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
					__magic_name__					:							int		   =							CsvDatasetReader(UpperCAmelCase,	cache_dir=UpperCAmelCase,	split=UpperCAmelCase       ).read()
					_check_csv_dataset(UpperCAmelCase,	UpperCAmelCase       )
					assert dataset.split == split if split else "train"







@pytest.mark.parametrize('''path_type''',	[str, list]       )
def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase,	UpperCAmelCase       )  ->str:






					"""simple docstring"""






					if issubclass(UpperCAmelCase,	UpperCAmelCase       ):
										__magic_name__					:							str		   =							csv_path
					elif issubclass(UpperCAmelCase,	UpperCAmelCase       ):
										__magic_name__					:							str		   =							[csv_path]
					__magic_name__					:							Optional[int]		   =							tmp_path / '''cache'''
					__magic_name__					:							Any		   =							{'''col_1''': '''int64''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
					__magic_name__					:							Tuple		   =							CsvDatasetReader(UpperCAmelCase,	cache_dir=UpperCAmelCase       ).read()
					_check_csv_dataset(UpperCAmelCase,	UpperCAmelCase       )







def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase,	UpperCAmelCase=("train",)       )  ->Tuple:






					"""simple docstring"""






					assert isinstance(UpperCAmelCase,	UpperCAmelCase       )
					for split in splits:
										__magic_name__					:							int		   =							dataset_dict[split]
										assert dataset.num_rows == 4
										assert dataset.num_columns == 3
										assert dataset.column_names == ["col_1", "col_2", "col_3"]
										for feature, expected_dtype in expected_features.items():
															assert dataset.features[feature].dtype == expected_dtype







@pytest.mark.parametrize('''keep_in_memory''',	[False, True]       )
def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase,	UpperCAmelCase       )  ->int:






					"""simple docstring"""






					__magic_name__					:							Tuple		   =							tmp_path / '''cache'''
					__magic_name__					:							int		   =							{'''col_1''': '''int64''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
					with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
										__magic_name__					:							Optional[Any]		   =							CsvDatasetReader({'''train''': csv_path},	cache_dir=UpperCAmelCase,	keep_in_memory=UpperCAmelCase       ).read()
					_check_csv_datasetdict(UpperCAmelCase,	UpperCAmelCase       )







@pytest.mark.parametrize(
    '''features''',	[
        None,
        {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''},
        {'''col_1''': '''string''', '''col_2''': '''string''', '''col_3''': '''string'''},
        {'''col_1''': '''int32''', '''col_2''': '''int32''', '''col_3''': '''int32'''},
        {'''col_1''': '''float32''', '''col_2''': '''float32''', '''col_3''': '''float32'''},
    ],	)
def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase,	UpperCAmelCase       )  ->List[str]:






					"""simple docstring"""






					__magic_name__					:							Any		   =							tmp_path / '''cache'''
					# CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
					__magic_name__					:							Dict		   =							{'''col_1''': '''int64''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
					__magic_name__					:							Dict		   =							features.copy() if features else default_expected_features
					__magic_name__					:							Any		   =							(
					    Features({feature: Value(UpperCAmelCase       ) for feature, dtype in features.items()}       ) if features is not None else None
					)
					__magic_name__					:							List[str]		   =							CsvDatasetReader({'''train''': csv_path},	features=UpperCAmelCase,	cache_dir=UpperCAmelCase       ).read()
					_check_csv_datasetdict(UpperCAmelCase,	UpperCAmelCase       )







@pytest.mark.parametrize('''split''',	[None, NamedSplit('''train'''       ), '''train''', '''test''']       )
def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase,	UpperCAmelCase       )  ->List[str]:






					"""simple docstring"""






					if split:
										__magic_name__					:							Dict		   =							{split: csv_path}
					else:
										__magic_name__					:							Any		   =							{'''train''': csv_path, '''test''': csv_path}
					__magic_name__					:							str		   =							tmp_path / '''cache'''
					__magic_name__					:							Union[str, Any]		   =							{'''col_1''': '''int64''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
					__magic_name__					:							int		   =							CsvDatasetReader(UpperCAmelCase,	cache_dir=UpperCAmelCase       ).read()
					_check_csv_datasetdict(UpperCAmelCase,	UpperCAmelCase,	splits=list(path.keys()       )       )
					assert all(dataset[split].split == split for split in path.keys()       )







def  lowerCAmelCase  (	UpperCAmelCase       )  ->int:






					"""simple docstring"""






					with open(UpperCAmelCase,	encoding='''utf-8'''       ) as csvfile:
										yield from csv.reader(UpperCAmelCase       )







def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase       )  ->int:






					"""simple docstring"""






					__magic_name__					:							Dict		   =							tmp_path / '''cache'''
					__magic_name__					:							List[Any]		   =							os.path.join(UpperCAmelCase,	'''tmp.csv'''       )
					__magic_name__					:							int		   =							CsvDatasetReader({'''train''': csv_path},	cache_dir=UpperCAmelCase       ).read()
					CsvDatasetWriter(dataset['''train'''],	UpperCAmelCase,	num_proc=1       ).write()

					__magic_name__					:							Optional[int]		   =							iter_csv_file(UpperCAmelCase       )
					__magic_name__					:							Any		   =							iter_csv_file(UpperCAmelCase       )

					for rowa, rowa in zip(UpperCAmelCase,	UpperCAmelCase       ):
										assert rowa == rowa







def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase       )  ->Union[str, Any]:






					"""simple docstring"""






					__magic_name__					:							Optional[int]		   =							tmp_path / '''cache'''
					__magic_name__					:							Optional[Any]		   =							os.path.join(UpperCAmelCase,	'''tmp.csv'''       )
					__magic_name__					:							List[Any]		   =							CsvDatasetReader({'''train''': csv_path},	cache_dir=UpperCAmelCase       ).read()
					CsvDatasetWriter(dataset['''train'''],	UpperCAmelCase,	num_proc=2       ).write()

					__magic_name__					:							Tuple		   =							iter_csv_file(UpperCAmelCase       )
					__magic_name__					:							List[str]		   =							iter_csv_file(UpperCAmelCase       )

					for rowa, rowa in zip(UpperCAmelCase,	UpperCAmelCase       ):
										assert rowa == rowa







def  lowerCAmelCase  (	UpperCAmelCase,	UpperCAmelCase       )  ->int:






					"""simple docstring"""






					__magic_name__					:							Optional[int]		   =							tmp_path / '''cache'''
					__magic_name__					:							Any		   =							os.path.join(UpperCAmelCase,	'''tmp.csv'''       )
					__magic_name__					:							Any		   =							CsvDatasetReader({'''train''': csv_path},	cache_dir=UpperCAmelCase       ).read()
					with pytest.raises(UpperCAmelCase       ):
										CsvDatasetWriter(dataset['''train'''],	UpperCAmelCase,	num_proc=0       )



