arithmetic.py 3.34 KB
Newer Older
1
2
3
4
5
6
7
8
9
"""
Language Models are Few-Shot Learners
https://arxiv.org/pdf/2005.14165.pdf

A small battery of 10 tests that involve asking language models a simple arithmetic
problem in natural language.

Homepage: https://github.com/openai/gpt-3/tree/master/data
"""
Jonathan Tow's avatar
Jonathan Tow committed
10
11
import inspect
import lm_eval.datasets.arithmetic.arithmetic
&'s avatar
& committed
12
13
from lm_eval.base import Task, rf
from lm_eval.metrics import mean
14
15


16
17
18
19
20
21
22
23
24
25
26
27
28
_CITATION = """
@inproceedings{NEURIPS2020_1457c0d6,
    author = {Brown, Tom and Mann, Benjamin and Ryder, Nick and Subbiah, Melanie and Kaplan, Jared D and Dhariwal, Prafulla and Neelakantan, Arvind and Shyam, Pranav and Sastry, Girish and Askell, Amanda and Agarwal, Sandhini and Herbert-Voss, Ariel and Krueger, Gretchen and Henighan, Tom and Child, Rewon and Ramesh, Aditya and Ziegler, Daniel and Wu, Jeffrey and Winter, Clemens and Hesse, Chris and Chen, Mark and Sigler, Eric and Litwin, Mateusz and Gray, Scott and Chess, Benjamin and Clark, Jack and Berner, Christopher and McCandlish, Sam and Radford, Alec and Sutskever, Ilya and Amodei, Dario},
    booktitle = {Advances in Neural Information Processing Systems},
    editor = {H. Larochelle and M. Ranzato and R. Hadsell and M. F. Balcan and H. Lin},
    pages = {1877--1901},
    publisher = {Curran Associates, Inc.},
    title = {Language Models are Few-Shot Learners},
    url = {https://proceedings.neurips.cc/paper/2020/file/1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf},
    volume = {33},
    year = {2020}
}
"""
29

30

31
class Arithmetic(Task):
Leo Gao's avatar
Leo Gao committed
32
    VERSION = 0
Jonathan Tow's avatar
Jonathan Tow committed
33
    DATASET_PATH = inspect.getfile(lm_eval.datasets.arithmetic.arithmetic)
34

35
    def has_training_docs(self):
Leo Gao's avatar
Leo Gao committed
36
        return False
37
38
39
40
41
42
43
44

    def has_validation_docs(self):
        return True

    def has_test_docs(self):
        return False

    def training_docs(self):
Leo Gao's avatar
Leo Gao committed
45
        return NotImplemented
46
47

    def validation_docs(self):
Jonathan Tow's avatar
Jonathan Tow committed
48
        return self.dataset["validation"]
49
50
51
52
53

    def test_docs(self):
        return NotImplemented
    
    def doc_to_text(self, doc):
Jonathan Tow's avatar
Jonathan Tow committed
54
        return doc["context"]
55

56
57
58
59
    def should_decontaminate(self):
        return True

    def doc_to_decontamination_query(self, doc):
jon-tow's avatar
jon-tow committed
60
        return doc["context"]
61

62
    def doc_to_target(self, doc):
Jonathan Tow's avatar
Jonathan Tow committed
63
        return doc["completion"]
64
65

    def construct_requests(self, doc, ctx):
Jonathan Tow's avatar
Jonathan Tow committed
66
        ll, is_prediction = rf.loglikelihood(ctx, doc["completion"])
67
68
69
        return is_prediction

    def process_results(self, doc, results):
Leo Gao's avatar
Leo Gao committed
70
        is_prediction, = results
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
        return {
            "acc": is_prediction
        }

    def aggregation(self):
        return {
            "acc": mean,
        }

    def higher_is_better(self):
        return {
            "acc": True
        }


class Arithmetic2DPlus(Arithmetic):
Jonathan Tow's avatar
Jonathan Tow committed
87
88
    DATASET_NAME = "arithmetic_2da"

89
90

class Arithmetic2DMinus(Arithmetic):
Jonathan Tow's avatar
Jonathan Tow committed
91
92
    DATASET_NAME = "arithmetic_2ds"

93
94

class Arithmetic3DPlus(Arithmetic):
Jonathan Tow's avatar
Jonathan Tow committed
95
96
    DATASET_NAME = "arithmetic_3da"

97
98

class Arithmetic3DMinus(Arithmetic):
Jonathan Tow's avatar
Jonathan Tow committed
99
100
    DATASET_NAME = "arithmetic_3ds"

101
102

class Arithmetic4DPlus(Arithmetic):
Jonathan Tow's avatar
Jonathan Tow committed
103
104
    DATASET_NAME = "arithmetic_4da"

105
106

class Arithmetic4DMinus(Arithmetic):
Jonathan Tow's avatar
Jonathan Tow committed
107
108
    DATASET_NAME = "arithmetic_4ds"

109
110

class Arithmetic5DPlus(Arithmetic):
Jonathan Tow's avatar
Jonathan Tow committed
111
112
    DATASET_NAME = "arithmetic_5da"

113
114

class Arithmetic5DMinus(Arithmetic):
Jonathan Tow's avatar
Jonathan Tow committed
115
116
    DATASET_NAME = "arithmetic_5ds"

117
118

class Arithmetic2DMultiplication(Arithmetic):
Jonathan Tow's avatar
Jonathan Tow committed
119
120
    DATASET_NAME = "arithmetic_2dm"

121
122

class Arithmetic1DComposite(Arithmetic):
Jonathan Tow's avatar
Jonathan Tow committed
123
    DATASET_NAME = "arithmetic_1dc"