




'''simple docstring'''






import os

import jsonlines
import numpy as np
from tqdm import tqdm


__UpperCAmelCase       =							2_048
__UpperCAmelCase       =							4_096
__UpperCAmelCase       =							42
__UpperCAmelCase       =							os.environ.pop('''PROCESS_TRAIN''', '''false''')
__UpperCAmelCase       =							{'''null''': 0, '''short''': 1, '''long''': 2, '''yes''': 3, '''no''': 4}


def  _snake_case       (					A	)						->   Optional[Any]:


							def choose_first(A			,   A=False	):
														assert isinstance(_lowerCamelCase			,   _lowerCamelCase	)
														if len(_lowerCamelCase	) == 1:
																					lowerCAmelCase__	       =     answer[0]
																					return {k: [answer[k]] for k in answer} if is_long_answer else answer
														for a in answer:
																					if is_long_answer:
																												lowerCAmelCase__	       =     {k: [a[k]] for k in a}
																					if len(a['''start_token''']	) > 0:
																												break
														return a

							lowerCAmelCase__	       =     {"""id""": example["""id"""]}
							lowerCAmelCase__	       =     example["""annotations"""]
							lowerCAmelCase__	       =     annotation["""yes_no_answer"""]
							if 0 in yes_no_answer or 1 in yes_no_answer:
														lowerCAmelCase__	       =     ["""yes"""] if 1 in yes_no_answer else ["""no"""]
														lowerCAmelCase__	       =     []
														lowerCAmelCase__	       =     []
														lowerCAmelCase__	       =     ["""<cls>"""]
							else:
														lowerCAmelCase__	       =     ["""short"""]
														lowerCAmelCase__	       =     choose_first(annotation['''short_answers''']	)
														if len(out['''start_token''']	) == 0:
																					# answer will be long if short is not available
																					lowerCAmelCase__	       =     ["""long"""]
																					lowerCAmelCase__	       =     choose_first(annotation['''long_answer''']			,   is_long_answer=_lowerCamelCase	)
																					lowerCAmelCase__	       =     []
														answer.update(_lowerCamelCase	)

							# disregard some samples
							if len(answer['''start_token''']	) > 1 or answer["start_token"] == answer["end_token"]:
														lowerCAmelCase__	       =     True
							else:
														lowerCAmelCase__	       =     False

							lowerCAmelCase__	       =     ["""start_token""", """end_token""", """start_byte""", """end_byte""", """text"""]
							if not all(isinstance(answer[k]			,   _lowerCamelCase	) for k in cols	):
														raise ValueError('''Issue in ID'''			,   example['''id''']	)

							return answer


def  _snake_case       (					A			,   A=False	)						->   List[str]:
							lowerCAmelCase__	       =     _get_single_answer(_lowerCamelCase	)
							# bytes are of no use
							del answer["start_byte"]
							del answer["end_byte"]

							# handle yes_no answers explicitly
							if answer["category"][0] in ["yes", "no"]:  # category is list with one element
														lowerCAmelCase__	       =     example["""document"""]["""tokens"""]
														lowerCAmelCase__	       =     []
														for i in range(len(doc['''token''']	)	):
																					if not doc["is_html"][i]:
																												context.append(doc['''token'''][i]	)
														return {
														    "context": " ".join(_lowerCamelCase	),
														    "answer": {
														        "start_token": -100,  # ignore index in cross-entropy
														        "end_token": -100,  # ignore index in cross-entropy
														        "category": answer["category"],
														        "span": answer["category"],  # extra
														    },
														}

							# later, help in removing all no answers
							if answer["start_token"] == [-1]:
														return {
														    "context": "None",
														    "answer": {
														        "start_token": -1,
														        "end_token": -1,
														        "category": "null",
														        "span": "None",  # extra
														    },
														}

							# handling normal samples

							lowerCAmelCase__	       =     ["""start_token""", """end_token"""]
							answer.update({k: answer[k][0] if len(answer[k]	) > 0 else answer[k] for k in cols}	)  # e.g. [10] == 10

							lowerCAmelCase__	       =     example["""document"""]["""tokens"""]
							lowerCAmelCase__	       =     answer["""start_token"""]
							lowerCAmelCase__	       =     answer["""end_token"""]

							lowerCAmelCase__	       =     []
							for i in range(len(doc['''token''']	)	):
														if not doc["is_html"][i]:
																					context.append(doc['''token'''][i]	)
														else:
																					if answer["start_token"] > i:
																												start_token -= 1
																					if answer["end_token"] > i:
																												end_token -= 1
							lowerCAmelCase__	       =     """ """.join(context[start_token:end_token]	)

							# checking above code
							if assertion:
														lowerCAmelCase__	       =     doc["""is_html"""][answer["""start_token"""] : answer["""end_token"""]]
														lowerCAmelCase__	       =     doc["""token"""][answer["""start_token"""] : answer["""end_token"""]]
														lowerCAmelCase__	       =     """ """.join([old[i] for i in range(len(_lowerCamelCase	)	) if not is_html[i]]	)
														if new != old:
																					print('''ID:'''			,   example['''id''']	)
																					print('''New:'''			,   _lowerCamelCase			,   end='''\n'''	)
																					print('''Old:'''			,   _lowerCamelCase			,   end='''\n\n'''	)

							return {
							    "context": " ".join(_lowerCamelCase	),
							    "answer": {
							        "start_token": start_token,
							        "end_token": end_token - 1,  # this makes it inclusive
							        "category": answer["category"],  # either long or short
							        "span": new,  # extra
							    },
							}


def  _snake_case       (					A			,   A			,   A=2048			,   A=4096			,   A=True	)						->   Dict:
							lowerCAmelCase__	       =     get_context_and_ans(_lowerCamelCase			,   assertion=_lowerCamelCase	)
							lowerCAmelCase__	       =     out["""answer"""]

							# later, removing these samples
							if answer["start_token"] == -1:
														return {
														    "example_id": example["id"],
														    "input_ids": [[-1]],
														    "labels": {
														        "start_token": [-1],
														        "end_token": [-1],
														        "category": ["null"],
														    },
														}

							lowerCAmelCase__	       =     tokenizer(example['''question''']['''text''']			,   out['''context''']	).input_ids
							lowerCAmelCase__	       =     input_ids.index(tokenizer.sep_token_id	) + 1

							# return yes/no
							if answer["category"][0] in ["yes", "no"]:  # category is list with one element
														lowerCAmelCase__	       =     []
														lowerCAmelCase__	       =     []
														lowerCAmelCase__	       =     input_ids[:q_len]
														lowerCAmelCase__	       =     range(_lowerCamelCase			,   len(_lowerCamelCase	)			,   max_length - doc_stride	)
														for i in doc_start_indices:
																					lowerCAmelCase__	       =     i + max_length - q_len
																					lowerCAmelCase__	       =     input_ids[i:end_index]
																					inputs.append(q_indices + slice	)
																					category.append(answer['''category'''][0]	)
																					if slice[-1] == tokenizer.sep_token_id:
																												break

														return {
														    "example_id": example["id"],
														    "input_ids": inputs,
														    "labels": {
														        "start_token": [-100] * len(_lowerCamelCase	),
														        "end_token": [-100] * len(_lowerCamelCase	),
														        "category": category,
														    },
														}

							lowerCAmelCase__	       =     out["""context"""].split()
							lowerCAmelCase__	       =     splitted_context[answer["""end_token"""]]
							lowerCAmelCase__	       =     len(
							    tokenizer(
							        ''' '''.join(splitted_context[: answer['''start_token''']]	)			,   add_special_tokens=_lowerCamelCase			,   ).input_ids	)
							lowerCAmelCase__	       =     len(
							    tokenizer(''' '''.join(splitted_context[: answer['''end_token''']]	)			,   add_special_tokens=_lowerCamelCase	).input_ids	)

							answer["start_token"] += q_len
							answer["end_token"] += q_len

							# fixing end token
							lowerCAmelCase__	       =     len(tokenizer(_lowerCamelCase			,   add_special_tokens=_lowerCamelCase	).input_ids	)
							if num_sub_tokens > 1:
														answer["end_token"] += num_sub_tokens - 1

							lowerCAmelCase__	       =     input_ids[answer["""start_token"""] : answer["""end_token"""] + 1]  # right & left are inclusive
							lowerCAmelCase__	       =     answer["""start_token"""]
							lowerCAmelCase__	       =     answer["""end_token"""]

							if assertion:
														lowerCAmelCase__	       =     tokenizer.decode(_lowerCamelCase	)
														if answer["span"] != new:
																					print('''ISSUE IN TOKENIZATION'''	)
																					print('''OLD:'''			,   answer['''span''']	)
																					print('''NEW:'''			,   _lowerCamelCase			,   end='''\n\n'''	)

							if len(_lowerCamelCase	) <= max_length:
														return {
														    "example_id": example["id"],
														    "input_ids": [input_ids],
														    "labels": {
														        "start_token": [answer["start_token"]],
														        "end_token": [answer["end_token"]],
														        "category": answer["category"],
														    },
														}

							lowerCAmelCase__	       =     input_ids[:q_len]
							lowerCAmelCase__	       =     range(_lowerCamelCase			,   len(_lowerCamelCase	)			,   max_length - doc_stride	)

							lowerCAmelCase__	       =     []
							lowerCAmelCase__	       =     []
							lowerCAmelCase__	       =     []
							lowerCAmelCase__	       =     []  # null, yes, no, long, short
							for i in doc_start_indices:
														lowerCAmelCase__	       =     i + max_length - q_len
														lowerCAmelCase__	       =     input_ids[i:end_index]
														inputs.append(q_indices + slice	)
														assert len(inputs[-1]	) <= max_length, "Issue in truncating length"

														if start_token >= i and end_token <= end_index - 1:
																					lowerCAmelCase__	       =     start_token - i + q_len
																					lowerCAmelCase__	       =     end_token - i + q_len
																					answers_category.append(answer['''category'''][0]	)  # ["short"] -> "short"
														else:
																					lowerCAmelCase__	       =     -100
																					lowerCAmelCase__	       =     -100
																					answers_category.append('''null'''	)
														lowerCAmelCase__	       =     inputs[-1][start_token : end_token + 1]

														answers_start_token.append(_lowerCamelCase	)
														answers_end_token.append(_lowerCamelCase	)
														if assertion:
																					if new != old and new != [tokenizer.cls_token_id]:
																												print('''ISSUE in strided for ID:'''			,   example['''id''']	)
																												print('''New:'''			,   tokenizer.decode(_lowerCamelCase	)	)
																												print('''Old:'''			,   tokenizer.decode(_lowerCamelCase	)			,   end='''\n\n'''	)
														if slice[-1] == tokenizer.sep_token_id:
																					break

							return {
							    "example_id": example["id"],
							    "input_ids": inputs,
							    "labels": {
							        "start_token": answers_start_token,
							        "end_token": answers_end_token,
							        "category": answers_category,
							    },
							}


def  _snake_case       (					A			,   A			,   A=2048			,   A=4096			,   A=False	)						->   int:
							lowerCAmelCase__	       =     get_strided_contexts_and_ans(
							    _lowerCamelCase			,   _lowerCamelCase			,   doc_stride=_lowerCamelCase			,   max_length=_lowerCamelCase			,   assertion=_lowerCamelCase			,   )

							return example


def  _snake_case       (					A			,   A	)						->   Optional[Any]:


							with jsonlines.open(_lowerCamelCase			,   '''a'''	) as writer:
														for example in tqdm(_lowerCamelCase			,   total=len(_lowerCamelCase	)			,   desc='''Saving samples ... '''	):
																					lowerCAmelCase__	       =     example["""labels"""]
																					for ids, start, end, cat in zip(
																					    example['''input_ids''']			,   labels['''start_token''']			,   labels['''end_token''']			,   labels['''category''']			,   ):
																												if start == -1 and end == -1:
																																			continue  # leave waste samples with no answer
																												if cat == "null" and np.random.rand() < 0.6:
																																			continue  # removing 50 % samples
																												writer.write(
																												    {
																												        '''input_ids''': ids,
																												        '''start_token''': start,
																												        '''end_token''': end,
																												        '''category''': CATEGORY_MAPPING[cat],
																												    }	)


if __name__ == "__main__":
		from datasets import load_dataset

		from transformers import BigBirdTokenizer

		__UpperCAmelCase       =							load_dataset('''natural_questions''')
		__UpperCAmelCase       =							BigBirdTokenizer.from_pretrained('''google/bigbird-roberta-base''')

		__UpperCAmelCase       =							data['''train''' if PROCESS_TRAIN == '''true''' else '''validation''']

		__UpperCAmelCase       =							{
		    '''tokenizer''': tokenizer,
		    '''doc_stride''': DOC_STRIDE,
		    '''max_length''': MAX_LENGTH,
		    '''assertion''': False,
		}
		__UpperCAmelCase       =							data.map(prepare_inputs, fn_kwargs=fn_kwargs)
		__UpperCAmelCase       =							data.remove_columns(['''annotations''', '''document''', '''id''', '''question'''])
		print(data)

		np.random.seed(SEED)
		__UpperCAmelCase       =							'''nq-training.jsonl''' if PROCESS_TRAIN == '''true''' else '''nq-validation.jsonl'''
		save_to_disk(data, file_name=cache_file_name)



'''simple docstring'''






def  _snake_case       (					A = 10	)						->   str:
							if not isinstance(A			,   A	) or n < 0:
														raise ValueError('''Invalid input'''	)
							lowerCAmelCase__	       =     10**n
							lowerCAmelCase__	       =     28433 * (pow(2			,   7830457			,   A	)) + 1
							return str(number % modulus	)


if __name__ == "__main__":
		from doctest import testmod

		testmod()
		print(f"""{solution(10) = }""")