






"""simple docstring"""
from abc import ABC, abstractmethod
from argparse import ArgumentParser



class        UpperCamelCase__ (				a_):







    """simple docstring"""







    @staticmethod
    @abstractmethod
    def 				a__			(		UpperCamelCase_ :	ArgumentParser ):





           '''simple docstring'''

           raise NotImplementedError()







    @abstractmethod
    def 				a__			(		self :	List[str] ):





           '''simple docstring'''

           raise NotImplementedError()







"""simple docstring"""
import os

import jsonlines
import numpy as np
from tqdm import tqdm


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


def 							A						(				__snake_case: Any  )    ->		Optional[int]:

       """simple docstring"""





       def choose_first(__snake_case: Union[str, Any] ,     __snake_case: List[str]=False  ):
              assert isinstance(__snake_case ,     __snake_case  )
              if len(__snake_case  ) == 1:
                     __magic_name__      				=    answer[0]
                     return {k: [answer[k]] for k in answer} if is_long_answer else answer
              for a in answer:
                     if is_long_answer:
                            __magic_name__      				=    {k: [a[k]] for k in a}
                     if len(a['start_token']  ) > 0:
                            break
              return a

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

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

       __magic_name__      				=    ['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 							A						(				__snake_case: Any ,     __snake_case: str=False  )    ->		Optional[Any]:

       """simple docstring"""





       __magic_name__      				=    _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
              __magic_name__      				=    example['document']['tokens']
              __magic_name__      				=    []
              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_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

       __magic_name__      				=    ['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

       __magic_name__      				=    example['document']['tokens']
       __magic_name__      				=    answer['start_token']
       __magic_name__      				=    answer['end_token']

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

       # checking above code
       if assertion:
              __magic_name__      				=    doc['is_html'][answer['start_token'] : answer['end_token']]
              __magic_name__      				=    doc['token'][answer['start_token'] : answer['end_token']]
              __magic_name__      				=    ' '.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 							A						(				__snake_case: List[Any] ,     __snake_case: int ,     __snake_case: Tuple=2_0_4_8 ,     __snake_case: List[str]=4_0_9_6 ,     __snake_case: Optional[int]=True  )    ->		Any:

       """simple docstring"""





       __magic_name__      				=    get_context_and_ans(__snake_case ,     assertion=__snake_case  )
       __magic_name__      				=    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"],
                  },
              }

       __magic_name__      				=    tokenizer(example['question']['text'] ,     out['context']  ).input_ids
       __magic_name__      				=    input_ids.index(tokenizer.sep_token_id  ) + 1

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

       __magic_name__      				=    out['context'].split()
       __magic_name__      				=    splitted_context[answer['end_token']]
       __magic_name__      				=    len(
           tokenizer(
               ' '.join(splitted_context[: answer['start_token']]  ) ,     add_special_tokens=__snake_case ,     ).input_ids  )
       __magic_name__      				=    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
       __magic_name__      				=    len(tokenizer(__snake_case ,     add_special_tokens=__snake_case  ).input_ids  )
       if num_sub_tokens > 1:
              answer["end_token"] += num_sub_tokens - 1

       __magic_name__      				=    input_ids[answer['start_token'] : answer['end_token'] + 1]  # right & left are inclusive
       __magic_name__      				=    answer['start_token']
       __magic_name__      				=    answer['end_token']

       if assertion:
              __magic_name__      				=    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"],
                  },
              }

       __magic_name__      				=    input_ids[:q_len]
       __magic_name__      				=    range(__snake_case ,     len(__snake_case  ) ,     max_length - doc_stride  )

       __magic_name__      				=    []
       __magic_name__      				=    []
       __magic_name__      				=    []
       __magic_name__      				=    []  # null, yes, no, long, short
       for i in doc_start_indices:
              __magic_name__      				=    i + max_length - q_len
              __magic_name__      				=    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:
                     __magic_name__      				=    start_token - i + q_len
                     __magic_name__      				=    end_token - i + q_len
                     answers_category.append(answer['category'][0]  )  # ["short"] -> "short"
              else:
                     __magic_name__      				=    -1_0_0
                     __magic_name__      				=    -1_0_0
                     answers_category.append('null'  )
              __magic_name__      				=    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 							A						(				__snake_case: Tuple ,     __snake_case: Dict ,     __snake_case: Any=2_0_4_8 ,     __snake_case: List[Any]=4_0_9_6 ,     __snake_case: Any=False  )    ->		Tuple:

       """simple docstring"""





       __magic_name__      				=    get_strided_contexts_and_ans(
           __snake_case ,     __snake_case ,     doc_stride=__snake_case ,     max_length=__snake_case ,     assertion=__snake_case ,     )

       return example


def 							A						(				__snake_case: Optional[int] ,     __snake_case: Union[str, Any]  )    ->		Dict:

       """simple docstring"""





       with jsonlines.open(__snake_case ,     'a'  ) as writer:
              for example in tqdm(__snake_case ,     total=len(__snake_case  ) ,     desc='Saving samples ... '  ):
                     __magic_name__      				=    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 : Any     		=		load_dataset("""natural_questions""")
 snake_case : List[Any]     		=		BigBirdTokenizer.from_pretrained("""google/bigbird-roberta-base""")

 snake_case : Tuple     		=		data["""train""" if PROCESS_TRAIN == """true""" else """validation"""]

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

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