arithmetic.py 3.22 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
Jonathan Tow's avatar
Jonathan Tow committed
14

15

16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
_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}
}
"""


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

    def doc_to_target(self, doc):
Jonathan Tow's avatar
Jonathan Tow committed
57
        return doc["completion"]
58
59

    def construct_requests(self, doc, ctx):
Jonathan Tow's avatar
Jonathan Tow committed
60
        ll, is_prediction = rf.loglikelihood(ctx, doc["completion"])
jon-tow's avatar
jon-tow committed
61
        return ll, is_prediction
62
63

    def process_results(self, doc, results):
jon-tow's avatar
jon-tow committed
64
65
        print(results)
        results = results
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
        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
82
83
    DATASET_NAME = "arithmetic_2da"

84
85

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

88
89

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

92
93

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

96
97

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

100
101

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

104
105

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

108
109

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

112
113

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

116
117

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