



"""simple docstring"""





import os

import jsonlines
import numpy as np
from tqdm import tqdm


__lowerCAmelCase       :					Optional[int]      					=					2048
__lowerCAmelCase       :					int      					=					4096
__lowerCAmelCase       :					Union[str, Any]      					=					42
__lowerCAmelCase       :					List[Any]      					=					os.environ.pop('''PROCESS_TRAIN''', '''false''')
__lowerCAmelCase       :					Union[str, Any]      					=					{'''null''': 0, '''short''': 1, '''long''': 2, '''yes''': 3, '''no''': 4}
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Optional[int]    ):


						'''simple docstring'''






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

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

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

						snake_case_   : List[Any]						=			["""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						(		__UpperCamelCase      :			int					, __UpperCamelCase      :			Union[str, Any]=False    ):


						'''simple docstring'''






						snake_case_   : Dict						=			_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
												snake_case_   : Any						=			example["""document"""]["""tokens"""]
												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(__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

						snake_case_   : Tuple						=			["""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

						snake_case_   : Union[str, Any]						=			example["""document"""]["""tokens"""]
						snake_case_   : Dict						=			answer["""start_token"""]
						snake_case_   : Union[str, Any]						=			answer["""end_token"""]

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

						# checking above code
						if assertion:
												snake_case_   : List[Any]						=			doc["""is_html"""][answer["""start_token"""] : answer["""end_token"""]]
												snake_case_   : Dict						=			doc["""token"""][answer["""start_token"""] : answer["""end_token"""]]
												snake_case_   : Optional[Any]						=			""" """.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						(		__UpperCamelCase      :			Union[str, Any]					, __UpperCamelCase      :			int					, __UpperCamelCase      :			Dict=2_0_4_8					, __UpperCamelCase      :			Tuple=4_0_9_6					, __UpperCamelCase      :			Optional[int]=True    ):


						'''simple docstring'''






						snake_case_   : Tuple						=			get_context_and_ans(__UpperCamelCase					, assertion=__UpperCamelCase    )
						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"],
												    },
												}

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

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

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

						snake_case_   : List[str]						=			input_ids[answer["""start_token"""] : answer["""end_token"""] + 1]  # right & left are inclusive
						snake_case_   : Optional[int]						=			answer["""start_token"""]
						snake_case_   : List[str]						=			answer["""end_token"""]

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

						snake_case_   : int						=			input_ids[:q_len]
						snake_case_   : Tuple						=			range(__UpperCamelCase					, len(__UpperCamelCase    )					, max_length - doc_stride    )

						snake_case_   : Dict						=			[]
						snake_case_   : Optional[Any]						=			[]
						snake_case_   : Dict						=			[]
						snake_case_   : Dict						=			[]  # null, yes, no, long, short
						for i in doc_start_indices:
												snake_case_   : List[Any]						=			i + max_length - q_len
												snake_case_   : Tuple						=			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:
																		snake_case_   : Any						=			start_token - i + q_len
																		snake_case_   : Tuple						=			end_token - i + q_len
																		answers_category.append(answer["""category"""][0]    )  # ["short"] -> "short"
												else:
																		snake_case_   : Union[str, Any]						=			-1_0_0
																		snake_case_   : Union[str, Any]						=			-1_0_0
																		answers_category.append("""null"""    )
												snake_case_   : List[str]						=			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						(		__UpperCamelCase      :			Dict					, __UpperCamelCase      :			str					, __UpperCamelCase      :			int=2_0_4_8					, __UpperCamelCase      :			Tuple=4_0_9_6					, __UpperCamelCase      :			Union[str, Any]=False    ):


						'''simple docstring'''






						snake_case_   : int						=			get_strided_contexts_and_ans(
						    __UpperCamelCase					, __UpperCamelCase					, doc_stride=__UpperCamelCase					, max_length=__UpperCamelCase					, assertion=__UpperCamelCase					, )

						return example
def 	__lowerCAmelCase						(		__UpperCamelCase      :			List[Any]					, __UpperCamelCase      :			List[Any]    ):


						'''simple docstring'''






						with jsonlines.open(__UpperCamelCase					, """a"""    ) as writer:
												for example in tqdm(__UpperCamelCase					, total=len(__UpperCamelCase    )					, desc="""Saving samples ... """    ):
																		snake_case_   : int						=			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       :					Tuple      					=					load_dataset('''natural_questions''')
					__lowerCAmelCase       :					Optional[Any]      					=					BigBirdTokenizer.from_pretrained('''google/bigbird-roberta-base''')

					__lowerCAmelCase       :					Any      					=					data['''train''' if PROCESS_TRAIN == '''true''' else '''validation''']

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

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






"""simple docstring"""





import json
import os
import unittest

from transformers import MgpstrTokenizer
from transformers.models.mgp_str.tokenization_mgp_str import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin



@require_tokenizers
class    _lowerCAmelCase    (				SCREAMING_SNAKE_CASE__  ,				unittest.TestCase		):





		"""simple docstring"""






		_lowerCamelCase			  =       MgpstrTokenizer
		_lowerCamelCase			  =       False
		_lowerCamelCase			  =       {}
		_lowerCamelCase			  =       False





		def    UpperCAmelCase__							(					self   ) ->						Tuple:







								'''simple docstring'''

								super().setUp()

								# fmt: off
								snake_case_   : Optional[Any]						=			["""[GO]""", """[s]""", """0""", """1""", """2""", """3""", """4""", """5""", """6""", """7""", """8""", """9""", """a""", """b""", """c""", """d""", """e""", """f""", """g""", """h""", """i""", """j""", """k""", """l""", """m""", """n""", """o""", """p""", """q""", """r""", """s""", """t""", """u""", """v""", """w""", """x""", """y""", """z"""]
								# fmt: on
								snake_case_   : str						=			dict(zip(_lowercase							,    range(len(_lowercase   )   )   )   )

								snake_case_   : Tuple						=			os.path.join(self.tmpdirname							,    VOCAB_FILES_NAMES["""vocab_file"""]   )
								with open(self.vocab_file							,    """w"""							,    encoding="""utf-8"""   ) as fp:
														fp.write(json.dumps(_lowercase   ) + """\n"""   )





		def    UpperCAmelCase__							(					self							,    **_lowercase   ) ->						Union[str, Any]:







								'''simple docstring'''

								return MgpstrTokenizer.from_pretrained(self.tmpdirname							,    **_lowercase   )





		def    UpperCAmelCase__							(					self							,    _lowercase   ) ->						Optional[Any]:







								'''simple docstring'''

								snake_case_   : Tuple						=			"""tester"""
								snake_case_   : Tuple						=			"""tester"""
								return input_text, output_text





		@unittest.skip("""MGP-STR always lower cases letters."""   )
		def    UpperCAmelCase__							(					self   ) ->						List[Any]:







								'''simple docstring'''

								pass





		def    UpperCAmelCase__							(					self   ) ->						str:







								'''simple docstring'''

								snake_case_   : Dict						=			self.get_tokenizers(do_lower_case=_lowercase   )
								for tokenizer in tokenizers:
														with self.subTest(f'{tokenizer.__class__.__name__}'   ):
																				snake_case_   : Optional[Any]						=			"""[SPECIAL_TOKEN]"""

																				tokenizer.add_special_tokens({"""cls_token""": special_token}   )
																				snake_case_   : Union[str, Any]						=			tokenizer.encode([special_token]							,    add_special_tokens=_lowercase   )
																				self.assertEqual(len(_lowercase   )							,    1   )

																				snake_case_   : List[Any]						=			tokenizer.decode(_lowercase							,    skip_special_tokens=_lowercase   )
																				self.assertTrue(special_token not in decoded   )





		def    UpperCAmelCase__							(					self   ) ->						Dict:







								'''simple docstring'''

								snake_case_   : List[Any]						=			self.get_tokenizers()
								for tokenizer in tokenizers:
														with self.subTest(f'{tokenizer.__class__.__name__}'   ):
																				snake_case_			,		snake_case_   : Union[str, Any]						=			self.get_input_output_texts(_lowercase   )

																				snake_case_   : List[Any]						=			tokenizer.tokenize(_lowercase   )
																				snake_case_   : Union[str, Any]						=			tokenizer.convert_tokens_to_ids(_lowercase   )
																				snake_case_   : Union[str, Any]						=			tokenizer.encode(_lowercase							,    add_special_tokens=_lowercase   )
																				self.assertListEqual(_lowercase							,    _lowercase   )

																				snake_case_   : List[str]						=			tokenizer.convert_ids_to_tokens(_lowercase   )
																				self.assertNotEqual(len(_lowercase   )							,    0   )
																				snake_case_   : str						=			tokenizer.decode(_lowercase   )
																				self.assertIsInstance(_lowercase							,    _lowercase   )

																				self.assertEqual(text_a.replace(""" """							,    """"""   )							,    _lowercase   )





		@unittest.skip("""MGP-STR tokenizer only handles one sequence."""   )
		def    UpperCAmelCase__							(					self   ) ->						int:







								'''simple docstring'''

								pass





		@unittest.skip("""inputs cannot be pretokenized in MgpstrTokenizer"""   )
		def    UpperCAmelCase__							(					self   ) ->						Any:







								'''simple docstring'''

								pass

