mathqa.py 2.04 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
"""
MathQA: Towards Interpretable Math Word Problem Solving with Operation-Based Formalisms
https://arxiv.org/pdf/1905.13319.pdf

MathQA is a large-scale dataset of 37k English multiple-choice math word problems
covering multiple math domain categories by modeling operation programs corresponding
to word problems in the AQuA dataset (Ling et al., 2017).

Homepage: https://math-qa.github.io/math-QA/
"""
11
import re
12
13
from lm_eval.base import MultipleChoiceTask

14

15
16
_CITATION = """
@misc{amini2019mathqa,
bzantium's avatar
bzantium committed
17
    title={MathQA: Towards Interpretable Math Word Problem Solving with Operation-Based Formalisms},
18
19
20
21
22
23
24
25
26
    author={Aida Amini and Saadia Gabriel and Peter Lin and Rik Koncel-Kedziorski and Yejin Choi and Hannaneh Hajishirzi},
    year={2019},
    eprint={1905.13319},
    archivePrefix={arXiv},
    primaryClass={cs.CL}
}
"""


Jonathan Tow's avatar
Jonathan Tow committed
27
class MathQA(MultipleChoiceTask):
Leo Gao's avatar
Leo Gao committed
28
    VERSION = 0
29
30
31
32
33
34
35
36
37
38
39
40
    DATASET_PATH = "math_qa"
    DATASET_NAME = None

    def has_training_docs(self):
        return True

    def has_validation_docs(self):
        return True

    def has_test_docs(self):
        return True

Jonathan Tow's avatar
Jonathan Tow committed
41
42
    def training_docs(self):
        if self._training_docs is None:
Jon Tow's avatar
Jon Tow committed
43
            self._training_docs = list(map(self._process_doc, self.dataset["train"]))
Jon Tow's avatar
Jon Tow committed
44
        return self._training_docs
Jonathan Tow's avatar
Jonathan Tow committed
45
46

    def validation_docs(self):
Jon Tow's avatar
Jon Tow committed
47
        return map(self._process_doc, self.dataset["validation"])
48

Jonathan Tow's avatar
Jonathan Tow committed
49
    def test_docs(self):
Jon Tow's avatar
Jon Tow committed
50
        return map(self._process_doc, self.dataset["test"])
Jonathan Tow's avatar
Jonathan Tow committed
51

Jon Tow's avatar
Jon Tow committed
52
    def _process_doc(self, doc):
bzantium's avatar
bzantium committed
53
54
55
56
57
        answer_idx = ["a", "b", "c", "d", "e"].index(doc["correct"])
        choices = [
            c[4:].rstrip(" ,")
            for c in re.findall(r"[abcd] \) .*?, |e \) .*?$", doc["options"])
        ]
58

59
        out_doc = {
bzantium's avatar
bzantium committed
60
            "query": "Question: " + doc["Problem"] + "\nAnswer:",
61
62
            "choices": choices,
            "gold": answer_idx,
63
64
        }
        return out_doc
65

66
67
    def doc_to_text(self, doc):
        return doc["query"]
bzantium's avatar
bzantium committed
68
69
70
71
72
73

    def should_decontaminate(self):
        return True

    def doc_to_decontamination_query(self, doc):
        return doc["query"]