



'''simple docstring'''

import os

import jsonlines
import numpy as np
from tqdm import tqdm


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






def       __lowercase  (_SCREAMING_SNAKE_CASE      :List[Any]       ):
				def choose_first(_SCREAMING_SNAKE_CASE      :Any    ,	_SCREAMING_SNAKE_CASE      :Optional[Any]=False       ):
								assert isinstance(SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_       )
								if len(SCREAMING_SNAKE_CASE_       ) == 1:
												SCREAMING_SNAKE_CASE  :						Optional[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[int]      					= {k: [a[k]] for k in a}
												if len(a['''start_token''']       ) > 0:
																break
								return a

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

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

				SCREAMING_SNAKE_CASE  :						Tuple      					= ['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       __lowercase  (_SCREAMING_SNAKE_CASE      :List[Any]    ,	_SCREAMING_SNAKE_CASE      :str=False       ):
				SCREAMING_SNAKE_CASE  :						int      					= _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  :						List[Any]      					= example['document']['tokens']
								SCREAMING_SNAKE_CASE  :						Tuple      					= []
								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": -1_00,  # ignore index in cross-entropy
								        "end_token": -1_00,  # 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  :						int      					= ['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  :						Tuple      					= example['document']['tokens']
				SCREAMING_SNAKE_CASE  :						int      					= answer['start_token']
				SCREAMING_SNAKE_CASE  :						List[Any]      					= answer['end_token']

				SCREAMING_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
				SCREAMING_SNAKE_CASE  :						Tuple      					= ' '.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  :						List[Any]      					= doc['token'][answer['start_token'] : answer['end_token']]
								SCREAMING_SNAKE_CASE  :						Any      					= ' '.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       __lowercase  (_SCREAMING_SNAKE_CASE      :str    ,	_SCREAMING_SNAKE_CASE      :List[Any]    ,	_SCREAMING_SNAKE_CASE      :int=20_48    ,	_SCREAMING_SNAKE_CASE      :List[str]=40_96    ,	_SCREAMING_SNAKE_CASE      :List[str]=True       ):
				# overlap will be of doc_stride - q_len

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

				SCREAMING_SNAKE_CASE  :						List[str]      					= out['context'].split()
				SCREAMING_SNAKE_CASE  :						int      					= splitted_context[answer['end_token']]
				SCREAMING_SNAKE_CASE  :						List[str]      					= len(
				    tokenizer(
				        ''' '''.join(splitted_context[: answer['''start_token''']]       )    ,	add_special_tokens=SCREAMING_SNAKE_CASE_    ,	).input_ids       )
				SCREAMING_SNAKE_CASE  :						int      					= 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  :						List[Any]      					= 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  :						Optional[Any]      					= input_ids[answer['start_token'] : answer['end_token'] + 1]  # right & left are inclusive
				SCREAMING_SNAKE_CASE  :						Optional[int]      					= answer['start_token']
				SCREAMING_SNAKE_CASE  :						Dict      					= answer['end_token']

				if assertion:
								SCREAMING_SNAKE_CASE  :						Union[str, 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  :						Optional[Any]      					= input_ids[:q_len]
				SCREAMING_SNAKE_CASE  :						Union[str, Any]      					= range(SCREAMING_SNAKE_CASE_    ,	len(SCREAMING_SNAKE_CASE_       )    ,	max_length - doc_stride       )

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

					snake_case_			=     load_dataset("""natural_questions""")
					snake_case_			=     BigBirdTokenizer.from_pretrained("""google/bigbird-roberta-base""")

					snake_case_			=     data["""train""" if PROCESS_TRAIN == """true""" else """validation"""]

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

					np.random.seed(SEED)
					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 ...configuration_utils import PretrainedConfig
from ...utils import logging


snake_case_			=     logging.get_logger(__name__)

snake_case_			=     {
    """microsoft/markuplm-base""": """https://huggingface.co/microsoft/markuplm-base/resolve/main/config.json""",
    """microsoft/markuplm-large""": """https://huggingface.co/microsoft/markuplm-large/resolve/main/config.json""",
}


class 	a__     (				_lowercase	):
	__magic_name__		:    List[Any]	    =     "markuplm"


	def __init__(self : List[Any],					__UpperCAmelCase : int=30522,					__UpperCAmelCase : Tuple=768,					__UpperCAmelCase : str=12,					__UpperCAmelCase : List[Any]=12,					__UpperCAmelCase : Tuple=3072,					__UpperCAmelCase : Union[str, Any]="gelu",					__UpperCAmelCase : Union[str, Any]=0.1,					__UpperCAmelCase : str=0.1,					__UpperCAmelCase : Union[str, Any]=512,					__UpperCAmelCase : Dict=2,					__UpperCAmelCase : int=0.02,					__UpperCAmelCase : Optional[Any]=1e-12,					__UpperCAmelCase : Union[str, Any]=0,					__UpperCAmelCase : Optional[Any]=0,					__UpperCAmelCase : Union[str, Any]=2,					__UpperCAmelCase : int=256,					__UpperCAmelCase : Tuple=1024,					__UpperCAmelCase : Optional[Any]=216,					__UpperCAmelCase : str=1001,					__UpperCAmelCase : Optional[Any]=32,					__UpperCAmelCase : Dict=50,					__UpperCAmelCase : Union[str, Any]="absolute",					__UpperCAmelCase : Tuple=True,					__UpperCAmelCase : List[str]=None,					**__UpperCAmelCase : Union[str, Any],					)     ->			Union[str, Any]:

					"""simple docstring"""

					super().__init__(
					    pad_token_id=__UpperCAmelCase,					bos_token_id=__UpperCAmelCase,					eos_token_id=__UpperCAmelCase,					**__UpperCAmelCase,					)
					SCREAMING_SNAKE_CASE  :						List[Any]      					= vocab_size
					SCREAMING_SNAKE_CASE  :						Dict      					= hidden_size
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= num_hidden_layers
					SCREAMING_SNAKE_CASE  :						Dict      					= num_attention_heads
					SCREAMING_SNAKE_CASE  :						Any      					= hidden_act
					SCREAMING_SNAKE_CASE  :						Tuple      					= intermediate_size
					SCREAMING_SNAKE_CASE  :						Tuple      					= hidden_dropout_prob
					SCREAMING_SNAKE_CASE  :						List[str]      					= attention_probs_dropout_prob
					SCREAMING_SNAKE_CASE  :						List[str]      					= max_position_embeddings
					SCREAMING_SNAKE_CASE  :						Dict      					= type_vocab_size
					SCREAMING_SNAKE_CASE  :						Tuple      					= initializer_range
					SCREAMING_SNAKE_CASE  :						Optional[Any]      					= layer_norm_eps
					SCREAMING_SNAKE_CASE  :						int      					= position_embedding_type
					SCREAMING_SNAKE_CASE  :						List[Any]      					= use_cache
					SCREAMING_SNAKE_CASE  :						int      					= classifier_dropout
					# additional properties
					SCREAMING_SNAKE_CASE  :						List[Any]      					= max_depth
					SCREAMING_SNAKE_CASE  :						Any      					= max_xpath_tag_unit_embeddings
					SCREAMING_SNAKE_CASE  :						int      					= max_xpath_subs_unit_embeddings
					SCREAMING_SNAKE_CASE  :						Dict      					= tag_pad_id
					SCREAMING_SNAKE_CASE  :						Optional[int]      					= subs_pad_id
					SCREAMING_SNAKE_CASE  :						int      					= xpath_unit_hidden_size




