
import os

import jsonlines
import numpy as np
from tqdm import tqdm


lowerCamelCase_  :	Tuple 				=					20_48
lowerCamelCase_  :	Union[str, Any] 				=					40_96
lowerCamelCase_  :	Dict 				=					42
lowerCamelCase_  :	List[Any] 				=					os.environ.pop("""PROCESS_TRAIN""", """false""")
lowerCamelCase_  :	Optional[Any] 				=					{'''null''': 0, '''short''': 1, '''long''': 2, '''yes''': 3, '''no''': 4}







def     A__   (  lowerCamelCase  )					->		Dict:



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

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

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

				UpperCamelCase_:							List[Any]		=  ['''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     A__   (  lowerCamelCase       ,      lowerCamelCase=False  )					->		Dict:
				UpperCamelCase_:							Optional[int]		=  _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
								UpperCamelCase_:							Optional[int]		=  example['''document''']['''tokens''']
								UpperCamelCase_:							str		=  []
								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": -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

				UpperCamelCase_:							Optional[Any]		=  ['''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

				UpperCamelCase_:							Dict		=  example['''document''']['''tokens''']
				UpperCamelCase_:							Union[str, Any]		=  answer['''start_token''']
				UpperCamelCase_:							Optional[Any]		=  answer['''end_token''']

				UpperCamelCase_:							List[str]		=  []
				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
				UpperCamelCase_:							List[str]		=  ''' '''.join(context[start_token:end_token]  )

				# checking above code
				if assertion:
								UpperCamelCase_:							List[str]		=  doc['''is_html'''][answer['''start_token'''] : answer['''end_token''']]
								UpperCamelCase_:							str		=  doc['''token'''][answer['''start_token'''] : answer['''end_token''']]
								UpperCamelCase_:							Union[str, 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     A__   (  lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase=20_48       ,      lowerCamelCase=40_96       ,      lowerCamelCase=True  )					->		Optional[int]:
				UpperCamelCase_:							Dict		=  get_context_and_ans(a__       ,      assertion=a__  )
				UpperCamelCase_:							Union[str, 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"],
								    },
								}

				UpperCamelCase_:							Optional[int]		=  tokenizer(example["""question"""]["""text"""]       ,      out["""context"""]  ).input_ids
				UpperCamelCase_:							Tuple		=  input_ids.index(tokenizer.sep_token_id  ) + 1

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

				UpperCamelCase_:							Optional[Any]		=  out['''context'''].split()
				UpperCamelCase_:							str		=  splitted_context[answer['''end_token''']]
				UpperCamelCase_:							Union[str, Any]		=  len(
				    tokenizer(
				        """ """.join(splitted_context[: answer["""start_token"""]]  )       ,      add_special_tokens=a__       ,      ).input_ids  )
				UpperCamelCase_:							List[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
				UpperCamelCase_:							Dict		=  len(tokenizer(a__       ,      add_special_tokens=a__  ).input_ids  )
				if num_sub_tokens > 1:
								answer["end_token"] += num_sub_tokens - 1

				UpperCamelCase_:							List[str]		=  input_ids[answer['''start_token'''] : answer['''end_token'''] + 1]  # right & left are inclusive
				UpperCamelCase_:							int		=  answer['''start_token''']
				UpperCamelCase_:							List[Any]		=  answer['''end_token''']

				if assertion:
								UpperCamelCase_:							List[str]		=  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"],
								    },
								}

				UpperCamelCase_:							str		=  input_ids[:q_len]
				UpperCamelCase_:							Optional[Any]		=  range(a__       ,      len(a__  )       ,      max_length - doc_stride  )

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

				return example







def     A__   (  lowerCamelCase       ,      lowerCamelCase  )					->		Any:



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

		lowerCamelCase_  :	Union[str, Any] 				=					load_dataset("""natural_questions""")
		lowerCamelCase_  :	Any 				=					BigBirdTokenizer.from_pretrained("""google/bigbird-roberta-base""")

		lowerCamelCase_  :	Dict 				=					data['''train''' if PROCESS_TRAIN == '''true''' else '''validation''']

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

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








from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase_  :	Optional[int] 				=					logging.get_logger(__name__)

lowerCamelCase_  :	Optional[int] 				=					{"""ctrl""": """https://huggingface.co/ctrl/resolve/main/config.json"""}

class     _UpperCamelCase			(   _A					):






			'''simple docstring'''

			__UpperCamelCase :  int    =	"""ctrl"""
			__UpperCamelCase :  Dict    =	["""past_key_values"""]
			__UpperCamelCase :  List[str]    =	{
			    """max_position_embeddings""": """n_positions""",
			    """hidden_size""": """n_embd""",
			    """num_attention_heads""": """n_head""",
			    """num_hidden_layers""": """n_layer""",
			}

			def __init__(    self	:    Dict  ,		snake_case_	:    Any=24_6534  ,		snake_case_	:    Dict=256  ,		snake_case_	:    str=1280  ,		snake_case_	:    Optional[int]=8192  ,		snake_case_	:    Union[str, Any]=48  ,		snake_case_	:    Any=16  ,		snake_case_	:    Optional[int]=0.1  ,		snake_case_	:    Any=0.1  ,		snake_case_	:    Any=1e-6  ,		snake_case_	:    Optional[Any]=0.02  ,		snake_case_	:    Optional[int]=True  ,		**snake_case_	:    Union[str, Any]  ,		):
							UpperCamelCase_:							Union[str, Any]		=  vocab_size
							UpperCamelCase_:							Union[str, Any]		=  n_positions
							UpperCamelCase_:							Optional[int]		=  n_embd
							UpperCamelCase_:							int		=  n_layer
							UpperCamelCase_:							str		=  n_head
							UpperCamelCase_:							Optional[int]		=  dff
							UpperCamelCase_:							Optional[Any]		=  resid_pdrop
							UpperCamelCase_:							Union[str, Any]		=  embd_pdrop
							UpperCamelCase_:							List[str]		=  layer_norm_epsilon
							UpperCamelCase_:							Optional[Any]		=  initializer_range

							UpperCamelCase_:							Optional[Any]		=  use_cache

							super().__init__(**snake_case_ )






