


import os

import jsonlines
import numpy as np
from tqdm import tqdm


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






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

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

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

							a__             =     ['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 							__lowercase   ( __lowerCAmelCase							:				Tuple					,    __lowerCAmelCase							:				str=False   ):
							a__             =     _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
														a__             =     example['document']['tokens']
														a__             =     []
														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": -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

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

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

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

							# checking above code
							if assertion:
														a__             =     doc['is_html'][answer['start_token'] : answer['end_token']]
														a__             =     doc['token'][answer['start_token'] : answer['end_token']]
														a__             =     ' '.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 							__lowercase   ( __lowerCAmelCase							:				Optional[int]					,    __lowerCAmelCase							:				Any					,    __lowerCAmelCase							:				List[str]=2_0_4_8					,    __lowerCAmelCase							:				int=4_0_9_6					,    __lowerCAmelCase							:				int=True   ):
							# overlap will be of doc_stride - q_len

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

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

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

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

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

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

							a__             =     input_ids[:q_len]
							a__             =     range(__lowerCAmelCase					,    len(__lowerCAmelCase   )					,    max_length - doc_stride   )

							a__             =     []
							a__             =     []
							a__             =     []
							a__             =     []  # null, yes, no, long, short
							for i in doc_start_indices:
														a__             =     i + max_length - q_len
														a__             =     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:
																					a__             =     start_token - i + q_len
																					a__             =     end_token - i + q_len
																					answers_category.append(answer['category'][0]   )  # ["short"] -> "short"
														else:
																					a__             =     -1_0_0
																					a__             =     -1_0_0
																					answers_category.append('null'   )
														a__             =     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 							__lowercase   ( __lowerCAmelCase							:				Dict					,    __lowerCAmelCase							:				Any					,    __lowerCAmelCase							:				List[str]=2_0_4_8					,    __lowerCAmelCase							:				Optional[int]=4_0_9_6					,    __lowerCAmelCase							:				Any=False   ):
							a__             =     get_strided_contexts_and_ans(
							    __lowerCAmelCase					,    __lowerCAmelCase					,    doc_stride=__lowerCAmelCase					,    max_length=__lowerCAmelCase					,    assertion=__lowerCAmelCase					,    )

							return example






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

						snake_case							:       Any     			=      data['''train''' if PROCESS_TRAIN == '''true''' else '''validation''']

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

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








import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_distilbert import DistilBertTokenizer


snake_case							:       Dict     			=      logging.get_logger(__name__)

snake_case							:       Dict     			=      {'''vocab_file''': '''vocab.txt''', '''tokenizer_file''': '''tokenizer.json'''}

snake_case							:       List[Any]     			=      {
    '''vocab_file''': {
        '''distilbert-base-uncased''': '''https://huggingface.co/distilbert-base-uncased/resolve/main/vocab.txt''',
        '''distilbert-base-uncased-distilled-squad''': (
            '''https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/vocab.txt'''
        ),
        '''distilbert-base-cased''': '''https://huggingface.co/distilbert-base-cased/resolve/main/vocab.txt''',
        '''distilbert-base-cased-distilled-squad''': (
            '''https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/vocab.txt'''
        ),
        '''distilbert-base-german-cased''': '''https://huggingface.co/distilbert-base-german-cased/resolve/main/vocab.txt''',
        '''distilbert-base-multilingual-cased''': (
            '''https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/vocab.txt'''
        ),
    },
    '''tokenizer_file''': {
        '''distilbert-base-uncased''': '''https://huggingface.co/distilbert-base-uncased/resolve/main/tokenizer.json''',
        '''distilbert-base-uncased-distilled-squad''': (
            '''https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/tokenizer.json'''
        ),
        '''distilbert-base-cased''': '''https://huggingface.co/distilbert-base-cased/resolve/main/tokenizer.json''',
        '''distilbert-base-cased-distilled-squad''': (
            '''https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/tokenizer.json'''
        ),
        '''distilbert-base-german-cased''': (
            '''https://huggingface.co/distilbert-base-german-cased/resolve/main/tokenizer.json'''
        ),
        '''distilbert-base-multilingual-cased''': (
            '''https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/tokenizer.json'''
        ),
    },
}

snake_case							:       int     			=      {
    '''distilbert-base-uncased''': 5_12,
    '''distilbert-base-uncased-distilled-squad''': 5_12,
    '''distilbert-base-cased''': 5_12,
    '''distilbert-base-cased-distilled-squad''': 5_12,
    '''distilbert-base-german-cased''': 5_12,
    '''distilbert-base-multilingual-cased''': 5_12,
}


snake_case							:       Union[str, Any]     			=      {
    '''distilbert-base-uncased''': {'''do_lower_case''': True},
    '''distilbert-base-uncased-distilled-squad''': {'''do_lower_case''': True},
    '''distilbert-base-cased''': {'''do_lower_case''': False},
    '''distilbert-base-cased-distilled-squad''': {'''do_lower_case''': False},
    '''distilbert-base-german-cased''': {'''do_lower_case''': False},
    '''distilbert-base-multilingual-cased''': {'''do_lower_case''': False},
}







class        snake_case_   (lowerCamelCase_      ):
					UpperCAmelCase__					:    Any          =			VOCAB_FILES_NAMES
					UpperCAmelCase__					:    str          =			PRETRAINED_VOCAB_FILES_MAP
					UpperCAmelCase__					:    int          =			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					UpperCAmelCase__					:    Tuple          =			PRETRAINED_INIT_CONFIGURATION
					UpperCAmelCase__					:    Optional[Any]          =			['''input_ids''', '''attention_mask''']
					UpperCAmelCase__					:    Optional[int]          =			DistilBertTokenizer
					def __init__(      self   :Dict			,__snake_case   :Dict=None			,__snake_case   :Optional[Any]=None			,__snake_case   :Optional[Any]=True			,__snake_case   :List[Any]="[UNK]"			,__snake_case   :str="[SEP]"			,__snake_case   :List[Any]="[PAD]"			,__snake_case   :Tuple="[CLS]"			,__snake_case   :Optional[int]="[MASK]"			,__snake_case   :Dict=True			,__snake_case   :Dict=None			,**__snake_case   :List[Any]			,)					->    Optional[int]:
												super().__init__(
												    __snake_case			,tokenizer_file=__snake_case			,do_lower_case=__snake_case			,unk_token=__snake_case			,sep_token=__snake_case			,pad_token=__snake_case			,cls_token=__snake_case			,mask_token=__snake_case			,tokenize_chinese_chars=__snake_case			,strip_accents=__snake_case			,**__snake_case			,)

												a__             =     json.loads(self.backend_tokenizer.normalizer.__getstate__()      )
												if (
												    normalizer_state.get('lowercase'			,__snake_case      ) != do_lower_case
												    or normalizer_state.get('strip_accents'			,__snake_case      ) != strip_accents
												    or normalizer_state.get('handle_chinese_chars'			,__snake_case      ) != tokenize_chinese_chars
												):
																			a__             =     getattr(__snake_case			,normalizer_state.pop('type'      )      )
																			a__             =     do_lower_case
																			a__             =     strip_accents
																			a__             =     tokenize_chinese_chars
																			a__             =     normalizer_class(**__snake_case      )

												a__             =     do_lower_case
					def  lowerCamelCase__(      self   :Any			,__snake_case   :List[str]			,__snake_case   :int=None      )					->    Dict:
												a__             =     [self.cls_token_id] + token_ids_a + [self.sep_token_id]

												if token_ids_a:
																			output += token_ids_a + [self.sep_token_id]

												return output
					def  lowerCamelCase__(      self   :List[str]			,__snake_case   :List[int]			,__snake_case   :Optional[List[int]] = None      )					->    List[int]:
												a__             =     [self.sep_token_id]
												a__             =     [self.cls_token_id]
												if token_ids_a is None:
																			return len(cls + token_ids_a + sep      ) * [0]
												return len(cls + token_ids_a + sep      ) * [0] + len(token_ids_a + sep      ) * [1]



					def  lowerCamelCase__(      self   :Union[str, Any]			,__snake_case   :str			,__snake_case   :Optional[str] = None      )					->    Tuple[str]:
												a__             =     self._tokenizer.model.save(__snake_case			,name=__snake_case      )
												return tuple(__snake_case      )




