


'''simple docstring'''

import os

import jsonlines
import numpy as np
from tqdm import tqdm


lowerCamelCase				:  Optional[int]    		=      2_048
lowerCamelCase				:  List[Any]    		=      4_096
lowerCamelCase				:  Dict    		=      42
lowerCamelCase				:  Optional[Any]    		=      os.environ.pop('PROCESS_TRAIN', 'false')
lowerCamelCase				:  str    		=      {'null': 0, 'short': 1, 'long': 2, 'yes': 3, 'no': 4}

def _SCREAMING_SNAKE_CASE (A			)			->			Optional[Any]:






       """simple docstring"""







       def choose_first(A    ,					A=False			):
              assert isinstance(_lowerCAmelCase    ,					_lowerCAmelCase			)
              if len(_lowerCAmelCase			) == 1:
                     lowercase__      =					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__      =					{k: [a[k]] for k in a}
                     if len(a['''start_token''']			) > 0:
                            break
              return a

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

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

       lowercase__      =					['''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 _SCREAMING_SNAKE_CASE (A    ,					A=False			)			->			Union[str, Any]:






       """simple docstring"""







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

       lowercase__      =					['''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__      =					example['''document''']['''tokens''']
       lowercase__      =					answer['''start_token''']
       lowercase__      =					answer['''end_token''']

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

       # checking above code
       if assertion:
              lowercase__      =					doc['''is_html'''][answer['''start_token'''] : answer['''end_token''']]
              lowercase__      =					doc['''token'''][answer['''start_token'''] : answer['''end_token''']]
              lowercase__      =					''' '''.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 _SCREAMING_SNAKE_CASE (A    ,					A    ,					A=2_048    ,					A=4_096    ,					A=True			)			->			Any:






       """simple docstring"""







       lowercase__      =					get_context_and_ans(_lowerCAmelCase    ,					assertion=_lowerCAmelCase			)
       lowercase__      =					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__      =					tokenizer(example['''question''']['''text''']    ,					out['''context''']			).input_ids
       lowercase__      =					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__      =					[]
              lowercase__      =					[]
              lowercase__      =					input_ids[:q_len]
              lowercase__      =					range(_lowerCAmelCase    ,					len(_lowerCAmelCase			)    ,					max_length - doc_stride			)
              for i in doc_start_indices:
                     lowercase__      =					i + max_length - q_len
                     lowercase__      =					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,
                  },
              }

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

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

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

       lowercase__      =					input_ids[:q_len]
       lowercase__      =					range(_lowerCAmelCase    ,					len(_lowerCAmelCase			)    ,					max_length - doc_stride			)

       lowercase__      =					[]
       lowercase__      =					[]
       lowercase__      =					[]
       lowercase__      =					[]  # null, yes, no, long, short
       for i in doc_start_indices:
              lowercase__      =					i + max_length - q_len
              lowercase__      =					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__      =					start_token - i + q_len
                     lowercase__      =					end_token - i + q_len
                     answers_category.append(answer['''category'''][0]			)  # ["short"] -> "short"
              else:
                     lowercase__      =					-100
                     lowercase__      =					-100
                     answers_category.append('''null'''			)
              lowercase__      =					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 _SCREAMING_SNAKE_CASE (A    ,					A    ,					A=2_048    ,					A=4_096    ,					A=False			)			->			Any:






       """simple docstring"""







       lowercase__      =					get_strided_contexts_and_ans(
           _lowerCAmelCase    ,					_lowerCAmelCase    ,					doc_stride=_lowerCAmelCase    ,					max_length=_lowerCAmelCase    ,					assertion=_lowerCAmelCase    ,					)

       return example

def _SCREAMING_SNAKE_CASE (A    ,					A			)			->			List[str]:






       """simple docstring"""







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

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

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

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






"""simple docstring"""
import os
import socket
from contextlib import contextmanager

import torch

from ..commands.config.default import write_basic_config  # noqa: F401
from ..state import PartialState
from .dataclasses import DistributedType
from .imports import is_deepspeed_available, is_tpu_available
from .transformer_engine import convert_model
from .versions import is_torch_version


if is_deepspeed_available():
			from deepspeed import DeepSpeedEngine

if is_tpu_available(check_device=False):
			import torch_xla.core.xla_model as xm

def  lowercase     (_lowerCAmelCase						):
						if is_torch_version("""<""" ,    """2.0.0"""						) or not hasattr(_lowerCAmelCase ,    """_dynamo"""						):
												return False
						return isinstance(_lowerCAmelCase ,    torch._dynamo.eval_frame.OptimizedModule						)

def  lowercase     (_lowerCAmelCase ,    _lowerCAmelCase = True						):
						__lowerCAmelCase     =  (torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)

						__lowerCAmelCase     =  is_compiled_module(_lowerCAmelCase						)
						if is_compiled:
												__lowerCAmelCase     =  model
												__lowerCAmelCase     =  model._orig_mod

						if is_deepspeed_available():
												options += (DeepSpeedEngine,)

						while isinstance(_lowerCAmelCase ,    _lowerCAmelCase						):
												__lowerCAmelCase     =  model.module

						if not keep_fpaa_wrapper:
												__lowerCAmelCase     =  getattr(_lowerCAmelCase ,    """forward"""						)
												__lowerCAmelCase     =  model.__dict__.pop("""_original_forward""" ,    _lowerCAmelCase						)
												if original_forward is not None:
																		while hasattr(_lowerCAmelCase ,    """__wrapped__"""						):
																								__lowerCAmelCase     =  forward.__wrapped__
																								if forward == original_forward:
																														break
																		__lowerCAmelCase     =  forward
												if getattr(_lowerCAmelCase ,    """_converted_to_transformer_engine""" ,    _lowerCAmelCase						):
																		convert_model(_lowerCAmelCase ,    to_transformer_engine=_lowerCAmelCase						)

						if is_compiled:
												__lowerCAmelCase     =  model
												__lowerCAmelCase     =  compiled_model

						return model

def  lowercase     ():
						PartialState().wait_for_everyone()

def  lowercase     (_lowerCAmelCase ,    _lowerCAmelCase						):
						if PartialState().distributed_type == DistributedType.TPU:
												xm.save(_lowerCAmelCase ,    _lowerCAmelCase						)
						elif PartialState().local_process_index == 0:
												torch.save(_lowerCAmelCase ,    _lowerCAmelCase						)

@contextmanager
def  lowercase     (**_lowerCAmelCase						):
						for key, value in kwargs.items():
												__lowerCAmelCase     =  str(_lowerCAmelCase						)

						yield

						for key in kwargs:
												if key.upper() in os.environ:
																		del os.environ[key.upper()]

def  lowercase     (_lowerCAmelCase						):
						if not hasattr(_lowerCAmelCase ,    """__qualname__"""						) and not hasattr(_lowerCAmelCase ,    """__name__"""						):
												__lowerCAmelCase     =  getattr(_lowerCAmelCase ,    """__class__""" ,    _lowerCAmelCase						)
						if hasattr(_lowerCAmelCase ,    """__qualname__"""						):
												return obj.__qualname__
						if hasattr(_lowerCAmelCase ,    """__name__"""						):
												return obj.__name__
						return str(_lowerCAmelCase						)

def  lowercase     (_lowerCAmelCase ,    _lowerCAmelCase						):
						for key, value in source.items():
												if isinstance(_lowerCAmelCase ,    _lowerCAmelCase						):
																		__lowerCAmelCase     =  destination.setdefault(_lowerCAmelCase ,    {}						)
																		merge_dicts(_lowerCAmelCase ,    _lowerCAmelCase						)
												else:
																		__lowerCAmelCase     =  value

						return destination



def  lowercase     (_lowerCAmelCase = None						):
						if port is None:
												__lowerCAmelCase     =  2_9500
						with socket.socket(socket.AF_INET ,    socket.SOCK_STREAM						) as s:
												return s.connect_ex(("""localhost""", port)						) == 0
