"""simple docstring"""





import os

import jsonlines
import numpy as np
from tqdm import tqdm


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



def      lowerCamelCase_(__SCREAMING_SNAKE_CASE  )->    Optional[int]:
							def choose_first(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE=False  ):
														assert isinstance(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE  )
														if len(__SCREAMING_SNAKE_CASE  ) == 1:
																					_SCREAMING_SNAKE_CASE		:						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		:						Tuple			  =    {"""id""": example["""id"""]}
							_SCREAMING_SNAKE_CASE		:						List[str]			  =    example["""annotations"""]
							_SCREAMING_SNAKE_CASE		:						List[str]			  =    annotation["""yes_no_answer"""]
							if 0 in yes_no_answer or 1 in yes_no_answer:
														_SCREAMING_SNAKE_CASE		:						List[Any]			  =    ["""yes"""] if 1 in yes_no_answer else ["""no"""]
														_SCREAMING_SNAKE_CASE		:						int			  =    []
														_SCREAMING_SNAKE_CASE		:						Any			  =    []
														_SCREAMING_SNAKE_CASE		:						str			  =    ["""<cls>"""]
							else:
														_SCREAMING_SNAKE_CASE		:						int			  =    ["""short"""]
														_SCREAMING_SNAKE_CASE		:						Any			  =    choose_first(annotation["""short_answers"""]  )
														if len(out["""start_token"""]  ) == 0:
																					# answer will be long if short is not available
																					_SCREAMING_SNAKE_CASE		:						Any			  =    ["""long"""]
																					_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    choose_first(annotation["""long_answer"""]     ,	is_long_answer=__SCREAMING_SNAKE_CASE  )
																					_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    []
														answer.update(__SCREAMING_SNAKE_CASE  )

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

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

							return answer



def      lowerCamelCase_(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE=False  )->    Optional[int]:
							_SCREAMING_SNAKE_CASE		:						List[str]			  =    _get_single_answer(__SCREAMING_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		:						Optional[Any]			  =    example["""document"""]["""tokens"""]
														_SCREAMING_SNAKE_CASE		:						Any			  =    []
														for i in range(len(doc["""token"""]  )  ):
																					if not doc["is_html"][i]:
																												context.append(doc["""token"""][i]  )
														return {
														    "context": " ".join(__SCREAMING_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		:						Any			  =    ["""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		:						List[str]			  =    answer["""start_token"""]
							_SCREAMING_SNAKE_CASE		:						int			  =    answer["""end_token"""]

							_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    []
							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		:						Optional[Any]			  =    """ """.join(context[start_token:end_token]  )

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

							return {
							    "context": " ".join(__SCREAMING_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      lowerCamelCase_(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE=2_048     ,	__SCREAMING_SNAKE_CASE=4_096     ,	__SCREAMING_SNAKE_CASE=True  )->    Union[str, Any]:
							# overlap will be of doc_stride - q_len

							_SCREAMING_SNAKE_CASE		:						List[str]			  =    get_context_and_ans(__SCREAMING_SNAKE_CASE     ,	assertion=__SCREAMING_SNAKE_CASE  )
							_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    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		:						Optional[int]			  =    tokenizer(example["""question"""]["""text"""]     ,	out["""context"""]  ).input_ids
							_SCREAMING_SNAKE_CASE		:						List[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		:						Union[str, Any]			  =    []
														_SCREAMING_SNAKE_CASE		:						List[Any]			  =    input_ids[:q_len]
														_SCREAMING_SNAKE_CASE		:						List[str]			  =    range(__SCREAMING_SNAKE_CASE     ,	len(__SCREAMING_SNAKE_CASE  )     ,	max_length - doc_stride  )
														for i in doc_start_indices:
																					_SCREAMING_SNAKE_CASE		:						int			  =    i + max_length - q_len
																					_SCREAMING_SNAKE_CASE		:						str			  =    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(__SCREAMING_SNAKE_CASE  ),
														        "end_token": [-100] * len(__SCREAMING_SNAKE_CASE  ),
														        "category": category,
														    },
														}

							_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    out["""context"""].split()
							_SCREAMING_SNAKE_CASE		:						str			  =    splitted_context[answer["""end_token"""]]
							_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    len(
							    tokenizer(
							        """ """.join(splitted_context[: answer["""start_token"""]]  )     ,	add_special_tokens=__SCREAMING_SNAKE_CASE     ,	).input_ids  )
							_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    len(
							    tokenizer(""" """.join(splitted_context[: answer["""end_token"""]]  )     ,	add_special_tokens=__SCREAMING_SNAKE_CASE  ).input_ids  )

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

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

							_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    input_ids[answer["""start_token"""] : answer["""end_token"""] + 1]  # right & left are inclusive
							_SCREAMING_SNAKE_CASE		:						List[str]			  =    answer["""start_token"""]
							_SCREAMING_SNAKE_CASE		:						Any			  =    answer["""end_token"""]

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

							if len(__SCREAMING_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		:						List[str]			  =    input_ids[:q_len]
							_SCREAMING_SNAKE_CASE		:						str			  =    range(__SCREAMING_SNAKE_CASE     ,	len(__SCREAMING_SNAKE_CASE  )     ,	max_length - doc_stride  )

							_SCREAMING_SNAKE_CASE		:						Dict			  =    []
							_SCREAMING_SNAKE_CASE		:						int			  =    []
							_SCREAMING_SNAKE_CASE		:						List[Any]			  =    []
							_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    []  # null, yes, no, long, short
							for i in doc_start_indices:
														_SCREAMING_SNAKE_CASE		:						str			  =    i + max_length - q_len
														_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    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		:						Any			  =    start_token - i + q_len
																					_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    end_token - i + q_len
																					answers_category.append(answer["""category"""][0]  )  # ["short"] -> "short"
														else:
																					_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    -100
																					_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    -100
																					answers_category.append("""null"""  )
														_SCREAMING_SNAKE_CASE		:						Any			  =    inputs[-1][start_token : end_token + 1]

														answers_start_token.append(__SCREAMING_SNAKE_CASE  )
														answers_end_token.append(__SCREAMING_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(__SCREAMING_SNAKE_CASE  )  )
																												print("""Old:"""     ,	tokenizer.decode(__SCREAMING_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      lowerCamelCase_(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE=2_048     ,	__SCREAMING_SNAKE_CASE=4_096     ,	__SCREAMING_SNAKE_CASE=False  )->    Any:
							_SCREAMING_SNAKE_CASE		:						List[str]			  =    get_strided_contexts_and_ans(
							    __SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	doc_stride=__SCREAMING_SNAKE_CASE     ,	max_length=__SCREAMING_SNAKE_CASE     ,	assertion=__SCREAMING_SNAKE_CASE     ,	)

							return example



def      lowerCamelCase_(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE  )->    Union[str, Any]:
							with jsonlines.open(__SCREAMING_SNAKE_CASE     ,	"""a"""  ) as writer:
														for example in tqdm(__SCREAMING_SNAKE_CASE     ,	total=len(__SCREAMING_SNAKE_CASE  )     ,	desc="""Saving samples ... """  ):
																					_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    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)



"""simple docstring"""





import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding




class  _snake_case		(    __snake_case ):







				"""simple docstring"""






				a			=						["image_processor", "tokenizer"]
				a			=						"ChineseCLIPImageProcessor"
				a			=						("BertTokenizer", "BertTokenizerFast")





				def __init__(      self							:							Dict      ,    _A							:							Tuple=None      ,    _A							:							List[Any]=None      ,    **_A							:							int):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:						Tuple			  =    None
											if "feature_extractor" in kwargs:
																		warnings.warn(
																		    """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
																		    """ instead."""      ,    _A      ,    )
																		_SCREAMING_SNAKE_CASE		:						str			  =    kwargs.pop("""feature_extractor""")

											_SCREAMING_SNAKE_CASE		:						int			  =    image_processor if image_processor is not None else feature_extractor
											if image_processor is None:
																		raise ValueError("""You need to specify an `image_processor`.""")
											if tokenizer is None:
																		raise ValueError("""You need to specify a `tokenizer`.""")

											super().__init__(_A      ,    _A)
											_SCREAMING_SNAKE_CASE		:						Dict			  =    self.image_processor





				def __call__(      self							:							Optional[int]      ,    _A							:							Optional[Any]=None      ,    _A							:							Any=None      ,    _A							:							Tuple=None      ,    **_A							:							int):



											"""simple docstring"""







											if text is None and images is None:
																		raise ValueError("""You have to specify either text or images. Both cannot be none.""")

											if text is not None:
																		_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    self.tokenizer(_A      ,    return_tensors=_A      ,    **_A)

											if images is not None:
																		_SCREAMING_SNAKE_CASE		:						List[Any]			  =    self.image_processor(_A      ,    return_tensors=_A      ,    **_A)

											if text is not None and images is not None:
																		_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    image_features.pixel_values
																		return encoding
											elif text is not None:
																		return encoding
											else:
																		return BatchEncoding(data=dict(**_A)      ,    tensor_type=_A)





				def 		_lowerCAmelCase		(      self							:							str      ,    *_A							:							Any      ,    **_A							:							Any):



											"""simple docstring"""







											return self.tokenizer.batch_decode(*_A      ,    **_A)





				def 		_lowerCAmelCase		(      self							:							Union[str, Any]      ,    *_A							:							List[Any]      ,    **_A							:							Any):



											"""simple docstring"""







											return self.tokenizer.decode(*_A      ,    **_A)





				@property
				def 		_lowerCAmelCase		(      self							:							str):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    self.tokenizer.model_input_names
											_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    self.image_processor.model_input_names
											return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))





				@property
				def 		_lowerCAmelCase		(      self							:							List[str]):



											"""simple docstring"""







											warnings.warn(
											    """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""      ,    _A      ,    )
											return self.image_processor_class


