create_eli5_dataset.py

from datasets import load_dataset

dataset = load_dataset("eli5")

print(dataset)

path = "/Users/danielk/ideaProjects/qoogle-experiments/data"
trainfile = open(path + "/eli5/train.tsv", "w")
testfile = open(path + "/eli5/test.tsv", "w")
devfile = open(path + "/eli5/dev.tsv", "w")

def safe(split, outfile, is_eval):
    for x in dataset[split]:
        if is_eval:
            ans = "///".join([a for a in x['answers']['text']]).replace("\t", " ").replace("\n", " ")
        else:
            ids_and_scores = list(enumerate(x['answers']['score']))
            ids_and_scores = sorted(ids_and_scores, key=lambda x: -x[1])
            max_score_idx = ids_and_scores[0][0]
            ans = x['answers']['text'][max_score_idx].replace("\t", " ").replace("\n", " ")
        question = x['title'].replace("\t", " ").replace("\n", " ")
        if len(question.split(" ")) > 510:
            question = " ".join(question.split(" ")[:510])
        if len(ans.split(" ")) > 510:
            ans = " ".join(ans.split(" ")[:510])
        outfile.write(f"{question} \t {ans} \n")


safe('validation_eli5', devfile, True)
safe('test_eli5', testfile, True)
safe('train_eli5', trainfile, False)