





'''simple docstring'''



import os

import jsonlines
import numpy as np
from tqdm import tqdm


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




def 				_lowerCAmelCase  (			lowerCamelCase_     :					int			):
			def choose_first(lowerCamelCase_     :					Any						,    lowerCamelCase_     :					str=False			):
						assert isinstance(_lowerCAmelCase						,    _lowerCAmelCase			)
						if len(_lowerCAmelCase			) == 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			  =							["""<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=_lowerCAmelCase			)
									__lowercase			  =							[]
						answer.update(_lowerCAmelCase			)

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

			return answer




def 				_lowerCAmelCase  (			lowerCamelCase_     :					List[Any]						,    lowerCamelCase_     :					int=False			):
			__lowercase			  =							_get_single_answer(_lowerCAmelCase			)
			# 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(_lowerCAmelCase			),
						    "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(_lowerCAmelCase			)			) if not is_html[i]]			)
						if new != old:
									print('''ID:'''						,    example['''id''']			)
									print('''New:'''						,    _lowerCAmelCase						,    end='''\n'''			)
									print('''Old:'''						,    _lowerCAmelCase						,    end='''\n\n'''			)

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

			__lowercase			  =							get_context_and_ans(_lowerCAmelCase						,    assertion=_lowerCAmelCase			)
			__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(_lowerCAmelCase						,    len(_lowerCAmelCase			)						,    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(_lowerCAmelCase			),
						        "end_token": [-1_0_0] * len(_lowerCAmelCase			),
						        "category": category,
						    },
						}

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

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

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

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

			return example




def 				_lowerCAmelCase  (			lowerCamelCase_     :					Any						,    lowerCamelCase_     :					Optional[Any]			):
			with jsonlines.open(_lowerCAmelCase						,    '''a'''			) as writer:
						for example in tqdm(_lowerCAmelCase						,    total=len(_lowerCAmelCase			)						,    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

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

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

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

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












'''simple docstring'''



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


_SCREAMING_SNAKE_CASE   =      {
    '''configuration_mvp''': ['''MVP_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''MvpConfig''', '''MvpOnnxConfig'''],
    '''tokenization_mvp''': ['''MvpTokenizer'''],
}

try:
	if not is_tokenizers_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_SCREAMING_SNAKE_CASE   =      ['''MvpTokenizerFast''']

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_SCREAMING_SNAKE_CASE   =      [
	    '''MVP_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''MvpForCausalLM''',
	    '''MvpForConditionalGeneration''',
	    '''MvpForQuestionAnswering''',
	    '''MvpForSequenceClassification''',
	    '''MvpModel''',
	    '''MvpPreTrainedModel''',
	]

if TYPE_CHECKING:
	from .configuration_mvp import MVP_PRETRAINED_CONFIG_ARCHIVE_MAP, MvpConfig, MvpOnnxConfig
	from .tokenization_mvp import MvpTokenizer

	try:
		if not is_tokenizers_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_mvp_fast import MvpTokenizerFast

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_mvp import (
		    MVP_PRETRAINED_MODEL_ARCHIVE_LIST,
		    MvpForCausalLM,
		    MvpForConditionalGeneration,
		    MvpForQuestionAnswering,
		    MvpForSequenceClassification,
		    MvpModel,
		    MvpPreTrainedModel,
		)

else:
	import sys

	_SCREAMING_SNAKE_CASE   =      _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





