
"""simple docstring"""






import unittest
from pathlib import Path
from shutil import copyfile

from transformers import SPIECE_UNDERLINE, is_sentencepiece_available
from transformers.models.speech_to_text import SpeechaTextTokenizer
from transformers.models.speech_to_text.tokenization_speech_to_text import VOCAB_FILES_NAMES, save_json
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


_lowercase							=				get_tests_dir('''fixtures/test_sentencepiece.model''')

if is_sentencepiece_available():
					import sentencepiece as sp


_lowercase							=				5
_lowercase							=				10

@require_sentencepiece
@require_tokenizers
class 					__a     ( __a					,   unittest.TestCase       ):



							'''simple docstring'''







							_lowerCamelCase					:			Dict          =   SpeechaTextTokenizer
							_lowerCamelCase					:			Optional[int]          =   False
							_lowerCamelCase					:			Union[str, Any]          =   True
							def   SCREAMING_SNAKE_CASE						(     self       )			->							List[Any]:




										'''simple docstring'''
										super().setUp()

										__lowercase       =       sp.SentencePieceProcessor()
										spm_model.Load(_lowerCamelCase       )
										__lowercase       =       ["<s>", "<pad>", "</s>", "<unk>"]

										vocab += [spm_model.IdToPiece(id_       ) for id_ in range(len(_lowerCamelCase       )       )]
										__lowercase       =       dict(zip(_lowerCamelCase  ,    range(len(_lowerCamelCase       )       )       )       )

										__lowercase       =       Path(self.tmpdirname       )
										save_json(_lowerCamelCase  ,    save_dir / VOCAB_FILES_NAMES["vocab_file"]       )
										if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
													copyfile(_lowerCamelCase  ,    save_dir / VOCAB_FILES_NAMES["spm_file"]       )

										__lowercase       =       SpeechaTextTokenizer.from_pretrained(self.tmpdirname       )
										tokenizer.save_pretrained(self.tmpdirname       )
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Optional[Any]:




										'''simple docstring'''
										__lowercase       =       "<pad>"
										__lowercase       =       1

										self.assertEqual(self.get_tokenizer()._convert_token_to_id(_lowerCamelCase       )  ,    _lowerCamelCase       )
										self.assertEqual(self.get_tokenizer()._convert_id_to_token(_lowerCamelCase       )  ,    _lowerCamelCase       )
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Optional[int]:




										'''simple docstring'''
										__lowercase       =       list(self.get_tokenizer().get_vocab().keys()       )

										self.assertEqual(vocab_keys[0]  ,    "<s>"       )
										self.assertEqual(vocab_keys[1]  ,    "<pad>"       )
										self.assertEqual(vocab_keys[-1]  ,    "j"       )
										self.assertEqual(len(_lowerCamelCase       )  ,    1_001       )
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Optional[Any]:




										'''simple docstring'''
										self.assertEqual(self.get_tokenizer().vocab_size  ,    1_001       )
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Dict:




										'''simple docstring'''
										__lowercase       =       SpeechaTextTokenizer.from_pretrained(self.tmpdirname       )

										__lowercase       =       tokenizer.tokenize("This is a test"       )
										self.assertListEqual(_lowerCamelCase  ,    ["▁This", "▁is", "▁a", "▁t", "est"]       )

										self.assertListEqual(
										    tokenizer.convert_tokens_to_ids(_lowerCamelCase       )  ,    [289, 50, 14, 174, 386]  ,    )

										__lowercase       =       tokenizer.tokenize("I was born in 92000, and this is falsé."       )
										self.assertListEqual(
										    _lowerCamelCase  ,    [SPIECE_UNDERLINE + "I", SPIECE_UNDERLINE + "was", SPIECE_UNDERLINE + "b", "or", "n", SPIECE_UNDERLINE + "in", SPIECE_UNDERLINE + "", "9", "2", "0", "0", "0", ",", SPIECE_UNDERLINE + "and", SPIECE_UNDERLINE + "this", SPIECE_UNDERLINE + "is", SPIECE_UNDERLINE + "f", "al", "s", "é", "."]  ,    )
										__lowercase       =       tokenizer.convert_tokens_to_ids(_lowerCamelCase       )
										self.assertListEqual(_lowerCamelCase  ,    [12, 25, 88, 59, 28, 23, 11, 4, 606, 351, 351, 351, 7, 16, 70, 50, 76, 84, 10, 4, 8]       )

										__lowercase       =       tokenizer.convert_ids_to_tokens(_lowerCamelCase       )
										self.assertListEqual(
										    _lowerCamelCase  ,    [SPIECE_UNDERLINE + "I", SPIECE_UNDERLINE + "was", SPIECE_UNDERLINE + "b", "or", "n", SPIECE_UNDERLINE + "in", SPIECE_UNDERLINE + "", "<unk>", "2", "0", "0", "0", ",", SPIECE_UNDERLINE + "and", SPIECE_UNDERLINE + "this", SPIECE_UNDERLINE + "is", SPIECE_UNDERLINE + "f", "al", "s", "<unk>", "."]  ,    )





							@slow
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Optional[Any]:




										'''simple docstring'''
										# fmt: off
										__lowercase       =       {"input_ids": [[3_791, 797, 31, 11, 64, 797, 31, 2_429, 433, 12, 1_176, 12, 20, 786, 915, 142, 2_413, 240, 37, 3_238, 797, 31, 11, 35, 93, 915, 142, 2_413, 240, 37, 5_540, 567, 1_276, 93, 37, 610, 40, 62, 455, 657, 1_042, 123, 780, 177, 37, 309, 241, 1_298, 514, 20, 292, 2_737, 114, 2_469, 241, 85, 64, 302, 548, 528, 423, 4, 509, 406, 423, 37, 601, 4, 777, 302, 548, 528, 423, 284, 4, 3_388, 511, 459, 4, 3_555, 40, 321, 302, 705, 4, 3_388, 511, 583, 326, 5, 5, 5, 62, 3_310, 560, 177, 2_680, 217, 1_508, 32, 31, 853, 418, 64, 583, 511, 1_605, 62, 35, 93, 560, 177, 2_680, 217, 1_508, 1_521, 64, 583, 511, 519, 62, 20, 1_515, 764, 20, 149, 261, 5_625, 7_972, 20, 5_540, 567, 1_276, 93, 3_925, 1_675, 11, 15, 802, 7_972, 576, 217, 1_508, 11, 35, 93, 1_253, 2_441, 15, 289, 652, 31, 416, 321, 3_842, 115, 40, 911, 8, 476, 619, 4, 380, 142, 423, 335, 240, 35, 93, 264, 8, 11, 335, 569, 420, 163, 5, 2], [260, 548, 528, 423, 20, 451, 20, 2_681, 1_153, 3_434, 20, 5_540, 37, 567, 126, 1_253, 2_441, 3_376, 449, 210, 431, 1_563, 177, 767, 5_540, 11, 1_203, 472, 11, 2_953, 685, 285, 364, 706, 1_153, 20, 6_799, 20, 2_869, 20, 4_464, 126, 40, 2_429, 20, 1_040, 866, 2_664, 418, 20, 318, 20, 1_726, 186, 20, 265, 522, 35, 93, 2_191, 4_634, 20, 1_040, 12, 6_799, 15, 228, 2_356, 142, 31, 11, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [2_575, 2_666, 684, 1_582, 1_176, 12, 627, 149, 619, 20, 4_902, 563, 11, 20, 149, 261, 3_420, 2_356, 174, 142, 4_714, 131, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], "attention_mask": [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
										# fmt: on

										self.tokenizer_integration_test_util(
										    expected_encoding=_lowerCamelCase  ,    model_name="facebook/s2t-small-mustc-en-de-st"  ,    revision="a14f04cf0776c02f62a8cb800cf7909e15ea23ad"  ,    )





@require_sentencepiece
class 					__a     ( unittest.TestCase       ):



							'''simple docstring'''







							_lowerCamelCase					:			Any          =   """valhalla/s2t_mustc_multilinguial_medium"""

							_lowerCamelCase					:			Tuple          =   """C'est trop cool"""
							_lowerCamelCase					:			Optional[Any]          =   """Esto es genial"""
							@classmethod
							def   SCREAMING_SNAKE_CASE						(     cls       )			->							str:




										'''simple docstring'''
										__lowercase       =       SpeechaTextTokenizer.from_pretrained(cls.checkpoint_name       )
										return cls
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Any:




										'''simple docstring'''
										self.assertEqual(self.tokenizer.lang_code_to_id["pt"]  ,    4       )
										self.assertEqual(self.tokenizer.lang_code_to_id["ru"]  ,    6       )
										self.assertEqual(self.tokenizer.lang_code_to_id["it"]  ,    9       )
										self.assertEqual(self.tokenizer.lang_code_to_id["de"]  ,    11       )
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Union[str, Any]:




										'''simple docstring'''
										self.assertEqual(self.tokenizer.vocab_size  ,    10_000       )
							def   SCREAMING_SNAKE_CASE						(     self       )			->							str:




										'''simple docstring'''
										self.assertIn(_lowerCamelCase  ,    self.tokenizer.all_special_ids       )
										__lowercase       =       [ES_CODE, 4, 1_601, 47, 7_647, 2]
										__lowercase       =       self.tokenizer.decode(_lowerCamelCase  ,    skip_special_tokens=_lowerCamelCase       )
										__lowercase       =       self.tokenizer.decode(generated_ids[1:]  ,    skip_special_tokens=_lowerCamelCase       )
										self.assertEqual(_lowerCamelCase  ,    _lowerCamelCase       )
										self.assertNotIn(self.tokenizer.eos_token  ,    _lowerCamelCase       )
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Any:




										'''simple docstring'''
										__lowercase       =       "fr"
										__lowercase       =       self.tokenizer(self.french_text       ).input_ids
										self.assertEqual(encoded[0]  ,    _lowerCamelCase       )
										self.assertEqual(encoded[-1]  ,    self.tokenizer.eos_token_id       )





							def   SCREAMING_SNAKE_CASE						(     self       )			->							Optional[Any]:




										'''simple docstring'''
										__lowercase       =       "fr"
										self.assertListEqual(self.tokenizer.prefix_tokens  ,    [FR_CODE]       )

										__lowercase       =       "es"
										self.assertListEqual(self.tokenizer.prefix_tokens  ,    [ES_CODE]       )







"""simple docstring"""






import os

import jsonlines
import numpy as np
from tqdm import tqdm


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



def      lowerCAmelCase__			(		__magic_name__      )						->str:
			def choose_first(__magic_name__			,   __magic_name__=False      ):
						assert isinstance(__magic_name__			,   __magic_name__      )
						if len(__magic_name__      ) == 1:
									__lowercase       =       answer[0]
									return {k: [answer[k]] for k in answer} if is_long_answer else answer
						for a in answer:
									if is_long_answer:
												__lowercase       =       {k: [a[k]] for k in a}
									if len(a["start_token"]      ) > 0:
												break
						return a

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

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

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

			return answer



def      lowerCAmelCase__			(		__magic_name__			,   __magic_name__=False      )						->int:
			__lowercase       =       _get_single_answer(__magic_name__      )
			# 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
						__lowercase       =       example["document"]["tokens"]
						__lowercase       =       []
						for i in range(len(doc["token"]      )      ):
									if not doc["is_html"][i]:
												context.append(doc["token"][i]      )
						return {
						    "context": " ".join(__magic_name__      ),
						    "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

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

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

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

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

			return {
			    "context": " ".join(__magic_name__      ),
			    "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__			(		__magic_name__			,   __magic_name__			,   __magic_name__=2_0_4_8			,   __magic_name__=4_0_9_6			,   __magic_name__=True      )						->Optional[Any]:
			# overlap will be of doc_stride - q_len

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

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

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

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

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

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

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

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

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

			__lowercase       =       input_ids[:q_len]
			__lowercase       =       range(__magic_name__			,   len(__magic_name__      )			,   max_length - doc_stride      )

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

						answers_start_token.append(__magic_name__      )
						answers_end_token.append(__magic_name__      )
						if assertion:
									if new != old and new != [tokenizer.cls_token_id]:
												print("ISSUE in strided for ID:"			,   example["id"]      )
												print("New:"			,   tokenizer.decode(__magic_name__      )      )
												print("Old:"			,   tokenizer.decode(__magic_name__      )			,   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__			(		__magic_name__			,   __magic_name__			,   __magic_name__=2_0_4_8			,   __magic_name__=4_0_9_6			,   __magic_name__=False      )						->List[str]:
			__lowercase       =       get_strided_contexts_and_ans(
			    __magic_name__			,   __magic_name__			,   doc_stride=__magic_name__			,   max_length=__magic_name__			,   assertion=__magic_name__			,   )

			return example



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

					_lowercase							=				load_dataset('''natural_questions''')
					_lowercase							=				BigBirdTokenizer.from_pretrained('''google/bigbird-roberta-base''')

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

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

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





