





import os

import jsonlines
import numpy as np
from tqdm import tqdm


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





def    __UpperCamelCase      (    a) ->Optional[Any]:
					def choose_first(a,     a=False):
										assert isinstance(_A,     _A)
										if len(_A) == 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=_A)
															lowerCamelCase__       =       []
										answer.update(_A)

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

					return answer





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

					return {
					    "context": " ".join(_A),
					    "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      (    a,     a,     a=2048,     a=4096,     a=True) ->List[Any]:
					lowerCamelCase__       =       get_context_and_ans(_A,     assertion=_A)
					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(_A,     len(_A),     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(_A),
										        "end_token": [-100] * len(_A),
										        "category": category,
										    },
										}

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

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

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

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

					return example





def    __UpperCamelCase      (    a,     a) ->Optional[int]:
					with jsonlines.open(_A,     "a") as writer:
										for example in tqdm(_A,     total=len(_A),     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

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

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

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

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












from __future__ import annotations

import math
import random
from collections.abc import Collection
from typing import overload

class SCREAMING_SNAKE_CASE_ :





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase = None	):
										if components is None:
															lowerCamelCase__       =       []
										lowerCamelCase__       =       list(_lowerCAmelCase	)





					def __len__( self	):
										return len(self.__components	)





					def __str__( self	):
										return "(" + ",".join(map(_lowerCAmelCase	,  self.__components	)	) + ")"





					def __add__( self	,  _lowerCAmelCase	):
										lowerCamelCase__       =       len(self	)
										if size == len(_lowerCAmelCase	):
															lowerCamelCase__       =       [self.__components[i] + other.component(_lowerCAmelCase	) for i in range(_lowerCAmelCase	)]
															return Vector(_lowerCAmelCase	)
										else:
															raise Exception("must have the same size"	)





					def __sub__( self	,  _lowerCAmelCase	):
										lowerCamelCase__       =       len(self	)
										if size == len(_lowerCAmelCase	):
															lowerCamelCase__       =       [self.__components[i] - other.component(_lowerCAmelCase	) for i in range(_lowerCAmelCase	)]
															return Vector(_lowerCAmelCase	)
										else:  # error case
															raise Exception("must have the same size"	)





					@overload
					def __mul__( self	,  _lowerCAmelCase	):
										...





					@overload
					def __mul__( self	,  _lowerCAmelCase	):
										...





					def __mul__( self	,  _lowerCAmelCase	):
										if isinstance(_lowerCAmelCase	,  (float, int)	):
															lowerCamelCase__       =       [c * other for c in self.__components]
															return Vector(_lowerCAmelCase	)
										elif isinstance(_lowerCAmelCase	,  _lowerCAmelCase	) and len(self	) == len(_lowerCAmelCase	):
															lowerCamelCase__       =       len(self	)
															lowerCamelCase__       =       [self.__components[i] * other.component(_lowerCAmelCase	) for i in range(_lowerCAmelCase	)]
															return sum(_lowerCAmelCase	)
										else:  # error case
															raise Exception("invalid operand!"	)





					def 					__magic_name__ ( self	):
										return Vector(self.__components	)





					def 					__magic_name__ ( self	,  _lowerCAmelCase	):
										if isinstance(_lowerCAmelCase	,  _lowerCAmelCase	) and -len(self.__components	) <= i < len(self.__components	):
															return self.__components[i]
										else:
															raise Exception("index out of range"	)





					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										assert -len(self.__components	) <= pos < len(self.__components	)
										lowerCamelCase__       =       value





					def 					__magic_name__ ( self	):
										if len(self.__components	) == 0:
															raise Exception("Vector is empty"	)
										lowerCamelCase__       =       [c**2 for c in self.__components]
										return math.sqrt(sum(_lowerCAmelCase	)	)






					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase = False	):
										lowerCamelCase__       =       self * other
										lowerCamelCase__       =       self.euclidean_length() * other.euclidean_length()
										if deg:
															return math.degrees(math.acos(num / den	)	)
										else:
															return math.acos(num / den	)





def    __UpperCamelCase      (    a) ->Vector:
					assert isinstance(a,     a)
					return Vector([0] * dimension)





def    __UpperCamelCase      (    a,     a) ->Vector:
					assert isinstance(a,     a) and (isinstance(a,     a))
					lowerCamelCase__       =       [0] * dimension
					lowerCamelCase__       =       1
					return Vector(a)





def    __UpperCamelCase      (    a,     a,     a) ->Vector:
					assert (
					    isinstance(a,     a)
					    and isinstance(a,     a)
					    and (isinstance(a,     (int, float)))
					)
					return x * scalar + y





def    __UpperCamelCase      (    a,     a,     a) ->Vector:
					random.seed(a)
					lowerCamelCase__       =       [random.randint(a,     a) for _ in range(a)]
					return Vector(a)

class SCREAMING_SNAKE_CASE_ :





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__       =       matrix
										lowerCamelCase__       =       w
										lowerCamelCase__       =       h





					def __str__( self	):
										lowerCamelCase__       =       ""
										for i in range(self.__height	):
															ans += "|"
															for j in range(self.__width	):
																				if j < self.__width - 1:
																									ans += str(self.__matrix[i][j]	) + ","
																				else:
																									ans += str(self.__matrix[i][j]	) + "|\n"
										return ans





					def __add__( self	,  _lowerCAmelCase	):
										if self.__width == other.width() and self.__height == other.height():
															lowerCamelCase__       =       []
															for i in range(self.__height	):
																				lowerCamelCase__       =       [
																				    self.__matrix[i][j] + other.component(_lowerCAmelCase	,  _lowerCAmelCase	)
																				    for j in range(self.__width	)
																				]
																				matrix.append(_lowerCAmelCase	)
															return Matrix(_lowerCAmelCase	,  self.__width	,  self.__height	)
										else:
															raise Exception("matrix must have the same dimension!"	)





					def __sub__( self	,  _lowerCAmelCase	):
										if self.__width == other.width() and self.__height == other.height():
															lowerCamelCase__       =       []
															for i in range(self.__height	):
																				lowerCamelCase__       =       [
																				    self.__matrix[i][j] - other.component(_lowerCAmelCase	,  _lowerCAmelCase	)
																				    for j in range(self.__width	)
																				]
																				matrix.append(_lowerCAmelCase	)
															return Matrix(_lowerCAmelCase	,  self.__width	,  self.__height	)
										else:
															raise Exception("matrices must have the same dimension!"	)





					@overload
					def __mul__( self	,  _lowerCAmelCase	):
										...





					@overload
					def __mul__( self	,  _lowerCAmelCase	):
										...





					def __mul__( self	,  _lowerCAmelCase	):
										if isinstance(_lowerCAmelCase	,  _lowerCAmelCase	):  # matrix-vector
															if len(_lowerCAmelCase	) == self.__width:
																				lowerCamelCase__       =       zero_vector(self.__height	)
																				for i in range(self.__height	):
																									lowerCamelCase__       =       [
																									    self.__matrix[i][j] * other.component(_lowerCAmelCase	)
																									    for j in range(self.__width	)
																									]
																									ans.change_component(_lowerCAmelCase	,  sum(_lowerCAmelCase	)	)
																				return ans
															else:
																				raise Exception(
																				    "vector must have the same size as the "
																				    "number of columns of the matrix!"	)
										elif isinstance(_lowerCAmelCase	,  (int, float)	):  # matrix-scalar
															lowerCamelCase__       =       [
															    [self.__matrix[i][j] * other for j in range(self.__width	)]
															    for i in range(self.__height	)
															]
															return Matrix(_lowerCAmelCase	,  self.__width	,  self.__height	)
										return None





					def 					__magic_name__ ( self	):
										return self.__height





					def 					__magic_name__ ( self	):
										return self.__width





					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										if 0 <= x < self.__height and 0 <= y < self.__width:
															return self.__matrix[x][y]
										else:
															raise Exception("change_component: indices out of bounds"	)





					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										if 0 <= x < self.__height and 0 <= y < self.__width:
															lowerCamelCase__       =       value
										else:
															raise Exception("change_component: indices out of bounds"	)





					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										if self.__height != self.__width:
															raise Exception("Matrix is not square"	)
										lowerCamelCase__       =       self.__matrix[:x] + self.__matrix[x + 1 :]
										for i in range(len(_lowerCAmelCase	)	):
															lowerCamelCase__       =       minor[i][:y] + minor[i][y + 1 :]
										return Matrix(_lowerCAmelCase	,  self.__width - 1	,  self.__height - 1	).determinant()





					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										if self.__height != self.__width:
															raise Exception("Matrix is not square"	)
										if 0 <= x < self.__height and 0 <= y < self.__width:
															return (-1) ** (x + y) * self.minor(_lowerCAmelCase	,  _lowerCAmelCase	)
										else:
															raise Exception("Indices out of bounds"	)






					def 					__magic_name__ ( self	):
										if self.__height != self.__width:
															raise Exception("Matrix is not square"	)
										if self.__height < 1:
															raise Exception("Matrix has no element"	)
										elif self.__height == 1:
															return self.__matrix[0][0]
										elif self.__height == 2:
															return (
															    self.__matrix[0][0] * self.__matrix[1][1]
															    - self.__matrix[0][1] * self.__matrix[1][0]
															)
										else:
															lowerCamelCase__       =       [
															    self.__matrix[0][y] * self.cofactor(0	,  _lowerCAmelCase	) for y in range(self.__width	)
															]
															return sum(_lowerCAmelCase	)





def    __UpperCamelCase      (    a) ->Matrix:
					lowerCamelCase__       =       [[0] * n for _ in range(a)]
					return Matrix(a,     a,     a)







def    __UpperCamelCase      (    a,     a,     a,     a) ->Matrix:
					random.seed(a)
					lowerCamelCase__       =       [
					    [random.randint(a,     a) for _ in range(a)] for _ in range(a)
					]
					return Matrix(a,     a,     a)

