






'''simple docstring'''



import os

import jsonlines
import numpy as np
from tqdm import tqdm


__UpperCamelCase      :       List[Any]      			=      2048
__UpperCamelCase      :       List[str]      			=      4096
__UpperCamelCase      :       Any      			=      42
__UpperCamelCase      :       Dict      			=      os.environ.pop("""PROCESS_TRAIN""", """false""")
__UpperCamelCase      :       Optional[Any]      			=      {"""null""": 0, """short""": 1, """long""": 2, """yes""": 3, """no""": 4}



def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	str    )  ->			Dict:

   """simple docstring"""
   def choose_first(SCREAMING_SNAKE_CASE__:	str,					SCREAMING_SNAKE_CASE__:	int=False    ):
      assert isinstance(SCREAMING_SNAKE_CASE__,					SCREAMING_SNAKE_CASE__    )
      if len(SCREAMING_SNAKE_CASE__    ) == 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=SCREAMING_SNAKE_CASE__    )
         __a				     =    []
      answer.update(SCREAMING_SNAKE_CASE__    )

   # 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],					SCREAMING_SNAKE_CASE__    ) for k in cols    ):
      raise ValueError('Issue in ID',					example['id']    )

   return answer



def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	Union[str, Any],					SCREAMING_SNAKE_CASE__:	int=False    )  ->			List[str]:

   """simple docstring"""
   __a				     =    _get_single_answer(SCREAMING_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
      __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(SCREAMING_SNAKE_CASE__    ),
          "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

   __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(SCREAMING_SNAKE_CASE__    )    ) if not is_html[i]]    )
      if new != old:
         print('ID:',					example['id']    )
         print('New:',					SCREAMING_SNAKE_CASE__,					end='\n'    )
         print('Old:',					SCREAMING_SNAKE_CASE__,					end='\n\n'    )

   return {
       "context": " ".join(SCREAMING_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    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	int,					SCREAMING_SNAKE_CASE__:	Dict,					SCREAMING_SNAKE_CASE__:	Tuple=2048,					SCREAMING_SNAKE_CASE__:	List[Any]=4096,					SCREAMING_SNAKE_CASE__:	Dict=True    )  ->			Dict:

   """simple docstring"""
   # overlap will be of doc_stride - q_len

   __a				     =    get_context_and_ans(SCREAMING_SNAKE_CASE__,					assertion=SCREAMING_SNAKE_CASE__    )
   __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(SCREAMING_SNAKE_CASE__,					len(SCREAMING_SNAKE_CASE__    ),					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": [-100] * len(SCREAMING_SNAKE_CASE__    ),
              "end_token": [-100] * len(SCREAMING_SNAKE_CASE__    ),
              "category": category,
          },
      }

   __a				     =    out['context'].split()
   __a				     =    splitted_context[answer['end_token']]
   __a				     =    len(
       tokenizer(
           ' '.join(splitted_context[: answer['start_token']]    ),					add_special_tokens=SCREAMING_SNAKE_CASE__,					).input_ids    )
   __a				     =    len(
       tokenizer(' '.join(splitted_context[: answer['end_token']]    ),					add_special_tokens=SCREAMING_SNAKE_CASE__    ).input_ids    )

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

   # fixing end token
   __a				     =    len(tokenizer(SCREAMING_SNAKE_CASE__,					add_special_tokens=SCREAMING_SNAKE_CASE__    ).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(SCREAMING_SNAKE_CASE__    )
      if answer["span"] != new:
         print('ISSUE IN TOKENIZATION'    )
         print('OLD:',					answer['span']    )
         print('NEW:',					SCREAMING_SNAKE_CASE__,					end='\n\n'    )

   if len(SCREAMING_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"],
          },
      }

   __a				     =    input_ids[:q_len]
   __a				     =    range(SCREAMING_SNAKE_CASE__,					len(SCREAMING_SNAKE_CASE__    ),					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				     =    -100
         __a				     =    -100
         answers_category.append('null'    )
      __a				     =    inputs[-1][start_token : end_token + 1]

      answers_start_token.append(SCREAMING_SNAKE_CASE__    )
      answers_end_token.append(SCREAMING_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(SCREAMING_SNAKE_CASE__    )    )
            print('Old:',					tokenizer.decode(SCREAMING_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    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	List[str],					SCREAMING_SNAKE_CASE__:	Dict,					SCREAMING_SNAKE_CASE__:	int=2048,					SCREAMING_SNAKE_CASE__:	List[Any]=4096,					SCREAMING_SNAKE_CASE__:	List[Any]=False    )  ->			int:

   """simple docstring"""
   __a				     =    get_strided_contexts_and_ans(
       SCREAMING_SNAKE_CASE__,					SCREAMING_SNAKE_CASE__,					doc_stride=SCREAMING_SNAKE_CASE__,					max_length=SCREAMING_SNAKE_CASE__,					assertion=SCREAMING_SNAKE_CASE__,					)

   return example



def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	Tuple,					SCREAMING_SNAKE_CASE__:	str    )  ->			Union[str, Any]:

   """simple docstring"""
   with jsonlines.open(SCREAMING_SNAKE_CASE__,					'a'    ) as writer:
      for example in tqdm(SCREAMING_SNAKE_CASE__,					total=len(SCREAMING_SNAKE_CASE__    ),					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

 __UpperCamelCase      :       Union[str, Any]      			=      load_dataset("""natural_questions""")
 __UpperCamelCase      :       str      			=      BigBirdTokenizer.from_pretrained("""google/bigbird-roberta-base""")

 __UpperCamelCase      :       Union[str, Any]      			=      data["""train""" if PROCESS_TRAIN == """true""" else """validation"""]

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

 np.random.seed(SEED)
 __UpperCamelCase      :       List[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_sentencepiece_available,
    is_speech_available,
    is_tf_available,
    is_torch_available,
)


__UpperCamelCase      :       Optional[int]      			=      {
    """configuration_speech_to_text""": ["""SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """Speech2TextConfig"""],
    """processing_speech_to_text""": ["""Speech2TextProcessor"""],
}

try:
 if not is_sentencepiece_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __UpperCamelCase      :       Union[str, Any]      			=      ["""Speech2TextTokenizer"""]

try:
 if not is_speech_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __UpperCamelCase      :       Optional[Any]      			=      ["""Speech2TextFeatureExtractor"""]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __UpperCamelCase      :       List[Any]      			=      [
     """TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST""",
     """TFSpeech2TextForConditionalGeneration""",
     """TFSpeech2TextModel""",
     """TFSpeech2TextPreTrainedModel""",
 ]

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __UpperCamelCase      :       Tuple      			=      [
     """SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST""",
     """Speech2TextForConditionalGeneration""",
     """Speech2TextModel""",
     """Speech2TextPreTrainedModel""",
 ]


if TYPE_CHECKING:
 from .configuration_speech_to_text import SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP, SpeechaTextConfig
 from .processing_speech_to_text import SpeechaTextProcessor

 try:
  if not is_sentencepiece_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_speech_to_text import SpeechaTextTokenizer

 try:
  if not is_speech_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .feature_extraction_speech_to_text import SpeechaTextFeatureExtractor

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_speech_to_text import (
      TF_SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFSpeechaTextForConditionalGeneration,
      TFSpeechaTextModel,
      TFSpeechaTextPreTrainedModel,
  )

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_speech_to_text import (
      SPEECH_TO_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
      SpeechaTextForConditionalGeneration,
      SpeechaTextModel,
      SpeechaTextPreTrainedModel,
  )

else:
 import sys

 __UpperCamelCase      :       List[str]      			=      _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)