




"""simple docstring"""







import os

import jsonlines
import numpy as np
from tqdm import tqdm


SCREAMING_SNAKE_CASE_	     =       2048
SCREAMING_SNAKE_CASE_	     =       4096
SCREAMING_SNAKE_CASE_	     =       42
SCREAMING_SNAKE_CASE_	     =       os.environ.pop('PROCESS_TRAIN', 'false')
SCREAMING_SNAKE_CASE_	     =       {'null': 0, 'short': 1, 'long': 2, 'yes': 3, 'no': 4}





def      __snake_case (				_lowercase    ):


    """simple docstring"""







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

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

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

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

    return answer





def      __snake_case (				_lowercase			,_lowercase=False    ):


    """simple docstring"""







    UpperCamelCase              =	_get_single_answer(UpperCamelCase__    )
    # 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              =	example["""document"""]["""tokens"""]
        UpperCamelCase              =	[]
        for i in range(len(doc['''token''']    )    ):
            if not doc["is_html"][i]:
                context.append(doc['''token'''][i]    )
        return {
            "context": " ".join(UpperCamelCase__    ),
            "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

    UpperCamelCase              =	["""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              =	example["""document"""]["""tokens"""]
    UpperCamelCase              =	answer["""start_token"""]
    UpperCamelCase              =	answer["""end_token"""]

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

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

    return {
        "context": " ".join(UpperCamelCase__    ),
        "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=2048			,_lowercase=4096			,_lowercase=True    ):


    """simple docstring"""







    UpperCamelCase              =	get_context_and_ans(UpperCamelCase__			,assertion=UpperCamelCase__    )
    UpperCamelCase              =	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              =	tokenizer(example['''question''']['''text''']			,out['''context''']    ).input_ids
    UpperCamelCase              =	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              =	[]
        UpperCamelCase              =	[]
        UpperCamelCase              =	input_ids[:q_len]
        UpperCamelCase              =	range(UpperCamelCase__			,len(UpperCamelCase__    )			,max_length - doc_stride    )
        for i in doc_start_indices:
            UpperCamelCase              =	i + max_length - q_len
            UpperCamelCase              =	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(UpperCamelCase__    ),
                "end_token": [-100] * len(UpperCamelCase__    ),
                "category": category,
            },
        }

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

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

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

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

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

    if len(UpperCamelCase__    ) <= 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              =	input_ids[:q_len]
    UpperCamelCase              =	range(UpperCamelCase__			,len(UpperCamelCase__    )			,max_length - doc_stride    )

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

        answers_start_token.append(UpperCamelCase__    )
        answers_end_token.append(UpperCamelCase__    )
        if assertion:
            if new != old and new != [tokenizer.cls_token_id]:
                print('''ISSUE in strided for ID:'''			,example['''id''']    )
                print('''New:'''			,tokenizer.decode(UpperCamelCase__    )    )
                print('''Old:'''			,tokenizer.decode(UpperCamelCase__    )			,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=2048			,_lowercase=4096			,_lowercase=False    ):


    """simple docstring"""







    UpperCamelCase              =	get_strided_contexts_and_ans(
        UpperCamelCase__			,UpperCamelCase__			,doc_stride=UpperCamelCase__			,max_length=UpperCamelCase__			,assertion=UpperCamelCase__			,)

    return example





def      __snake_case (				_lowercase			,_lowercase    ):


    """simple docstring"""







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



def  _SCREAMING_SNAKE_CASE     (      UpperCamelCase__    :					list[float]				):




     """simple docstring"""


     if len(UpperCamelCase__				) < 2:
          raise ValueError("""Monogons and Digons are not polygons in the Euclidean space"""				)
     if any(i <= 0 for i in nums				):
          raise ValueError("""All values must be greater than 0"""				)
     a_     :							int     						=  nums.copy()
     copy_nums.sort()
     return copy_nums[-1] < sum(copy_nums[:-1]				)


if __name__ == "__main__":
 import doctest

 doctest.testmod()
