





'''simple docstring'''



import os

import jsonlines
import numpy as np
from tqdm import tqdm


_SCREAMING_SNAKE_CASE					=    2_0_4_8
_SCREAMING_SNAKE_CASE					=    4_0_9_6
_SCREAMING_SNAKE_CASE					=    4_2
_SCREAMING_SNAKE_CASE					=    os.environ.pop('''PROCESS_TRAIN''', '''false''')
_SCREAMING_SNAKE_CASE					=    {'''null''': 0, '''short''': 1, '''long''': 2, '''yes''': 3, '''no''': 4}




def    _lowerCAmelCase						(				lowerCamelCase_			:  Optional[Any]				):
   def choose_first(lowerCamelCase_			:  Dict		,       lowerCamelCase_			:  str=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    _lowerCAmelCase						(				lowerCamelCase_			:  Union[str, Any]		,       lowerCamelCase_			:  Optional[int]=False				):
   __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": -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

   __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    _lowerCAmelCase						(				lowerCamelCase_			:  Tuple		,       lowerCamelCase_			:  List[Any]		,       lowerCamelCase_			:  Union[str, Any]=2_0_4_8		,       lowerCamelCase_			:  str=4_0_9_6		,       lowerCamelCase_			:  Union[str, Any]=True				):
   # overlap will be of doc_stride - q_len

   __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": [-1_0_0] * len(lowerCamelCase_				),
              "end_token": [-1_0_0] * 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										=							-1_0_0
         __lowercase										=							-1_0_0
         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    _lowerCAmelCase						(				lowerCamelCase_			:  Tuple		,       lowerCamelCase_			:  str		,       lowerCamelCase_			:  List[str]=2_0_4_8		,       lowerCamelCase_			:  Optional[int]=4_0_9_6		,       lowerCamelCase_			:  Tuple=False				):
   __lowercase										=							get_strided_contexts_and_ans(
       lowerCamelCase_		,       lowerCamelCase_		,       doc_stride=lowerCamelCase_		,       max_length=lowerCamelCase_		,       assertion=lowerCamelCase_		,       )

   return example




def    _lowerCAmelCase						(				lowerCamelCase_			:  List[str]		,       lowerCamelCase_			:  Dict				):
   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

 _SCREAMING_SNAKE_CASE					=    load_dataset('''natural_questions''')
 _SCREAMING_SNAKE_CASE					=    BigBirdTokenizer.from_pretrained('''google/bigbird-roberta-base''')

 _SCREAMING_SNAKE_CASE					=    data['''train''' if PROCESS_TRAIN == '''true''' else '''validation''']

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

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












'''simple docstring'''



from scipy.stats import pearsonr, spearmanr
from sklearn.metrics import fa_score, matthews_corrcoef

import datasets


_SCREAMING_SNAKE_CASE					=    '''\
@inproceedings{wang2019glue,
  title={{GLUE}: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding},
  author={Wang, Alex and Singh, Amanpreet and Michael, Julian and Hill, Felix and Levy, Omer and Bowman, Samuel R.},
  note={In the Proceedings of ICLR.},
  year={2019}
}
'''

_SCREAMING_SNAKE_CASE					=    '''\
GLUE, the General Language Understanding Evaluation benchmark
(https://gluebenchmark.com/) is a collection of resources for training,
evaluating, and analyzing natural language understanding systems.
'''

_SCREAMING_SNAKE_CASE					=    '''
Compute GLUE evaluation metric associated to each GLUE dataset.
Args:
    predictions: list of predictions to score.
        Each translation should be tokenized into a list of tokens.
    references: list of lists of references for each translation.
        Each reference should be tokenized into a list of tokens.
Returns: depending on the GLUE subset, one or several of:
    "accuracy": Accuracy
    "f1": F1 score
    "pearson": Pearson Correlation
    "spearmanr": Spearman Correlation
    "matthews_correlation": Matthew Correlation
Examples:

    >>> glue_metric = datasets.load_metric(\'glue\', \'sst2\')  # \'sst2\' or any of ["mnli", "mnli_mismatched", "mnli_matched", "qnli", "rte", "wnli", "hans"]
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'accuracy\': 1.0}

    >>> glue_metric = datasets.load_metric(\'glue\', \'mrpc\')  # \'mrpc\' or \'qqp\'
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'accuracy\': 1.0, \'f1\': 1.0}

    >>> glue_metric = datasets.load_metric(\'glue\', \'stsb\')
    >>> references = [0., 1., 2., 3., 4., 5.]
    >>> predictions = [0., 1., 2., 3., 4., 5.]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print({"pearson": round(results["pearson"], 2), "spearmanr": round(results["spearmanr"], 2)})
    {\'pearson\': 1.0, \'spearmanr\': 1.0}

    >>> glue_metric = datasets.load_metric(\'glue\', \'cola\')
    >>> references = [0, 1]
    >>> predictions = [0, 1]
    >>> results = glue_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'matthews_correlation\': 1.0}
'''




def    _lowerCAmelCase						(				lowerCamelCase_			:  Any		,       lowerCamelCase_			:  int				):
			return float((preds == labels).mean()				)




def    _lowerCAmelCase						(				lowerCamelCase_			:  List[str]		,       lowerCamelCase_			:  str				):
			__lowercase										=							simple_accuracy(lowerCamelCase_		,       lowerCamelCase_				)
			__lowercase										=							float(fa_score(y_true=lowerCamelCase_		,       y_pred=lowerCamelCase_				)				)
			return {
			    "accuracy": acc,
			    "f1": fa,
			}




def    _lowerCAmelCase						(				lowerCamelCase_			:  List[Any]		,       lowerCamelCase_			:  Any				):
			__lowercase										=							float(pearsonr(lowerCamelCase_		,       lowerCamelCase_				)[0]				)
			__lowercase										=							float(spearmanr(lowerCamelCase_		,       lowerCamelCase_				)[0]				)
			return {
			    "pearson": pearson_corr,
			    "spearmanr": spearman_corr,
			}






@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION ,      _KWARGS_DESCRIPTION						)
class 							__lowercase       (      datasets.Metric						):







	'''simple docstring'''

	def 						_UpperCAmelCase     (self )							->							str:







				'''simple docstring'''





				if self.config_name not in [
				    "sst2",
				    "mnli",
				    "mnli_mismatched",
				    "mnli_matched",
				    "cola",
				    "stsb",
				    "mrpc",
				    "qqp",
				    "qnli",
				    "rte",
				    "wnli",
				    "hans",
				]:
							raise KeyError(
							    '''You should supply a configuration name selected in '''
							    '''["sst2", "mnli", "mnli_mismatched", "mnli_matched", '''
							    '''"cola", "stsb", "mrpc", "qqp", "qnli", "rte", "wnli", "hans"]''' )
				return datasets.MetricInfo(
				    description=_DESCRIPTION     ,citation=_CITATION     ,inputs_description=_KWARGS_DESCRIPTION     ,features=datasets.Features(
				        {
				            '''predictions''': datasets.Value('''int64''' if self.config_name != '''stsb''' else '''float32''' ),
				            '''references''': datasets.Value('''int64''' if self.config_name != '''stsb''' else '''float32''' ),
				        } )     ,codebase_urls=[]     ,reference_urls=[]     ,format='''numpy'''     ,)


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase     ,_lowerCamelCase )							->							Dict:







				'''simple docstring'''





				if self.config_name == "cola":
							return {"matthews_correlation": matthews_corrcoef(_lowerCamelCase     ,_lowerCamelCase )}
				elif self.config_name == "stsb":
							return pearson_and_spearman(_lowerCamelCase     ,_lowerCamelCase )
				elif self.config_name in ["mrpc", "qqp"]:
							return acc_and_fa(_lowerCamelCase     ,_lowerCamelCase )
				elif self.config_name in ["sst2", "mnli", "mnli_mismatched", "mnli_matched", "qnli", "rte", "wnli", "hans"]:
							return {"accuracy": simple_accuracy(_lowerCamelCase     ,_lowerCamelCase )}
				else:
							raise KeyError(
							    '''You should supply a configuration name selected in '''
							    '''["sst2", "mnli", "mnli_mismatched", "mnli_matched", '''
							    '''"cola", "stsb", "mrpc", "qqp", "qnli", "rte", "wnli", "hans"]''' )





