


import os

import jsonlines
import numpy as np
from tqdm import tqdm


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

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

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

   UpperCAmelCase_							:      int						=   ['''start_token''', '''end_token''', '''start_byte''', '''end_byte''', '''text''']
   if not all(isinstance(answer[k]			,       __lowercase       ) for k in cols       ):
      raise ValueError('''Issue in ID'''			,       example['''id''']       )

   return answer
def      snake_case_						( __lowercase			,       __lowercase=False       ):
   UpperCAmelCase_							:      Dict						=   _get_single_answer(__lowercase       )
   # 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
      UpperCAmelCase_							:      Optional[int]						=   example['''document''']['''tokens''']
      UpperCAmelCase_							:      Dict						=   []
      for i in range(len(doc['''token''']       )       ):
         if not doc["is_html"][i]:
            context.append(doc['''token'''][i]       )
      return {
          "context": " ".join(__lowercase       ),
          "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

   UpperCAmelCase_							:      Dict						=   ['''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

   UpperCAmelCase_							:      Tuple						=   example['''document''']['''tokens''']
   UpperCAmelCase_							:      Any						=   answer['''start_token''']
   UpperCAmelCase_							:      List[Any]						=   answer['''end_token''']

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

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

   return {
       "context": " ".join(__lowercase       ),
       "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_						( __lowercase			,       __lowercase			,       __lowercase=2_0_4_8			,       __lowercase=4_0_9_6			,       __lowercase=True       ):
   # overlap will be of doc_stride - q_len

   UpperCAmelCase_							:      Optional[Any]						=   get_context_and_ans(__lowercase			,       assertion=__lowercase       )
   UpperCAmelCase_							:      str						=   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"],
          },
      }

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

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

   UpperCAmelCase_							:      Union[str, Any]						=   out['''context'''].split()
   UpperCAmelCase_							:      int						=   splitted_context[answer['''end_token''']]
   UpperCAmelCase_							:      Optional[Any]						=   len(
       tokenizer(
           ''' '''.join(splitted_context[: answer['''start_token''']]       )			,       add_special_tokens=__lowercase			,       ).input_ids       )
   UpperCAmelCase_							:      Optional[int]						=   len(
       tokenizer(''' '''.join(splitted_context[: answer['''end_token''']]       )			,       add_special_tokens=__lowercase       ).input_ids       )

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

   # fixing end token
   UpperCAmelCase_							:      List[str]						=   len(tokenizer(__lowercase			,       add_special_tokens=__lowercase       ).input_ids       )
   if num_sub_tokens > 1:
      answer["end_token"] += num_sub_tokens - 1

   UpperCAmelCase_							:      Tuple						=   input_ids[answer['''start_token'''] : answer['''end_token'''] + 1]  # right & left are inclusive
   UpperCAmelCase_							:      Dict						=   answer['''start_token''']
   UpperCAmelCase_							:      int						=   answer['''end_token''']

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

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

   UpperCAmelCase_							:      Optional[Any]						=   input_ids[:q_len]
   UpperCAmelCase_							:      int						=   range(__lowercase			,       len(__lowercase       )			,       max_length - doc_stride       )

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

      answers_start_token.append(__lowercase       )
      answers_end_token.append(__lowercase       )
      if assertion:
         if new != old and new != [tokenizer.cls_token_id]:
            print('''ISSUE in strided for ID:'''			,       example['''id''']       )
            print('''New:'''			,       tokenizer.decode(__lowercase       )       )
            print('''Old:'''			,       tokenizer.decode(__lowercase       )			,       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_						( __lowercase			,       __lowercase			,       __lowercase=2_0_4_8			,       __lowercase=4_0_9_6			,       __lowercase=False       ):
   UpperCAmelCase_							:      List[str]						=   get_strided_contexts_and_ans(
       __lowercase			,       __lowercase			,       doc_stride=__lowercase			,       max_length=__lowercase			,       assertion=__lowercase			,       )

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

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

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

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



import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__UpperCamelCase				:  Union[str, Any]      					=							logging.get_logger(__name__)

__UpperCamelCase				:  str      					=							{
    'microsoft/unispeech-sat-base-100h-libri-ft': (
        'https://huggingface.co/microsoft/unispeech-sat-base-100h-libri-ft/resolve/main/config.json'
    ),
    # See all UniSpeechSat models at https://huggingface.co/models?filter=unispeech_sat
}





class  lowerCAmelCase__(    snake_case__						):






    '''simple docstring'''





    A_  :     List[Any]			= 'unispeech-sat'





    def __init__(	self				:   int       ,   __snake_case				:   Optional[int]=32       ,   __snake_case				:   Dict=768       ,   __snake_case				:   Optional[Any]=12       ,   __snake_case				:   Optional[int]=12       ,   __snake_case				:   Dict=3_072       ,   __snake_case				:   List[str]="gelu"       ,   __snake_case				:   Any=0.1       ,   __snake_case				:   Tuple=0.1       ,   __snake_case				:   Optional[int]=0.1       ,   __snake_case				:   Tuple=0.0       ,   __snake_case				:   List[Any]=0.0       ,   __snake_case				:   Tuple=0.1       ,   __snake_case				:   Union[str, Any]=0.1       ,   __snake_case				:   Dict=0.02       ,   __snake_case				:   Optional[Any]=1E-5       ,   __snake_case				:   Optional[int]="group"       ,   __snake_case				:   str="gelu"       ,   __snake_case				:   Union[str, Any]=(512, 512, 512, 512, 512, 512, 512)       ,   __snake_case				:   str=(5, 2, 2, 2, 2, 2, 2)       ,   __snake_case				:   Tuple=(10, 3, 3, 3, 3, 2, 2)       ,   __snake_case				:   int=False       ,   __snake_case				:   Optional[int]=128       ,   __snake_case				:   Any=16       ,   __snake_case				:   Union[str, Any]=False       ,   __snake_case				:   Union[str, Any]=True       ,   __snake_case				:   List[Any]=0.05       ,   __snake_case				:   Dict=10       ,   __snake_case				:   int=2       ,   __snake_case				:   Optional[Any]=0.0       ,   __snake_case				:   Optional[int]=10       ,   __snake_case				:   List[Any]=0       ,   __snake_case				:   Optional[int]=320       ,   __snake_case				:   int=2       ,   __snake_case				:   Any=0.1       ,   __snake_case				:   Optional[int]=100       ,   __snake_case				:   Tuple=256       ,   __snake_case				:   List[str]=256       ,   __snake_case				:   List[Any]=0.1       ,   __snake_case				:   Tuple="mean"       ,   __snake_case				:   List[Any]=False       ,   __snake_case				:   List[str]=False       ,   __snake_case				:   Optional[Any]=256       ,   __snake_case				:   Tuple=(512, 512, 512, 512, 1_500)       ,   __snake_case				:   Optional[int]=(5, 3, 3, 1, 1)       ,   __snake_case				:   Any=(1, 2, 3, 1, 1)       ,   __snake_case				:   int=512       ,   __snake_case				:   Optional[int]=0       ,   __snake_case				:   Dict=1       ,   __snake_case				:   Tuple=2       ,   __snake_case				:   Union[str, Any]=504       ,   **__snake_case				:   List[str]       ,   ):





       '''simple docstring'''
       super().__init__(**__snake_case       ,   pad_token_id=__snake_case       ,   bos_token_id=__snake_case       ,   eos_token_id=__snake_case	)
       UpperCAmelCase_							:      Union[str, Any]						=   hidden_size
       UpperCAmelCase_							:      int						=   feat_extract_norm
       UpperCAmelCase_							:      Dict						=   feat_extract_activation
       UpperCAmelCase_							:      Union[str, Any]						=   list(__snake_case	)
       UpperCAmelCase_							:      List[str]						=   list(__snake_case	)
       UpperCAmelCase_							:      Any						=   list(__snake_case	)
       UpperCAmelCase_							:      Any						=   conv_bias
       UpperCAmelCase_							:      List[str]						=   num_conv_pos_embeddings
       UpperCAmelCase_							:      Dict						=   num_conv_pos_embedding_groups
       UpperCAmelCase_							:      Optional[int]						=   len(self.conv_dim	)
       UpperCAmelCase_							:      List[str]						=   num_hidden_layers
       UpperCAmelCase_							:      Dict						=   intermediate_size
       UpperCAmelCase_							:      Union[str, Any]						=   hidden_act
       UpperCAmelCase_							:      str						=   num_attention_heads
       UpperCAmelCase_							:      Union[str, Any]						=   hidden_dropout
       UpperCAmelCase_							:      List[str]						=   attention_dropout
       UpperCAmelCase_							:      Optional[Any]						=   activation_dropout
       UpperCAmelCase_							:      Dict						=   feat_proj_dropout
       UpperCAmelCase_							:      Optional[Any]						=   final_dropout
       UpperCAmelCase_							:      List[Any]						=   layerdrop
       UpperCAmelCase_							:      int						=   layer_norm_eps
       UpperCAmelCase_							:      int						=   initializer_range
       UpperCAmelCase_							:      List[Any]						=   vocab_size
       UpperCAmelCase_							:      int						=   num_clusters
       UpperCAmelCase_							:      int						=   do_stable_layer_norm
       UpperCAmelCase_							:      Any						=   use_weighted_layer_sum

       if (
           (len(self.conv_stride	) != self.num_feat_extract_layers)
           or (len(self.conv_kernel	) != self.num_feat_extract_layers)
           or (len(self.conv_dim	) != self.num_feat_extract_layers)
       ):
          raise ValueError(
              '''Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` =='''
              ''' `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) ='''
              f''' {len(self.conv_dim	)}`, `len(config.conv_stride) = {len(self.conv_stride	)}`,'''
              f''' `len(config.conv_kernel) = {len(self.conv_kernel	)}`.'''	)

       # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
       UpperCAmelCase_							:      int						=   apply_spec_augment
       UpperCAmelCase_							:      Optional[Any]						=   mask_time_prob
       UpperCAmelCase_							:      str						=   mask_time_length
       UpperCAmelCase_							:      Any						=   mask_time_min_masks
       UpperCAmelCase_							:      str						=   mask_feature_prob
       UpperCAmelCase_							:      str						=   mask_feature_length
       UpperCAmelCase_							:      Tuple						=   mask_feature_min_masks

       # parameters for pretraining with codevector quantized representations
       UpperCAmelCase_							:      Optional[int]						=   num_codevectors_per_group
       UpperCAmelCase_							:      int						=   num_codevector_groups
       UpperCAmelCase_							:      List[str]						=   contrastive_logits_temperature
       UpperCAmelCase_							:      int						=   feat_quantizer_dropout
       UpperCAmelCase_							:      List[str]						=   num_negatives
       UpperCAmelCase_							:      Any						=   codevector_dim
       UpperCAmelCase_							:      Tuple						=   proj_codevector_dim
       UpperCAmelCase_							:      Union[str, Any]						=   diversity_loss_weight

       # ctc loss
       UpperCAmelCase_							:      Any						=   ctc_loss_reduction
       UpperCAmelCase_							:      Optional[Any]						=   ctc_zero_infinity

       # SequenceClassification-specific parameter. Feel free to ignore for other classes.
       UpperCAmelCase_							:      Optional[int]						=   classifier_proj_size

       # XVector-specific parameters. Feel free to ignore for other classes.
       UpperCAmelCase_							:      Dict						=   list(__snake_case	)
       UpperCAmelCase_							:      Dict						=   list(__snake_case	)
       UpperCAmelCase_							:      Dict						=   list(__snake_case	)
       UpperCAmelCase_							:      Union[str, Any]						=   xvector_output_dim





    @property
    def 					_lowerCamelCase (	self				:   List[str]	):





       '''simple docstring'''
       return functools.reduce(operator.mul       ,   self.conv_stride       ,   1	)