






import os

import jsonlines
import numpy as np
from tqdm import tqdm


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







def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							str       ):





	def choose_first(snake_case_ :							str					,  snake_case_ :							List[Any]=False       ):
		assert isinstance(a__					,  a__       )
		if len(a__       ) == 1:
			snake_case__			:				Tuple       =  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__			:				int       =  {k: [a[k]] for k in a}
			if len(a["start_token"]       ) > 0:
				break
		return a

	snake_case__			:				List[Any]       =  {"id": example["id"]}
	snake_case__			:				Optional[int]       =  example["annotations"]
	snake_case__			:				Union[str, Any]       =  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__			:				Tuple       =  []
		snake_case__			:				Any       =  []
		snake_case__			:				Union[str, Any]       =  ["<cls>"]
	else:
		snake_case__			:				Optional[int]       =  ["short"]
		snake_case__			:				Union[str, 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__			:				Dict       =  choose_first(annotation["long_answer"]					,  is_long_answer=a__       )
			snake_case__			:				Dict       =  []
		answer.update(a__       )

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

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

	return answer







def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							Dict					,  snake_case_ :							Union[str, Any]=False       ):
	snake_case__			:				Any       =  _get_single_answer(a__       )
	# 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__			:				Tuple       =  []
		for i in range(len(doc["token"]       )       ):
			if not doc["is_html"][i]:
				context.append(doc["token"][i]       )
		return {
		    "context": " ".join(a__       ),
		    "answer": {
		        "start_token": -100,  # ignore index in cross-entropy
		        "end_token": -100,  # 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__			:				Optional[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

	snake_case__			:				List[str]       =  example["document"]["tokens"]
	snake_case__			:				Any       =  answer["start_token"]
	snake_case__			:				str       =  answer["end_token"]

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

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

	return {
	    "context": " ".join(a__       ),
	    "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 					SCREAMING_SNAKE_CASE    (			snake_case_ :							Optional[int]					,  snake_case_ :							Optional[Any]					,  snake_case_ :							Union[str, Any]=2048					,  snake_case_ :							Optional[int]=4096					,  snake_case_ :							Optional[Any]=True       ):
	snake_case__			:				List[Any]       =  get_context_and_ans(a__					,  assertion=a__       )
	snake_case__			:				Optional[Any]       =  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__			:				Tuple       =  tokenizer(example["question"]["text"]					,  out["context"]       ).input_ids
	snake_case__			:				List[Any]       =  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__			:				Dict       =  []
		snake_case__			:				Any       =  []
		snake_case__			:				Tuple       =  input_ids[:q_len]
		snake_case__			:				Tuple       =  range(a__					,  len(a__       )					,  max_length - doc_stride       )
		for i in doc_start_indices:
			snake_case__			:				Tuple       =  i + max_length - q_len
			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": [-100] * len(a__       ),
		        "end_token": [-100] * len(a__       ),
		        "category": category,
		    },
		}

	snake_case__			:				int       =  out["context"].split()
	snake_case__			:				Optional[int]       =  splitted_context[answer["end_token"]]
	snake_case__			:				Tuple       =  len(
	    tokenizer(
	        " ".join(splitted_context[: answer["start_token"]]       )					,  add_special_tokens=a__					,  ).input_ids       )
	snake_case__			:				str       =  len(
	    tokenizer(" ".join(splitted_context[: answer["end_token"]]       )					,  add_special_tokens=a__       ).input_ids       )

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

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

	snake_case__			:				Tuple       =  input_ids[answer["start_token"] : answer["end_token"] + 1]  # right & left are inclusive
	snake_case__			:				List[Any]       =  answer["start_token"]
	snake_case__			:				Tuple       =  answer["end_token"]

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

	if len(a__       ) <= 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__			:				Tuple       =  input_ids[:q_len]
	snake_case__			:				List[str]       =  range(a__					,  len(a__       )					,  max_length - doc_stride       )

	snake_case__			:				Optional[Any]       =  []
	snake_case__			:				List[Any]       =  []
	snake_case__			:				int       =  []
	snake_case__			:				Dict       =  []  # null, yes, no, long, short
	for i in doc_start_indices:
		snake_case__			:				Union[str, Any]       =  i + max_length - q_len
		snake_case__			:				Dict       =  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__			:				List[Any]       =  start_token - i + q_len
			snake_case__			:				List[str]       =  end_token - i + q_len
			answers_category.append(answer["category"][0]       )  # ["short"] -> "short"
		else:
			snake_case__			:				Optional[int]       =  -100
			snake_case__			:				Any       =  -100
			answers_category.append("null"       )
		snake_case__			:				List[Any]       =  inputs[-1][start_token : end_token + 1]

		answers_start_token.append(a__       )
		answers_end_token.append(a__       )
		if assertion:
			if new != old and new != [tokenizer.cls_token_id]:
				print("ISSUE in strided for ID:"					,  example["id"]       )
				print("New:"					,  tokenizer.decode(a__       )       )
				print("Old:"					,  tokenizer.decode(a__       )					,  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 					SCREAMING_SNAKE_CASE    (			snake_case_ :							Optional[Any]					,  snake_case_ :							Tuple					,  snake_case_ :							Optional[Any]=2048					,  snake_case_ :							Tuple=4096					,  snake_case_ :							Optional[Any]=False       ):
	snake_case__			:				List[str]       =  get_strided_contexts_and_ans(
	    a__					,  a__					,  doc_stride=a__					,  max_length=a__					,  assertion=a__					,  )

	return example







def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							Optional[int]					,  snake_case_ :							Optional[int]       ):





	with jsonlines.open(a__					,  "a"       ) as writer:
		for example in tqdm(a__					,  total=len(a__       )					,  desc="Saving samples ... "       ):
			snake_case__			:				Dict       =  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 :      str										=				load_dataset("""natural_questions""")
				__lowerCamelCase :      Any										=				BigBirdTokenizer.from_pretrained("""google/bigbird-roberta-base""")

				__lowerCamelCase :      int										=				data["""train""" if PROCESS_TRAIN == """true""" else """validation"""]

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

				np.random.seed(SEED)
				__lowerCamelCase :      List[Any]										=				"""nq-training.jsonl""" if PROCESS_TRAIN == """true""" else """nq-validation.jsonl"""
				save_to_disk(data, file_name=cache_file_name)











from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__lowerCamelCase :      List[str]										=				logging.get_logger(__name__)

__lowerCamelCase :      List[Any]										=				{
    """facebook/xmod-base""": """https://huggingface.co/facebook/xmod-base/resolve/main/config.json""",
    """facebook/xmod-large-prenorm""": """https://huggingface.co/facebook/xmod-large-prenorm/resolve/main/config.json""",
    """facebook/xmod-base-13-125k""": """https://huggingface.co/facebook/xmod-base-13-125k/resolve/main/config.json""",
    """facebook/xmod-base-30-125k""": """https://huggingface.co/facebook/xmod-base-30-125k/resolve/main/config.json""",
    """facebook/xmod-base-30-195k""": """https://huggingface.co/facebook/xmod-base-30-195k/resolve/main/config.json""",
    """facebook/xmod-base-60-125k""": """https://huggingface.co/facebook/xmod-base-60-125k/resolve/main/config.json""",
    """facebook/xmod-base-60-265k""": """https://huggingface.co/facebook/xmod-base-60-265k/resolve/main/config.json""",
    """facebook/xmod-base-75-125k""": """https://huggingface.co/facebook/xmod-base-75-125k/resolve/main/config.json""",
    """facebook/xmod-base-75-269k""": """https://huggingface.co/facebook/xmod-base-75-269k/resolve/main/config.json""",
}



class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					a_     =		"xmod"



					def __init__( self					:  List[str]							,		__A					:  List[str]=3_0_5_2_2							,		__A					:  Tuple=7_6_8							,		__A					:  str=1_2							,		__A					:  List[Any]=1_2							,		__A					:  List[str]=3_0_7_2							,		__A					:  List[str]="gelu"							,		__A					:  List[Any]=0.1							,		__A					:  Tuple=0.1							,		__A					:  str=5_1_2							,		__A					:  Union[str, Any]=2							,		__A					:  List[Any]=0.0_2							,		__A					:  List[str]=1e-1_2							,		__A					:  Tuple=1							,		__A					:  List[Any]=0							,		__A					:  Optional[Any]=2							,		__A					:  Optional[int]="absolute"							,		__A					:  Optional[int]=True							,		__A					:  Dict=None							,		__A					:  Optional[int]=False							,		__A					:  Dict=2							,		__A					:  List[str]=False							,		__A					:  Dict=True							,		__A					:  Union[str, Any]=True							,		__A					:  Tuple=("en_XX",)							,		__A					:  Optional[Any]=None							,		**__A					:  Tuple							,		):
						super().__init__(pad_token_id=__A							,		bos_token_id=__A							,		eos_token_id=__A							,		**__A	)

						snake_case__			:				Tuple       =  vocab_size
						snake_case__			:				Union[str, Any]       =  hidden_size
						snake_case__			:				int       =  num_hidden_layers
						snake_case__			:				str       =  num_attention_heads
						snake_case__			:				Tuple       =  hidden_act
						snake_case__			:				Optional[Any]       =  intermediate_size
						snake_case__			:				List[str]       =  hidden_dropout_prob
						snake_case__			:				Optional[int]       =  attention_probs_dropout_prob
						snake_case__			:				Union[str, Any]       =  max_position_embeddings
						snake_case__			:				str       =  type_vocab_size
						snake_case__			:				List[str]       =  initializer_range
						snake_case__			:				Dict       =  layer_norm_eps
						snake_case__			:				str       =  position_embedding_type
						snake_case__			:				List[str]       =  use_cache
						snake_case__			:				Tuple       =  classifier_dropout
						snake_case__			:				Any       =  pre_norm
						snake_case__			:				List[str]       =  adapter_reduction_factor
						snake_case__			:				List[Any]       =  adapter_layer_norm
						snake_case__			:				str       =  adapter_reuse_layer_norm
						snake_case__			:				Union[str, Any]       =  ln_before_adapter
						snake_case__			:				Tuple       =  list(__A	)
						snake_case__			:				int       =  default_language



class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					@property
					def 	_lowercase					( self					:  int	):
						if self.task == "multiple-choice":
							snake_case__			:				Any       =  {0: "batch", 1: "choice", 2: "sequence"}
						else:
							snake_case__			:				Any       =  {0: "batch", 1: "sequence"}
						return OrderedDict(
						    [
						        ("input_ids", dynamic_axis),
						        ("attention_mask", dynamic_axis),
						    ]	)



