





'''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_			:  Optional[Any]				):
			def choose_first(lowerCamelCase_			:  Dict		,       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										=							__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_			:  Union[str, Any]		,       lowerCamelCase_			:  Optional[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_			:  Tuple		,       lowerCamelCase_			:  List[Any]		,       lowerCamelCase_			:  Union[str, Any]=2_0_4_8		,       lowerCamelCase_			:  str=4_0_9_6		,       lowerCamelCase_			:  Union[str, Any]=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_			:  str		,       lowerCamelCase_			:  List[str]=2_0_4_8		,       lowerCamelCase_			:  Optional[int]=4_0_9_6		,       lowerCamelCase_			:  Tuple=False				):
			__lowercase										=							get_strided_contexts_and_ans(
			    lowerCamelCase_		,       lowerCamelCase_		,       doc_stride=lowerCamelCase_		,       max_length=lowerCamelCase_		,       assertion=lowerCamelCase_		,       )

			return example




def    _lowerCAmelCase						(				lowerCamelCase_			:  List[str]		,       lowerCamelCase_			:  Dict				):
			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_torch_available, is_vision_available


_SCREAMING_SNAKE_CASE					=    {'''configuration_van''': ['''VAN_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''VanConfig''']}


try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_SCREAMING_SNAKE_CASE					=    [
	    '''VAN_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''VanForImageClassification''',
	    '''VanModel''',
	    '''VanPreTrainedModel''',
	]

if TYPE_CHECKING:
	from .configuration_van import VAN_PRETRAINED_CONFIG_ARCHIVE_MAP, VanConfig

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_van import (
		    VAN_PRETRAINED_MODEL_ARCHIVE_LIST,
		    VanForImageClassification,
		    VanModel,
		    VanPreTrainedModel,
		)

else:
	import sys

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





