
import os

import jsonlines
import numpy as np
from tqdm import tqdm


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

def 		UpperCamelCase	(      lowercase_   )		->	List[Any]:


  '''simple docstring'''






  def choose_first(lowercase_ , lowercase_=False   ):
    assert isinstance(_snake_case , _snake_case   )
    if len(_snake_case   ) == 1:
      lowercase__  :   List[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:
        lowercase__  :   Union[str, Any]    		=      {k: [a[k]] for k in a}
      if len(a["""start_token"""]   ) > 0:
        break
    return a

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

  # disregard some samples
  if len(answer["""start_token"""]   ) > 1 or answer["start_token"] == answer["end_token"]:
    lowercase__  :   Optional[int]    		=      True
  else:
    lowercase__  :   str    		=      False

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

  return answer

def 		UpperCamelCase	(      lowercase_ , lowercase_=False   )		->	List[Any]:


  '''simple docstring'''






  lowercase__  :   Tuple    		=      _get_single_answer(_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
    lowercase__  :   Optional[int]    		=      example["""document"""]["""tokens"""]
    lowercase__  :   str    		=      []
    for i in range(len(doc["""token"""]   )   ):
      if not doc["is_html"][i]:
        context.append(doc["""token"""][i]   )
    return {
        "context": " ".join(_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

  lowercase__  :   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

  lowercase__  :   str    		=      example["""document"""]["""tokens"""]
  lowercase__  :   Union[str, Any]    		=      answer["""start_token"""]
  lowercase__  :   Optional[int]    		=      answer["""end_token"""]

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

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

  return {
      "context": " ".join(_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 		UpperCamelCase	(      lowercase_ , lowercase_ , lowercase_=20_48 , lowercase_=40_96 , lowercase_=True   )		->	Dict:


  '''simple docstring'''






  lowercase__  :   Tuple    		=      get_context_and_ans(_snake_case , assertion=_snake_case   )
  lowercase__  :   List[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"],
        },
    }

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

  lowercase__  :   Dict    		=      out["""context"""].split()
  lowercase__  :   List[str]    		=      splitted_context[answer["""end_token"""]]
  lowercase__  :   Optional[int]    		=      len(
      tokenizer(
          """ """.join(splitted_context[: answer["""start_token"""]]   ) , add_special_tokens=_snake_case , ).input_ids   )
  lowercase__  :   Any    		=      len(
      tokenizer(""" """.join(splitted_context[: answer["""end_token"""]]   ) , add_special_tokens=_snake_case   ).input_ids   )

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

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

  lowercase__  :   Optional[int]    		=      input_ids[answer["""start_token"""] : answer["""end_token"""] + 1]  # right & left are inclusive
  lowercase__  :   Any    		=      answer["""start_token"""]
  lowercase__  :   List[Any]    		=      answer["""end_token"""]

  if assertion:
    lowercase__  :   List[Any]    		=      tokenizer.decode(_snake_case   )
    if answer["span"] != new:
      print("""ISSUE IN TOKENIZATION"""   )
      print("""OLD:""" , answer["""span"""]   )
      print("""NEW:""" , _snake_case , end="""\n\n"""   )

  if len(_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"],
        },
    }

  lowercase__  :   List[Any]    		=      input_ids[:q_len]
  lowercase__  :   str    		=      range(_snake_case , len(_snake_case   ) , max_length - doc_stride   )

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

    answers_start_token.append(_snake_case   )
    answers_end_token.append(_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(_snake_case   )   )
        print("""Old:""" , tokenizer.decode(_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 		UpperCamelCase	(      lowercase_ , lowercase_ , lowercase_=20_48 , lowercase_=40_96 , lowercase_=False   )		->	int:


  '''simple docstring'''






  lowercase__  :   List[Any]    		=      get_strided_contexts_and_ans(
      _snake_case , _snake_case , doc_stride=_snake_case , max_length=_snake_case , assertion=_snake_case , )

  return example

def 		UpperCamelCase	(      lowercase_ , lowercase_   )		->	Any:


  '''simple docstring'''






  with jsonlines.open(_snake_case , """a"""   ) as writer:
    for example in tqdm(_snake_case , total=len(_snake_case   ) , desc="""Saving samples ... """   ):
      lowercase__  :   Tuple    		=      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__    :		List[str]          =   load_dataset("""natural_questions""")
     lowerCamelCase__    :		List[str]          =   BigBirdTokenizer.from_pretrained("""google/bigbird-roberta-base""")

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

     lowerCamelCase__    :		List[str]          =   {
         """tokenizer""": tokenizer,
         """doc_stride""": DOC_STRIDE,
         """max_length""": MAX_LENGTH,
         """assertion""": False,
     }
     lowerCamelCase__    :		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__    :		Optional[int]          =   """nq-training.jsonl""" if PROCESS_TRAIN == """true""" else """nq-validation.jsonl"""
     save_to_disk(data, file_name=cache_file_name)








def 		UpperCamelCase	(      lowercase_ = 10_00   )		->	int:


  '''simple docstring'''






  return sum(2 * a * ((a - 1) // 2) for a in range(3 , n + 1   )   )


if __name__ == "__main__":
     print(solution())






