
"""simple docstring"""




import os

import jsonlines
import numpy as np
from tqdm import tqdm


UpperCAmelCase__   :					List[str]												=				2_0_4_8
UpperCAmelCase__   :					Optional[Any]												=				4_0_9_6
UpperCAmelCase__   :					int												=				4_2
UpperCAmelCase__   :					List[str]												=				os.environ.pop('PROCESS_TRAIN', 'false')
UpperCAmelCase__   :					Optional[Any]												=				{'null': 0, 'short': 1, 'long': 2, 'yes': 3, 'no': 4}

def lowercase_						(     _snake_case		):
							def choose_first(_snake_case     ,_snake_case=False		):
														assert isinstance(_snake_case     ,_snake_case		)
														if len(_snake_case		) == 1:
																					SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     answer[0]
																					return {k: [answer[k]] for k in answer} if is_long_answer else answer
														for a in answer:
																					if is_long_answer:
																												SCREAMING_SNAKE_CASE__		: Optional[Any]     =     {k: [a[k]] for k in a}
																					if len(a["""start_token"""]		) > 0:
																												break
														return a

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

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

							SCREAMING_SNAKE_CASE__		: List[str]     =     ["""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 lowercase_						(     _snake_case     ,_snake_case=False		):
							SCREAMING_SNAKE_CASE__		: List[str]     =     _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
														SCREAMING_SNAKE_CASE__		: Any     =     example["""document"""]["""tokens"""]
														SCREAMING_SNAKE_CASE__		: Optional[int]     =     []
														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

							SCREAMING_SNAKE_CASE__		: int     =     ["""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

							SCREAMING_SNAKE_CASE__		: List[Any]     =     example["""document"""]["""tokens"""]
							SCREAMING_SNAKE_CASE__		: Dict     =     answer["""start_token"""]
							SCREAMING_SNAKE_CASE__		: Tuple     =     answer["""end_token"""]

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

							# checking above code
							if assertion:
														SCREAMING_SNAKE_CASE__		: List[str]     =     doc["""is_html"""][answer["""start_token"""] : answer["""end_token"""]]
														SCREAMING_SNAKE_CASE__		: Dict     =     doc["""token"""][answer["""start_token"""] : answer["""end_token"""]]
														SCREAMING_SNAKE_CASE__		: Any     =     """ """.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 lowercase_						(     _snake_case     ,_snake_case     ,_snake_case=2_048     ,_snake_case=4_096     ,_snake_case=True		):
							# overlap will be of doc_stride - q_len

							SCREAMING_SNAKE_CASE__		: Dict     =     get_context_and_ans(_snake_case     ,assertion=_snake_case		)
							SCREAMING_SNAKE_CASE__		: List[Any]     =     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"],
														    },
														}

							SCREAMING_SNAKE_CASE__		: str     =     tokenizer(example["""question"""]["""text"""]     ,out["""context"""]		).input_ids
							SCREAMING_SNAKE_CASE__		: Optional[Any]     =     input_ids.index(tokenizer.sep_token_id		) + 1

							# return yes/no
							if answer["category"][0] in ["yes", "no"]:  # category is list with one element
														SCREAMING_SNAKE_CASE__		: Optional[Any]     =     []
														SCREAMING_SNAKE_CASE__		: List[Any]     =     []
														SCREAMING_SNAKE_CASE__		: Any     =     input_ids[:q_len]
														SCREAMING_SNAKE_CASE__		: Dict     =     range(_snake_case     ,len(_snake_case		)     ,max_length - doc_stride		)
														for i in doc_start_indices:
																					SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     i + max_length - q_len
																					SCREAMING_SNAKE_CASE__		: Any     =     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,
														    },
														}

							SCREAMING_SNAKE_CASE__		: Dict     =     out["""context"""].split()
							SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     splitted_context[answer["""end_token"""]]
							SCREAMING_SNAKE_CASE__		: Any     =     len(
							    tokenizer(
							        """ """.join(splitted_context[: answer["""start_token"""]]		)     ,add_special_tokens=_snake_case     ,).input_ids		)
							SCREAMING_SNAKE_CASE__		: int     =     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
							SCREAMING_SNAKE_CASE__		: List[str]     =     len(tokenizer(_snake_case     ,add_special_tokens=_snake_case		).input_ids		)
							if num_sub_tokens > 1:
														answer["end_token"] += num_sub_tokens - 1

							SCREAMING_SNAKE_CASE__		: str     =     input_ids[answer["""start_token"""] : answer["""end_token"""] + 1]  # right & left are inclusive
							SCREAMING_SNAKE_CASE__		: List[str]     =     answer["""start_token"""]
							SCREAMING_SNAKE_CASE__		: Tuple     =     answer["""end_token"""]

							if assertion:
														SCREAMING_SNAKE_CASE__		: int     =     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"],
														    },
														}

							SCREAMING_SNAKE_CASE__		: Dict     =     input_ids[:q_len]
							SCREAMING_SNAKE_CASE__		: int     =     range(_snake_case     ,len(_snake_case		)     ,max_length - doc_stride		)

							SCREAMING_SNAKE_CASE__		: Tuple     =     []
							SCREAMING_SNAKE_CASE__		: str     =     []
							SCREAMING_SNAKE_CASE__		: Dict     =     []
							SCREAMING_SNAKE_CASE__		: Tuple     =     []  # null, yes, no, long, short
							for i in doc_start_indices:
														SCREAMING_SNAKE_CASE__		: str     =     i + max_length - q_len
														SCREAMING_SNAKE_CASE__		: str     =     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:
																					SCREAMING_SNAKE_CASE__		: int     =     start_token - i + q_len
																					SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     end_token - i + q_len
																					answers_category.append(answer["""category"""][0]		)  # ["short"] -> "short"
														else:
																					SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     -100
																					SCREAMING_SNAKE_CASE__		: List[str]     =     -100
																					answers_category.append("""null"""		)
														SCREAMING_SNAKE_CASE__		: str     =     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 lowercase_						(     _snake_case     ,_snake_case     ,_snake_case=2_048     ,_snake_case=4_096     ,_snake_case=False		):
							SCREAMING_SNAKE_CASE__		: Tuple     =     get_strided_contexts_and_ans(
							    _snake_case     ,_snake_case     ,doc_stride=_snake_case     ,max_length=_snake_case     ,assertion=_snake_case     ,)

							return example

def lowercase_						(     _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 ... """		):
																					SCREAMING_SNAKE_CASE__		: Tuple     =     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__   :					Any												=				load_dataset('natural_questions')
					UpperCAmelCase__   :					Any												=				BigBirdTokenizer.from_pretrained('google/bigbird-roberta-base')

					UpperCAmelCase__   :					List[str]												=				data['train' if PROCESS_TRAIN == 'true' else 'validation']

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

					np.random.seed(SEED)
					UpperCAmelCase__   :					Union[str, Any]												=				'nq-training.jsonl' if PROCESS_TRAIN == 'true' else 'nq-validation.jsonl'
					save_to_disk(data, file_name=cache_file_name)






"""simple docstring"""




from manim import *



class 			lowerCAmelCase_		(a__						):


				"""simple docstring"""



				def 			__magic_name__					(self					)      ->	List[str]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Dict     =     Rectangle(height=0.5     , width=0.5					)
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     Rectangle(height=0.46     , width=0.46					).set_stroke(width=0					)

											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     [mem.copy() for i in range(6					)]
											SCREAMING_SNAKE_CASE__		: str     =     [mem.copy() for i in range(6					)]
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     VGroup(*SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0					)
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     VGroup(*SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0					)
											SCREAMING_SNAKE_CASE__		: Tuple     =     VGroup(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0					)
											SCREAMING_SNAKE_CASE__		: int     =     Text("""CPU"""     , font_size=24					)
											SCREAMING_SNAKE_CASE__		: Any     =     Group(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0.5     , aligned_edge=SCREAMING_SNAKE_CASE__					)
											cpu.move_to([-2.5, -0.5, 0]					)
											self.add(SCREAMING_SNAKE_CASE__					)

											SCREAMING_SNAKE_CASE__		: Optional[int]     =     [mem.copy() for i in range(1					)]
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     VGroup(*SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0					)
											SCREAMING_SNAKE_CASE__		: str     =     Text("""GPU"""     , font_size=24					)
											SCREAMING_SNAKE_CASE__		: str     =     Group(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0.5     , aligned_edge=SCREAMING_SNAKE_CASE__					)
											gpu.align_to(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)
											gpu.set_x(gpu.get_x() - 1					)

											self.add(SCREAMING_SNAKE_CASE__					)

											SCREAMING_SNAKE_CASE__		: List[Any]     =     [mem.copy() for i in range(6					)]
											SCREAMING_SNAKE_CASE__		: List[Any]     =     VGroup(*SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0					)

											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     Text("""Model"""     , font_size=24					)
											SCREAMING_SNAKE_CASE__		: Any     =     Group(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0.5     , aligned_edge=SCREAMING_SNAKE_CASE__					)
											model.move_to([3, -1.0, 0]					)

											self.play(
											    Create(SCREAMING_SNAKE_CASE__     , run_time=1					)     , Create(SCREAMING_SNAKE_CASE__     , run_time=1					)     , Create(SCREAMING_SNAKE_CASE__     , run_time=1					)     , )

											SCREAMING_SNAKE_CASE__		: Tuple     =     MarkupText(
											    F'''First, an empty model skeleton is loaded\ninto <span fgcolor=\'{YELLOW}\'>memory</span> without using much RAM.'''     , font_size=24     , )

											SCREAMING_SNAKE_CASE__		: Dict     =     Square(side_length=2.2					)
											key.move_to([-5, 2, 0]					)

											SCREAMING_SNAKE_CASE__		: Tuple     =     MarkupText(
											    F'''<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'''     , font_size=18     , )

											key_text.move_to([-5, 2.4, 0]					)

											step_a.move_to([2, 2, 0]					)
											self.play(Write(SCREAMING_SNAKE_CASE__     , run_time=2.5					)     , Write(SCREAMING_SNAKE_CASE__					)     , Write(SCREAMING_SNAKE_CASE__					)					)

											self.add(SCREAMING_SNAKE_CASE__					)

											SCREAMING_SNAKE_CASE__		: Optional[int]     =     []
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     []
											SCREAMING_SNAKE_CASE__		: str     =     []
											for i, rect in enumerate(SCREAMING_SNAKE_CASE__					):
																		SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     Rectangle(height=0.46     , width=0.46					).set_stroke(width=0.0					).set_fill(SCREAMING_SNAKE_CASE__     , opacity=0.7					)
																		cpu_target.move_to(SCREAMING_SNAKE_CASE__					)
																		cpu_target.generate_target()
																		SCREAMING_SNAKE_CASE__		: Tuple     =     0.46 / 4
																		SCREAMING_SNAKE_CASE__		: int     =     0.46 / 3

																		if i == 0:
																									cpu_target.target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT					)     , buff=0.02     , direction=SCREAMING_SNAKE_CASE__					)
																									cpu_target.target.set_x(cpu_target.target.get_x() + 0.1					)
																		elif i == 3:
																									cpu_target.target.next_to(cpu_targs[0].target     , direction=SCREAMING_SNAKE_CASE__     , buff=0.0					)
																		else:
																									cpu_target.target.next_to(cpu_targs[i - 1].target     , direction=SCREAMING_SNAKE_CASE__     , buff=0.0					)
																		cpu_targs.append(SCREAMING_SNAKE_CASE__					)

																		first_animations.append(rect.animate(run_time=0.5					).set_stroke(SCREAMING_SNAKE_CASE__					)					)
																		second_animations.append(MoveToTarget(SCREAMING_SNAKE_CASE__     , run_time=1.5					)					)

											self.play(*SCREAMING_SNAKE_CASE__					)
											self.play(*SCREAMING_SNAKE_CASE__					)

											self.wait()




