




import os

import jsonlines
import numpy as np
from tqdm import tqdm


UpperCAmelCase_						=2048
UpperCAmelCase_						=4096
UpperCAmelCase_						=42
UpperCAmelCase_						=os.environ.pop("""PROCESS_TRAIN""", """false""")
UpperCAmelCase_						={"""null""": 0, """short""": 1, """long""": 2, """yes""": 3, """no""": 4}



def       UpperCAmelCase       (     _snake_case							):
							def choose_first(_snake_case						, _snake_case=False							):
														assert isinstance(_snake_case						, _snake_case							)
														if len(_snake_case							) == 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   =		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=_snake_case							)
																					lowerCAmelCase   =		[]
														answer.update(_snake_case							)

							# 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]						, _snake_case							) for k in cols							):
														raise ValueError('''Issue in ID'''						, example['''id''']							)

							return answer



def       UpperCAmelCase       (     _snake_case						, _snake_case=False							):
							lowerCAmelCase   =		_get_single_answer(_snake_case							)
							# 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(_snake_case							),
														    "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(_snake_case							)							) if not is_html[i]]							)
														if new != old:
																					print('''ID:'''						, example['''id''']							)
																					print('''New:'''						, _snake_case						, end='''\n'''							)
																					print('''Old:'''						, _snake_case						, end='''\n\n'''							)

							return {
							    "context": " ".join(_snake_case							),
							    "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       UpperCAmelCase       (     _snake_case						, _snake_case						, _snake_case=2048						, _snake_case=4096						, _snake_case=True							):
							# overlap will be of doc_stride - q_len

							lowerCAmelCase   =		get_context_and_ans(_snake_case						, assertion=_snake_case							)
							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(_snake_case						, len(_snake_case							)						, 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(_snake_case							),
														        "end_token": [-100] * len(_snake_case							),
														        "category": category,
														    },
														}

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

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

							# fixing end token
							lowerCAmelCase   =		len(tokenizer(_snake_case						, add_special_tokens=_snake_case							).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(_snake_case							)
														if answer["span"] != new:
																					print('''ISSUE IN TOKENIZATION'''							)
																					print('''OLD:'''						, answer['''span''']							)
																					print('''NEW:'''						, _snake_case						, end='''\n\n'''							)

							if len(_snake_case							) <= 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(_snake_case						, len(_snake_case							)						, 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(_snake_case							)
														answers_end_token.append(_snake_case							)
														if assertion:
																					if new != old and new != [tokenizer.cls_token_id]:
																												print('''ISSUE in strided for ID:'''						, example['''id''']							)
																												print('''New:'''						, tokenizer.decode(_snake_case							)							)
																												print('''Old:'''						, tokenizer.decode(_snake_case							)						, 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       UpperCAmelCase       (     _snake_case						, _snake_case						, _snake_case=2048						, _snake_case=4096						, _snake_case=False							):
							lowerCAmelCase   =		get_strided_contexts_and_ans(
							    _snake_case						, _snake_case						, doc_stride=_snake_case						, max_length=_snake_case						, assertion=_snake_case						, )

							return example



def       UpperCAmelCase       (     _snake_case						, _snake_case							):
							with jsonlines.open(_snake_case						, '''a'''							) as writer:
														for example in tqdm(_snake_case						, total=len(_snake_case							)						, 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)









from typing import Optional

import pyspark

from .. import Features, NamedSplit
from ..download import DownloadMode
from ..packaged_modules.spark.spark import Spark
from .abc import AbstractDatasetReader


class 							__UpperCamelCase			( __UpperCAmelCase  ):





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_       , UpperCAmelCase_ = None       , UpperCAmelCase_ = None       , UpperCAmelCase_ = True       , UpperCAmelCase_ = None       , UpperCAmelCase_ = False       , UpperCAmelCase_ = None       , UpperCAmelCase_ = True       , UpperCAmelCase_ = "arrow"       , **UpperCAmelCase_       , ):
											super().__init__(
											    split=UpperCAmelCase_       , features=UpperCAmelCase_       , cache_dir=UpperCAmelCase_       , keep_in_memory=UpperCAmelCase_       , streaming=UpperCAmelCase_       , **UpperCAmelCase_       , )
											lowerCAmelCase   =		load_from_cache_file
											lowerCAmelCase   =		file_format
											lowerCAmelCase   =		Spark(
											    df=UpperCAmelCase_       , features=UpperCAmelCase_       , cache_dir=UpperCAmelCase_       , working_dir=UpperCAmelCase_       , **UpperCAmelCase_       , )






				def       __snake_case		( self						):
											if self.streaming:
																		return self.builder.as_streaming_dataset(split=self.split						)
											lowerCAmelCase   =		None if self._load_from_cache_file else DownloadMode.FORCE_REDOWNLOAD
											self.builder.download_and_prepare(
											    download_mode=UpperCAmelCase_       , file_format=self._file_format       , )
											return self.builder.as_dataset(split=self.split						)



