



"""simple docstring"""






import os

import jsonlines
import numpy as np
from tqdm import tqdm


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






def 	snake_case_  (				A_ :		Dict						):


       '''simple docstring'''





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

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

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

       _lowerCamelCase						:    Union[str, 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 	snake_case_  (				A_ :		List[Any],    A_ :		Dict=False						):


       '''simple docstring'''





       _lowerCamelCase						:    Union[str, 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
              _lowerCamelCase						:    Optional[int]         =			example['''document''']['''tokens''']
              _lowerCamelCase						:    Optional[Any]         =			[]
              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

       _lowerCamelCase						:    List[str]         =			['''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

       _lowerCamelCase						:    Optional[int]         =			example['''document''']['''tokens''']
       _lowerCamelCase						:    Dict         =			answer['''start_token''']
       _lowerCamelCase						:    Optional[Any]         =			answer['''end_token''']

       _lowerCamelCase						:    Union[str, 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
       _lowerCamelCase						:    Tuple         =			''' '''.join(context[start_token:end_token]						)

       # checking above code
       if assertion:
              _lowerCamelCase						:    Dict         =			doc['''is_html'''][answer['''start_token'''] : answer['''end_token''']]
              _lowerCamelCase						:    Tuple         =			doc['''token'''][answer['''start_token'''] : answer['''end_token''']]
              _lowerCamelCase						:    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 	snake_case_  (				A_ :		int,    A_ :		List[Any],    A_ :		Dict=20_48,    A_ :		Dict=40_96,    A_ :		Optional[int]=True						):


       '''simple docstring'''





       _lowerCamelCase						:    Union[str, Any]         =			get_context_and_ans(A_,    assertion=A_						)
       _lowerCamelCase						:    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"],
                  },
              }

       _lowerCamelCase						:    Any         =			tokenizer(example['''question''']['''text'''],    out['''context''']						).input_ids
       _lowerCamelCase						:    Union[str, 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
              _lowerCamelCase						:    List[str]         =			[]
              _lowerCamelCase						:    int         =			[]
              _lowerCamelCase						:    Tuple         =			input_ids[:q_len]
              _lowerCamelCase						:    str         =			range(A_,    len(A_						),    max_length - doc_stride						)
              for i in doc_start_indices:
                     _lowerCamelCase						:    List[str]         =			i + max_length - q_len
                     _lowerCamelCase						:    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,
                  },
              }

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

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

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

       _lowerCamelCase						:    Tuple         =			input_ids[:q_len]
       _lowerCamelCase						:    Union[str, Any]         =			range(A_,    len(A_						),    max_length - doc_stride						)

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


       '''simple docstring'''





       _lowerCamelCase						:    Union[str, Any]         =			get_strided_contexts_and_ans(
           A_,    A_,    doc_stride=A_,    max_length=A_,    assertion=A_,    )

       return example






def 	snake_case_  (				A_ :		List[str],    A_ :		Union[str, Any]						):


       '''simple docstring'''





       with jsonlines.open(A_,    '''a'''						) as writer:
              for example in tqdm(A_,    total=len(A_						),    desc='''Saving samples ... '''						):
                     _lowerCamelCase						:    Optional[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__      =    load_dataset('''natural_questions''')
   lowerCAmelCase__      =    BigBirdTokenizer.from_pretrained('''google/bigbird-roberta-base''')

   lowerCAmelCase__      =    data['''train''' if PROCESS_TRAIN == '''true''' else '''validation''']

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

   np.random.seed(SEED)
   lowerCAmelCase__      =    '''nq-training.jsonl''' if PROCESS_TRAIN == '''true''' else '''nq-validation.jsonl'''
   save_to_disk(data, file_name=cache_file_name)







"""simple docstring"""






from __future__ import annotations

from itertools import permutations
from random import randint
from timeit import repeat






def 	snake_case_  (				):


       '''simple docstring'''





       _lowerCamelCase						:    Optional[Any]         =			[randint(-10_00,    10_00						) for i in range(10						)]
       _lowerCamelCase						:    Tuple         =			randint(-50_00,    50_00						)
       return (arr, r)


lowerCAmelCase__      =    make_dataset()






def 	snake_case_  (				A_ :		list[int],    A_ :		int						):


       '''simple docstring'''





       for triplet in permutations(A_,    3						):
              if sum(A_						) == target:
                     return tuple(sorted(A_						)						)
       return (0, 0, 0)






def 	snake_case_  (				A_ :		list[int],    A_ :		int						):


       '''simple docstring'''





       arr.sort()
       _lowerCamelCase						:    Optional[Any]         =			len(A_						)
       for i in range(n - 1						):
              _lowerCamelCase      ,   _lowerCamelCase						:    int         =			i + 1, n - 1
              while left < right:
                     if arr[i] + arr[left] + arr[right] == target:
                            return (arr[i], arr[left], arr[right])
                     elif arr[i] + arr[left] + arr[right] < target:
                            left += 1
                     elif arr[i] + arr[left] + arr[right] > target:
                            right -= 1
       return (0, 0, 0)






def 	snake_case_  (				):


       '''simple docstring'''





       _lowerCamelCase						:    Tuple         =			'''
from __main__ import dataset, triplet_sum1, triplet_sum2
'''
       _lowerCamelCase						:    List[Any]         =			'''
triplet_sum1(*dataset)
'''
       _lowerCamelCase						:    str         =			'''
triplet_sum2(*dataset)
'''
       _lowerCamelCase						:    Optional[int]         =			repeat(setup=A_,    stmt=A_,    repeat=5,    number=1_00_00						)
       _lowerCamelCase						:    Union[str, Any]         =			repeat(setup=A_,    stmt=A_,    repeat=5,    number=1_00_00						)
       return (min(A_						), min(A_						))


if __name__ == "__main__":
   from doctest import testmod

   testmod()
   lowerCAmelCase__      =    solution_times()
   print(F"""The time for naive implementation is {times[0]}.""")
   print(F"""The time for optimized implementation is {times[1]}.""")


