


import os

import jsonlines
import numpy as np
from tqdm import tqdm


_SCREAMING_SNAKE_CASE			: Union[str, Any]			      =    20_48
_SCREAMING_SNAKE_CASE			: Tuple			      =    40_96
_SCREAMING_SNAKE_CASE			: Tuple			      =    42
_SCREAMING_SNAKE_CASE			: Optional[Any]			      =    os.environ.pop("PROCESS_TRAIN", "false")
_SCREAMING_SNAKE_CASE			: List[Any]			      =    {"null": 0, "short": 1, "long": 2, "yes": 3, "no": 4}
def 			UpperCAmelCase__						(UpperCamelCase_				):



     """simple docstring"""






     def choose_first(UpperCamelCase_							,UpperCamelCase_=False				):
          assert isinstance(A__							,A__				)
          if len(A__				) == 1:
               snake_case		 =				answer[0]
               return {k: [answer[k]] for k in answer} if is_long_answer else answer
          for a in answer:
               if is_long_answer:
                    snake_case		 =				{k: [a[k]] for k in a}
               if len(a['''start_token''']				) > 0:
                    break
          return a

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

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

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

     return answer
def 			UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_=False				):



     """simple docstring"""






     snake_case		 =				_get_single_answer(A__				)
     # 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
          snake_case		 =				example['''document''']['''tokens''']
          snake_case		 =				[]
          for i in range(len(doc['''token''']				)				):
               if not doc["is_html"][i]:
                    context.append(doc['''token'''][i]				)
          return {
              "context": " ".join(A__				),
              "answer": {
                  "start_token": -1_00,  # ignore index in cross-entropy
                  "end_token": -1_00,  # 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

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

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

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

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

     return {
         "context": " ".join(A__				),
         "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__						(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_=20_48							,UpperCamelCase_=40_96							,UpperCamelCase_=True				):



     """simple docstring"""






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

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

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

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

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

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

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

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

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

     snake_case		 =				input_ids[:q_len]
     snake_case		 =				range(A__							,len(A__				)							,max_length - doc_stride				)

     snake_case		 =				[]
     snake_case		 =				[]
     snake_case		 =				[]
     snake_case		 =				[]  # null, yes, no, long, short
     for i in doc_start_indices:
          snake_case		 =				i + max_length - q_len
          snake_case		 =				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:
               snake_case		 =				start_token - i + q_len
               snake_case		 =				end_token - i + q_len
               answers_category.append(answer['''category'''][0]				)  # ["short"] -> "short"
          else:
               snake_case		 =				-1_00
               snake_case		 =				-1_00
               answers_category.append('''null'''				)
          snake_case		 =				inputs[-1][start_token : end_token + 1]

          answers_start_token.append(A__				)
          answers_end_token.append(A__				)
          if assertion:
               if new != old and new != [tokenizer.cls_token_id]:
                    print('''ISSUE in strided for ID:'''							,example['''id''']				)
                    print('''New:'''							,tokenizer.decode(A__				)				)
                    print('''Old:'''							,tokenizer.decode(A__				)							,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__						(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_=20_48							,UpperCamelCase_=40_96							,UpperCamelCase_=False				):



     """simple docstring"""






     snake_case		 =				get_strided_contexts_and_ans(
         A__							,A__							,doc_stride=A__							,max_length=A__							,assertion=A__							,)

     return example
def 			UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_				):



     """simple docstring"""






     with jsonlines.open(A__							,'''a'''				) as writer:
          for example in tqdm(A__							,total=len(A__				)							,desc='''Saving samples ... '''				):
               snake_case		 =				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

 _SCREAMING_SNAKE_CASE			: Tuple			      =    load_dataset("natural_questions")
 _SCREAMING_SNAKE_CASE			: str			      =    BigBirdTokenizer.from_pretrained("google/bigbird-roberta-base")

 _SCREAMING_SNAKE_CASE			: List[Any]			      =    data["train" if PROCESS_TRAIN == "true" else "validation"]

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

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







import logging

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


lowercase_									=   logging.getLogger(__name__)





class  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE							):




  def   A__     (				self							,   lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase=None							,   lowerCAmelCase=None					)  -> List[str]:
        '''simple docstring'''






        _lowercase	       =self.layer[current_layer](lowerCAmelCase							,   lowerCAmelCase							,   head_mask[current_layer]					)

        _lowercase	       =layer_outputs[0]

        return hidden_states





@add_start_docstrings(
    """The bare Bert Model transformer with PABEE outputting raw hidden-states without any specific head on top."""			,			SCREAMING_SNAKE_CASE			,			)
class  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE							):




  def __init__(				self							,   lowerCAmelCase					)  -> Union[str, Any]:
        '''simple docstring'''






        super().__init__(lowerCAmelCase					)

        _lowercase	       =BertEncoderWithPabee(lowerCAmelCase					)

        self.init_weights()
        _lowercase	       =0
        _lowercase	       =0
        _lowercase	       =0

        _lowercase	       =0




  def   A__     (				self							,   lowerCAmelCase					)  -> Optional[Any]:
        '''simple docstring'''






        _lowercase	       =threshold




  def   A__     (				self							,   lowerCAmelCase					)  -> List[Any]:
        '''simple docstring'''






        _lowercase	       =patience




  def   A__     (				self					)  -> Dict:
        '''simple docstring'''






        _lowercase	       =0
        _lowercase	       =0




  def   A__     (				self					)  -> int:
        '''simple docstring'''






        _lowercase	       =self.inference_layers_num / self.inference_instances_num
        _lowercase	       =(
            F'''*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ='''
            F''' {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***'''
        )
        print(lowerCAmelCase					)




  @add_start_docstrings_to_model_forward(lowerCAmelCase					)
  def   A__     (				self							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=False							,   )  -> str:
        '''simple docstring'''






        if input_ids is not None and inputs_embeds is not None:
              raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time'					)
        elif input_ids is not None:
              _lowercase	       =input_ids.size()
        elif inputs_embeds is not None:
              _lowercase	       =inputs_embeds.size()[:-1]
        else:
              raise ValueError('You have to specify either input_ids or inputs_embeds'					)

        _lowercase	       =input_ids.device if input_ids is not None else inputs_embeds.device

        if attention_mask is None:
              _lowercase	       =torch.ones(lowerCAmelCase							,   device=lowerCAmelCase					)
        if token_type_ids is None:
              _lowercase	       =torch.zeros(lowerCAmelCase							,   dtype=torch.long							,   device=lowerCAmelCase					)

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        _lowercase	       =self.get_extended_attention_mask(lowerCAmelCase							,   lowerCAmelCase							,   lowerCAmelCase					)

        # If a 2D ou 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if self.config.is_decoder and encoder_hidden_states is not None:
              _lowercase       ,       _lowercase       ,       _lowercase	       =encoder_hidden_states.size()
              _lowercase	       =(encoder_batch_size, encoder_sequence_length)
              if encoder_attention_mask is None:
                    _lowercase	       =torch.ones(lowerCAmelCase							,   device=lowerCAmelCase					)
              _lowercase	       =self.invert_attention_mask(lowerCAmelCase					)
        else:
              _lowercase	       =None

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        _lowercase	       =self.get_head_mask(lowerCAmelCase							,   self.config.num_hidden_layers					)

        _lowercase	       =self.embeddings(
            input_ids=lowerCAmelCase							,   position_ids=lowerCAmelCase							,   token_type_ids=lowerCAmelCase							,   inputs_embeds=lowerCAmelCase					)
        _lowercase	       =embedding_output

        if self.training:
              _lowercase	       =[]
              for i in range(self.config.num_hidden_layers					):
                    _lowercase	       =self.encoder.adaptive_forward(
                        lowerCAmelCase							,   current_layer=lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   head_mask=lowerCAmelCase					)

                    _lowercase	       =self.pooler(lowerCAmelCase					)
                    _lowercase	       =output_layers[i](output_dropout(lowerCAmelCase					)					)
                    res.append(lowerCAmelCase					)
        elif self.patience == 0:  # Use all layers for inference
              _lowercase	       =self.encoder(
                  lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   head_mask=lowerCAmelCase							,   encoder_hidden_states=lowerCAmelCase							,   encoder_attention_mask=lowerCAmelCase							,   )
              _lowercase	       =self.pooler(encoder_outputs[0]					)
              _lowercase	       =[output_layers[self.config.num_hidden_layers - 1](lowerCAmelCase					)]
        else:
              _lowercase	       =0
              _lowercase	       =None
              _lowercase	       =0
              for i in range(self.config.num_hidden_layers					):
                    calculated_layer_num += 1
                    _lowercase	       =self.encoder.adaptive_forward(
                        lowerCAmelCase							,   current_layer=lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   head_mask=lowerCAmelCase					)

                    _lowercase	       =self.pooler(lowerCAmelCase					)
                    _lowercase	       =output_layers[i](lowerCAmelCase					)
                    if regression:
                          _lowercase	       =logits.detach()
                          if patient_result is not None:
                                _lowercase	       =patient_result.detach()
                          if (patient_result is not None) and torch.abs(patient_result - labels					) < self.regression_threshold:
                                patient_counter += 1
                          else:
                                _lowercase	       =0
                    else:
                          _lowercase	       =logits.detach().argmax(dim=1					)
                          if patient_result is not None:
                                _lowercase	       =patient_result.detach().argmax(dim=1					)
                          if (patient_result is not None) and torch.all(labels.eq(lowerCAmelCase					)					):
                                patient_counter += 1
                          else:
                                _lowercase	       =0

                    _lowercase	       =logits
                    if patient_counter == self.patience:
                          break
              _lowercase	       =[patient_result]
              self.inference_layers_num += calculated_layer_num
              self.inference_instances_num += 1

        return res





@add_start_docstrings(
    """Bert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of
    the pooled output) e.g. for GLUE tasks. """			,			SCREAMING_SNAKE_CASE			,			)
class  __lowerCAmelCase     (					SCREAMING_SNAKE_CASE							):




  def __init__(				self							,   lowerCAmelCase					)  -> List[Any]:
        '''simple docstring'''






        super().__init__(lowerCAmelCase					)
        _lowercase	       =config.num_labels

        _lowercase	       =BertModelWithPabee(lowerCAmelCase					)
        _lowercase	       =nn.Dropout(config.hidden_dropout_prob					)
        _lowercase	       =nn.ModuleList(
            [nn.Linear(config.hidden_size							,   self.config.num_labels					) for _ in range(config.num_hidden_layers					)]					)

        self.init_weights()




  @add_start_docstrings_to_model_forward(lowerCAmelCase					)
  def   A__     (				self							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   lowerCAmelCase=None							,   )  -> Union[str, Any]:
        '''simple docstring'''






        _lowercase	       =self.bert(
            input_ids=lowerCAmelCase							,   attention_mask=lowerCAmelCase							,   token_type_ids=lowerCAmelCase							,   position_ids=lowerCAmelCase							,   head_mask=lowerCAmelCase							,   inputs_embeds=lowerCAmelCase							,   output_dropout=self.dropout							,   output_layers=self.classifiers							,   regression=self.num_labels == 1							,   )

        _lowercase	       =(logits[-1],)

        if labels is not None:
              _lowercase	       =None
              _lowercase	       =0
              for ix, logits_item in enumerate(lowerCAmelCase					):
                    if self.num_labels == 1:
                          #  We are doing regression
                          _lowercase	       =MSELoss()
                          _lowercase	       =loss_fct(logits_item.view(-1					)							,   labels.view(-1					)					)
                    else:
                          _lowercase	       =CrossEntropyLoss()
                          _lowercase	       =loss_fct(logits_item.view(-1							,   self.num_labels					)							,   labels.view(-1					)					)
                    if total_loss is None:
                          _lowercase	       =loss
                    else:
                          total_loss += loss * (ix + 1)
                    total_weights += ix + 1
              _lowercase	       =(total_loss / total_weights,) + outputs

        return outputs


