goldenswag / utils.py
marthemidtgaard's picture
Deploy app
7266988 verified
Raw History Blame Contribute Delete
664 Bytes
import re
import datasets
def preprocess(text):
text = text.strip()
# NOTE: Brackets are artifacts of the WikiHow dataset portion of HellaSwag.
text = text.replace(" [tittel]", ". ")
text = re.sub("\\[.*?\\]", "", text)
text = text.replace(" ", " ")
return text
def process_docs(dataset: datasets.Dataset) -> datasets.Dataset:
def _process_doc(doc):
out_doc = {
"query": preprocess(doc["activity_label"] + ": " + doc["ctx"]),
"choices": [preprocess(ending) for ending in doc["endings"]],
"gold": int(doc["label"]),
}
return out_doc
return dataset.map(_process_doc)