answer will be long if short is not available disregard some samples gives new context after removing html new answer tokens as per new context answer getsingleanswerexample bytes are of no use del answerstartbyte del answerendbyte handle yesno answers explicitly if answercategory0 in yes no category is list with one element doc exampledocumenttokens context for i in rangelendoctoken if not docishtmli context appenddoctokeni return context joincontext answer starttoken 100 ignore index in crossentropy endtoken 100 ignore index in crossentropy category answercategory span answercategory extra later help in removing all no answers if answerstarttoken 1 return context none answer starttoken 1 endtoken 1 category null span none extra handling normal samples cols starttoken endtoken answer updatek answerk0 if lenanswerk 0 else answerk for k in cols e g 10 10 doc exampledocumenttokens starttoken answerstarttoken endtoken answerendtoken context for i in rangelendoctoken if not docishtmli context appenddoctokeni else if answerstarttoken i starttoken 1 if answerendtoken i endtoken 1 new joincontextstarttoken endtoken checking above code if assertion overlap will be of docstride qlen later removing these samples return yesno fixing end token this won t match exactly because of extra gaps visaully inspect everything new tokenizer decodeold if answerspan new printissue in tokenization printold answerspan printnew new endnn if leninputids maxlength return exampleid exampleid inputids inputids labels starttoken answerstarttoken endtoken answerendtoken category answercategory qindices inputids qlen docstartindices rangeqlen leninputids maxlength docstride inputs answersstarttoken answersendtoken answerscategory null yes no long short for i in docstartindices endindex i maxlength qlen slice inputidsi endindex inputs appendqindices slice assert leninputs1 maxlength issue in truncating length if starttoken i and endtoken endindex 1 starttoken starttoken i qlen endtoken endtoken i qlen answerscategory appendanswercategory0 short short else starttoken 100 endtoken 100 answerscategory appendnull new inputs1starttoken endtoken 1 answersstarttoken appendstarttoken answersendtoken appendendtoken if assertion running area from datasets import loaddataset from transformers import bigbirdtokenizer data loaddatasetnaturalquestions tokenizer bigbirdtokenizer frompretrainedgooglebigbirdrobertabase data datatrain if processtrain true else validation fnkwargs tokenizer tokenizer docstride docstride maxlength maxlength assertion false data data mapprepareinputs fnkwargsfnkwargs data data removecolumnsannotations document id question printdata np random seedseed cachefilename nqtraining jsonl if processtrain true else nqvalidation jsonl savetodiskdata filenamecachefilename answer will be long if short is not available disregard some samples gives new context after removing html new answer tokens as per new context bytes are of no use handle yes_no answers explicitly category is list with one element ignore index in cross entropy ignore index in cross entropy extra later help in removing all no answers extra handling normal samples e g 10 10 checking above code checking if above code is working as expected for all the samples this makes it inclusive either long or short extra overlap will be of doc_stride q_len later removing these samples return yes no category is list with one element fixing end token right left are inclusive this won t match exactly because of extra gaps visaully inspect everything null yes no long short short short checking if above code is working as expected for all the samples leave waste samples with no answer removing 50 samples running area
import os

import jsonlines
import numpy as np
from tqdm import tqdm


DOC_STRIDE = 2048
MAX_LENGTH = 4096
SEED = 42
PROCESS_TRAIN = os.environ.pop("PROCESS_TRAIN", "false")
CATEGORY_MAPPING = {"null": 0, "short": 1, "long": 2, "yes": 3, "no": 4}


def _get_single_answer(example):
    def choose_first(answer, is_long_answer=False):
        assert isinstance(answer, list)
        if len(answer) == 1:
            answer = 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

    answer = {"id": example["id"]}
    annotation = example["annotations"]
    yes_no_answer = annotation["yes_no_answer"]
    if 0 in yes_no_answer or 1 in yes_no_answer:
        answer["category"] = ["yes"] if 1 in yes_no_answer else ["no"]
        answer["start_token"] = answer["end_token"] = []
        answer["start_byte"] = answer["end_byte"] = []
        answer["text"] = ["<cls>"]
    else:
        answer["category"] = ["short"]
        out = choose_first(annotation["short_answers"])
        if len(out["start_token"]) == 0:
            
            answer["category"] = ["long"]
            out = choose_first(annotation["long_answer"], is_long_answer=True)
            out["text"] = []
        answer.update(out)

    
    if len(answer["start_token"]) > 1 or answer["start_token"] == answer["end_token"]:
        answer["remove_it"] = True
    else:
        answer["remove_it"] = False

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

    return answer


def get_context_and_ans(example, assertion=False):
    
    answer = _get_single_answer(example)
    
    del answer["start_byte"]
    del answer["end_byte"]

    
    if answer["category"][0] in ["yes", "no"]:  
        doc = example["document"]["tokens"]
        context = []
        for i in range(len(doc["token"])):
            if not doc["is_html"][i]:
                context.append(doc["token"][i])
        return {
            "context": " ".join(context),
            "answer": {
                "start_token": -100,  
                "end_token": -100,  
                "category": answer["category"],
                "span": answer["category"],  
            },
        }

    
    if answer["start_token"] == [-1]:
        return {
            "context": "None",
            "answer": {
                "start_token": -1,
                "end_token": -1,
                "category": "null",
                "span": "None",  
            },
        }

    

    cols = ["start_token", "end_token"]
    answer.update({k: answer[k][0] if len(answer[k]) > 0 else answer[k] for k in cols})  

    doc = example["document"]["tokens"]
    start_token = answer["start_token"]
    end_token = answer["end_token"]

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

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

    return {
        "context": " ".join(context),
        "answer": {
            "start_token": start_token,
            "end_token": end_token - 1,  
            "category": answer["category"],  
            "span": new,  
        },
    }


def get_strided_contexts_and_ans(example, tokenizer, doc_stride=2048, max_length=4096, assertion=True):
    

    out = get_context_and_ans(example, assertion=assertion)
    answer = out["answer"]

    
    if answer["start_token"] == -1:
        return {
            "example_id": example["id"],
            "input_ids": [[-1]],
            "labels": {
                "start_token": [-1],
                "end_token": [-1],
                "category": ["null"],
            },
        }

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

    
    if answer["category"][0] in ["yes", "no"]:  
        inputs = []
        category = []
        q_indices = input_ids[:q_len]
        doc_start_indices = range(q_len, len(input_ids), max_length - doc_stride)
        for i in doc_start_indices:
            end_index = i + max_length - q_len
            slice = 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(category),
                "end_token": [-100] * len(category),
                "category": category,
            },
        }

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

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

    
    num_sub_tokens = len(tokenizer(complete_end_token, add_special_tokens=False).input_ids)
    if num_sub_tokens > 1:
        answer["end_token"] += num_sub_tokens - 1

    old = input_ids[answer["start_token"] : answer["end_token"] + 1]  
    start_token = answer["start_token"]
    end_token = answer["end_token"]

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

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

    q_indices = input_ids[:q_len]
    doc_start_indices = range(q_len, len(input_ids), max_length - doc_stride)

    inputs = []
    answers_start_token = []
    answers_end_token = []
    answers_category = []  
    for i in doc_start_indices:
        end_index = i + max_length - q_len
        slice = 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:
            start_token = start_token - i + q_len
            end_token = end_token - i + q_len
            answers_category.append(answer["category"][0])  
        else:
            start_token = -100
            end_token = -100
            answers_category.append("null")
        new = inputs[-1][start_token : end_token + 1]

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

    return example


def save_to_disk(hf_data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(hf_data, total=len(hf_data), desc="Saving samples ... "):
            labels = 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  
                if cat == "null" and np.random.rand() < 0.6:
                    continue  
                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

    data = load_dataset("natural_questions")
    tokenizer = BigBirdTokenizer.from_pretrained("google/bigbird-roberta-base")

    data = data["train" if PROCESS_TRAIN == "true" else "validation"]

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

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