



'''simple docstring'''



import os

import jsonlines
import numpy as np
from tqdm import tqdm


lowerCAmelCase_     :       int								=						2048
lowerCAmelCase_     :       List[str]								=						4096
lowerCAmelCase_     :       List[Any]								=						42
lowerCAmelCase_     :       int								=						os.environ.pop('PROCESS_TRAIN', 'false')
lowerCAmelCase_     :       Optional[Any]								=						{'null': 0, 'short': 1, 'long': 2, 'yes': 3, 'no': 4}
def 					UpperCAmelCase   (						A  :				Optional[Any]				):
       def choose_first(A  :				Dict					,		A  :				str=False				):
              assert isinstance(lowerCamelCase_					,		lowerCamelCase_				)
              if len(lowerCamelCase_				) == 1:
                     SCREAMING_SNAKE_CASE					:      str     					=			answer[0]
                     return {k: [answer[k]] for k in answer} if is_long_answer else answer
              for a in answer:
                     if is_long_answer:
                            SCREAMING_SNAKE_CASE					:      Tuple     					=			{k: [a[k]] for k in a}
                     if len(a['''start_token''']				) > 0:
                            break
              return a

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

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

       SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			['''start_token''', '''end_token''', '''start_byte''', '''end_byte''', '''text''']
       if not all(isinstance(answer[k]					,		lowerCamelCase_				) for k in cols				):
              raise ValueError('''Issue in ID'''					,		example['''id''']				)

       return answer
def 					UpperCAmelCase   (						A  :				Union[str, Any]					,		A  :				Optional[int]=False				):
       SCREAMING_SNAKE_CASE					:      int     					=			_get_single_answer(lowerCamelCase_				)
       # 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
              SCREAMING_SNAKE_CASE					:      List[Any]     					=			example['''document''']['''tokens''']
              SCREAMING_SNAKE_CASE					:      List[Any]     					=			[]
              for i in range(len(doc['''token''']				)				):
                     if not doc["is_html"][i]:
                            context.append(doc['''token'''][i]				)
              return {
                  "context": " ".join(lowerCamelCase_				),
                  "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

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

       SCREAMING_SNAKE_CASE					:      Tuple     					=			example['''document''']['''tokens''']
       SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			answer['''start_token''']
       SCREAMING_SNAKE_CASE					:      Any     					=			answer['''end_token''']

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

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

       return {
           "context": " ".join(lowerCamelCase_				),
           "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   (						A  :				Tuple					,		A  :				List[Any]					,		A  :				Union[str, Any]=2048					,		A  :				str=4096					,		A  :				Union[str, Any]=True				):
       # overlap will be of doc_stride - q_len

       SCREAMING_SNAKE_CASE					:      str     					=			get_context_and_ans(lowerCamelCase_					,		assertion=lowerCamelCase_				)
       SCREAMING_SNAKE_CASE					:      Dict     					=			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"],
                  },
              }

       SCREAMING_SNAKE_CASE					:      Any     					=			tokenizer(example['''question''']['''text''']					,		out['''context''']				).input_ids
       SCREAMING_SNAKE_CASE					:      Dict     					=			input_ids.index(tokenizer.sep_token_id				) + 1

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

       SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			out['''context'''].split()
       SCREAMING_SNAKE_CASE					:      List[Any]     					=			splitted_context[answer['''end_token''']]
       SCREAMING_SNAKE_CASE					:      List[Any]     					=			len(
           tokenizer(
               ''' '''.join(splitted_context[: answer['''start_token''']]				)					,		add_special_tokens=lowerCamelCase_					,		).input_ids				)
       SCREAMING_SNAKE_CASE					:      List[Any]     					=			len(
           tokenizer(''' '''.join(splitted_context[: answer['''end_token''']]				)					,		add_special_tokens=lowerCamelCase_				).input_ids				)

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

       # fixing end token
       SCREAMING_SNAKE_CASE					:      Optional[int]     					=			len(tokenizer(lowerCamelCase_					,		add_special_tokens=lowerCamelCase_				).input_ids				)
       if num_sub_tokens > 1:
              answer["end_token"] += num_sub_tokens - 1

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

       if assertion:
              SCREAMING_SNAKE_CASE					:      Optional[Any]     					=			tokenizer.decode(lowerCamelCase_				)
              if answer["span"] != new:
                     print('''ISSUE IN TOKENIZATION'''				)
                     print('''OLD:'''					,		answer['''span''']				)
                     print('''NEW:'''					,		lowerCamelCase_					,		end='''\n\n'''				)

       if len(lowerCamelCase_				) <= 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"],
                  },
              }

       SCREAMING_SNAKE_CASE					:      Union[str, Any]     					=			input_ids[:q_len]
       SCREAMING_SNAKE_CASE					:      List[str]     					=			range(lowerCamelCase_					,		len(lowerCamelCase_				)					,		max_length - doc_stride				)

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

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

       return example
def 					UpperCAmelCase   (						A  :				List[str]					,		A  :				Dict				):
       with jsonlines.open(lowerCamelCase_					,		'''a'''				) as writer:
              for example in tqdm(lowerCamelCase_					,		total=len(lowerCamelCase_				)					,		desc='''Saving samples ... '''				):
                     SCREAMING_SNAKE_CASE					:      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_     :       List[str]								=						load_dataset('natural_questions')
       lowerCAmelCase_     :       Dict								=						BigBirdTokenizer.from_pretrained('google/bigbird-roberta-base')

       lowerCAmelCase_     :       int								=						data['train' if PROCESS_TRAIN == 'true' else 'validation']

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

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









'''simple docstring'''



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


lowerCAmelCase_     :       Tuple								=						{
    'configuration_tapas': ['TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP', 'TapasConfig'],
    'tokenization_tapas': ['TapasTokenizer'],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowerCAmelCase_     :       Dict								=						[
							    'TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST',
							    'TapasForMaskedLM',
							    'TapasForQuestionAnswering',
							    'TapasForSequenceClassification',
							    'TapasModel',
							    'TapasPreTrainedModel',
							    'load_tf_weights_in_tapas',
							]
try:
							if not is_tf_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowerCAmelCase_     :       str								=						[
							    'TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST',
							    'TFTapasForMaskedLM',
							    'TFTapasForQuestionAnswering',
							    'TFTapasForSequenceClassification',
							    'TFTapasModel',
							    'TFTapasPreTrainedModel',
							]


if TYPE_CHECKING:
							from .configuration_tapas import TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP, TapasConfig
							from .tokenization_tapas import TapasTokenizer

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_tapas import (
														    TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
														    TapasForMaskedLM,
														    TapasForQuestionAnswering,
														    TapasForSequenceClassification,
														    TapasModel,
														    TapasPreTrainedModel,
														    load_tf_weights_in_tapas,
														)

							try:
														if not is_tf_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_tf_tapas import (
														    TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
														    TFTapasForMaskedLM,
														    TFTapasForQuestionAnswering,
														    TFTapasForSequenceClassification,
														    TFTapasModel,
														    TFTapasPreTrainedModel,
														)


else:
							import sys

							lowerCAmelCase_     :       List[Any]								=						_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




