

import os

import jsonlines
import numpy as np
from tqdm import tqdm


lowerCAmelCase__		      =       2_0_4_8
lowerCAmelCase__		      =       4_0_9_6
lowerCAmelCase__		      =       4_2
lowerCAmelCase__		      =       os.environ.pop("""PROCESS_TRAIN""", """false""")
lowerCAmelCase__		      =       {"""null""": 0, """short""": 1, """long""": 2, """yes""": 3, """no""": 4}



def     lowerCAmelCase__							(		SCREAMING_SNAKE_CASE_:							Union[str, Any]					)	->					Dict:




				'''simple docstring'''





				def choose_first(SCREAMING_SNAKE_CASE_:							Dict		,				SCREAMING_SNAKE_CASE_:							Tuple=False					):
								assert isinstance(UpperCamelCase__		,				UpperCamelCase__					)
								if len(UpperCamelCase__					) == 1:
												A__						=	answer[0]
												return {k: [answer[k]] for k in answer} if is_long_answer else answer
								for a in answer:
												if is_long_answer:
																A__						=	{k: [a[k]] for k in a}
												if len(a["start_token"]					) > 0:
																break
								return a

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

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

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

				return answer



def     lowerCAmelCase__							(		SCREAMING_SNAKE_CASE_:							Optional[Any]		,				SCREAMING_SNAKE_CASE_:							Union[str, Any]=False					)	->					Tuple:




				'''simple docstring'''





				A__						=	_get_single_answer(UpperCamelCase__					)
				# 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
								A__						=	example["document"]["tokens"]
								A__						=	[]
								for i in range(len(doc["token"]					)					):
												if not doc["is_html"][i]:
																context.append(doc["token"][i]					)
								return {
								    "context": " ".join(UpperCamelCase__					),
								    "answer": {
								        "start_token": -1_0_0,  # ignore index in cross-entropy
								        "end_token": -1_0_0,  # 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

				A__						=	["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

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

				A__						=	[]
				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
				A__						=	" ".join(context[start_token:end_token]					)

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

				return {
				    "context": " ".join(UpperCamelCase__					),
				    "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     lowerCAmelCase__							(		SCREAMING_SNAKE_CASE_:							List[str]		,				SCREAMING_SNAKE_CASE_:							Any		,				SCREAMING_SNAKE_CASE_:							Optional[int]=2_0_4_8		,				SCREAMING_SNAKE_CASE_:							Optional[Any]=4_0_9_6		,				SCREAMING_SNAKE_CASE_:							Union[str, Any]=True					)	->					Dict:




				'''simple docstring'''





				A__						=	get_context_and_ans(UpperCamelCase__		,				assertion=UpperCamelCase__					)
				A__						=	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"],
								    },
								}

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

				# return yes/no
				if answer["category"][0] in ["yes", "no"]:  # category is list with one element
								A__						=	[]
								A__						=	[]
								A__						=	input_ids[:q_len]
								A__						=	range(UpperCamelCase__		,				len(UpperCamelCase__					)		,				max_length - doc_stride					)
								for i in doc_start_indices:
												A__						=	i + max_length - q_len
												A__						=	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": [-1_0_0] * len(UpperCamelCase__					),
								        "end_token": [-1_0_0] * len(UpperCamelCase__					),
								        "category": category,
								    },
								}

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

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

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

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

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

				if len(UpperCamelCase__					) <= 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"],
								    },
								}

				A__						=	input_ids[:q_len]
				A__						=	range(UpperCamelCase__		,				len(UpperCamelCase__					)		,				max_length - doc_stride					)

				A__						=	[]
				A__						=	[]
				A__						=	[]
				A__						=	[]  # null, yes, no, long, short
				for i in doc_start_indices:
								A__						=	i + max_length - q_len
								A__						=	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:
												A__						=	start_token - i + q_len
												A__						=	end_token - i + q_len
												answers_category.append(answer["category"][0]					)  # ["short"] -> "short"
								else:
												A__						=	-1_0_0
												A__						=	-1_0_0
												answers_category.append("null"					)
								A__						=	inputs[-1][start_token : end_token + 1]

								answers_start_token.append(UpperCamelCase__					)
								answers_end_token.append(UpperCamelCase__					)
								if assertion:
												if new != old and new != [tokenizer.cls_token_id]:
																print("ISSUE in strided for ID:"		,				example["id"]					)
																print("New:"		,				tokenizer.decode(UpperCamelCase__					)					)
																print("Old:"		,				tokenizer.decode(UpperCamelCase__					)		,				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     lowerCAmelCase__							(		SCREAMING_SNAKE_CASE_:							List[str]		,				SCREAMING_SNAKE_CASE_:							Union[str, Any]		,				SCREAMING_SNAKE_CASE_:							Dict=2_0_4_8		,				SCREAMING_SNAKE_CASE_:							List[str]=4_0_9_6		,				SCREAMING_SNAKE_CASE_:							Dict=False					)	->					Optional[Any]:




				'''simple docstring'''





				A__						=	get_strided_contexts_and_ans(
				    UpperCamelCase__		,				UpperCamelCase__		,				doc_stride=UpperCamelCase__		,				max_length=UpperCamelCase__		,				assertion=UpperCamelCase__		,				)

				return example



def     lowerCAmelCase__							(		SCREAMING_SNAKE_CASE_:							Any		,				SCREAMING_SNAKE_CASE_:							Tuple					)	->					List[str]:




				'''simple docstring'''





				with jsonlines.open(UpperCamelCase__		,				"a"					) as writer:
								for example in tqdm(UpperCamelCase__		,				total=len(UpperCamelCase__					)		,				desc="Saving samples ... "					):
												A__						=	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

	lowerCAmelCase__		      =       load_dataset("""natural_questions""")
	lowerCAmelCase__		      =       BigBirdTokenizer.from_pretrained("""google/bigbird-roberta-base""")

	lowerCAmelCase__		      =       data["""train""" if PROCESS_TRAIN == """true""" else """validation"""]

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

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




from typing import Dict, List, Optional

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


_UpperCAmelCase		:   Any										=       logging.get_logger(__name__)


_UpperCAmelCase		:   Dict										=       {
    """nielsr/canine-s""": 2048,
}

# Unicode defines 1,114,112 total “codepoints”
_UpperCAmelCase		:   Tuple										=       111_4112

# Below: Constants defining canonical codepoints for special, pseudo-characters.
# Copied from https://github.com/google-research/language/blob/master/language/canine/special_codepoints.py
_UpperCAmelCase		:   List[str]										=       0

_UpperCAmelCase		:   Any										=       0xE000
_UpperCAmelCase		:   Dict										=       0xE001
_UpperCAmelCase		:   Optional[int]										=       0xE002
_UpperCAmelCase		:   Tuple										=       0xE003
_UpperCAmelCase		:   Tuple										=       0xE004

# Maps special codepoints to human-readable names.
_UpperCAmelCase		:   Dict[int, str]										=       {
    # Special symbols are represented using codepoints values that are valid,
    # but designated as "Private Use", meaning that they will never be assigned
    # characters by the Unicode Consortium, and are thus safe for use here.
    #
    # NOTE: Do *NOT* add any sort of [UNK_CHAR] here. They are explicitly
    # excluded and should fail with a hard error.
    CLS: "[CLS]",
    SEP: "[SEP]",
    BOS: "[BOS]",
    MASK: "[MASK]",
    PAD: "[PAD]",
    RESERVED: "[RESERVED]",
}

# Maps special codepoint human-readable names to their codepoint values.
_UpperCAmelCase		:   Dict[str, int]										=       {name: codepoint for codepoint, name in SPECIAL_CODEPOINTS.items()}







class  lowercase				(       lowercase_			):
				__SCREAMING_SNAKE_CASE					: str       			=			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES




				def __init__(					self      ,   snake_case=chr(snake_case       )      ,   snake_case=chr(snake_case       )      ,   snake_case=chr(snake_case       )      ,   snake_case=chr(snake_case       )      ,   snake_case=chr(snake_case       )      ,   snake_case=chr(snake_case       )      ,   snake_case=False      ,   snake_case=2048      ,   **snake_case      ,   ):
											snake_case_							=    AddedToken(snake_case      ,   lstrip=snake_case      ,   rstrip=snake_case       ) if isinstance(snake_case      ,   snake_case       ) else bos_token
											snake_case_							=    AddedToken(snake_case      ,   lstrip=snake_case      ,   rstrip=snake_case       ) if isinstance(snake_case      ,   snake_case       ) else eos_token
											snake_case_							=    AddedToken(snake_case      ,   lstrip=snake_case      ,   rstrip=snake_case       ) if isinstance(snake_case      ,   snake_case       ) else sep_token
											snake_case_							=    AddedToken(snake_case      ,   lstrip=snake_case      ,   rstrip=snake_case       ) if isinstance(snake_case      ,   snake_case       ) else cls_token
											snake_case_							=    AddedToken(snake_case      ,   lstrip=snake_case      ,   rstrip=snake_case       ) if isinstance(snake_case      ,   snake_case       ) else pad_token

											# Mask token behave like a normal word, i.e. include the space before it
											snake_case_							=    AddedToken(snake_case      ,   lstrip=snake_case      ,   rstrip=snake_case       ) if isinstance(snake_case      ,   snake_case       ) else mask_token

											super().__init__(
											    bos_token=snake_case      ,   eos_token=snake_case      ,   sep_token=snake_case      ,   cls_token=snake_case      ,   pad_token=snake_case      ,   mask_token=snake_case      ,   add_prefix_space=snake_case      ,   model_max_length=snake_case      ,   **snake_case      ,   )

											# Creates a mapping for looking up the IDs of special symbols.
											snake_case_							=    {}
											for codepoint, name in SPECIAL_CODEPOINTS.items():
																		snake_case_							=    codepoint

											# Creates a mapping for looking up the string forms of special symbol IDs.
											snake_case_							=    {
											    codepoint: name for name, codepoint in self._special_codepoints.items()
											}

											snake_case_							=    UNICODE_VOCAB_SIZE
											snake_case_							=    len(self._special_codepoints       )




				@property
				def     a							(					self       ):
											return self._unicode_vocab_size




				def     a							(					self      ,   snake_case       ):
											return list(snake_case       )




				def     a							(					self      ,   snake_case       ):
											try:
																		return ord(snake_case       )
											except TypeError:
																		raise ValueError(F'''invalid token: \'{token}\''''       )




				def     a							(					self      ,   snake_case       ):
											try:
																		if index in SPECIAL_CODEPOINTS:
																									return SPECIAL_CODEPOINTS[index]
																		return chr(snake_case       )
											except TypeError:
																		raise ValueError(F'''invalid id: {index}'''       )




				def     a							(					self      ,   snake_case       ):
											return "".join(snake_case       )




				def     a							(					self      ,   snake_case      ,   snake_case = None       ):
											snake_case_							=    [self.sep_token_id]
											snake_case_							=    [self.cls_token_id]

											snake_case_							=    cls + token_ids_a + sep
											if token_ids_a is not None:
																		result += token_ids_a + sep
											return result




				def     a							(					self      ,   snake_case      ,   snake_case = None      ,   snake_case = False       ):
											if already_has_special_tokens:
																		return super().get_special_tokens_mask(
																		    token_ids_a=snake_case      ,   token_ids_a=snake_case      ,   already_has_special_tokens=snake_case       )

											snake_case_							=    [1] + ([0] * len(snake_case       )) + [1]
											if token_ids_a is not None:
																		result += ([0] * len(snake_case       )) + [1]
											return result




				def     a							(					self      ,   snake_case      ,   snake_case = None       ):
											snake_case_							=    [self.sep_token_id]
											snake_case_							=    [self.cls_token_id]

											snake_case_							=    len(cls + token_ids_a + sep       ) * [0]
											if token_ids_a is not None:
																		result += len(token_ids_a + sep       ) * [1]
											return result




				def     a							(					self      ,   snake_case      ,   snake_case = None       ):
											return ()




