base.py 9.37 KB
Newer Older
Leo Gao's avatar
Leo Gao committed
1
2
import abc
import random
thefazzer's avatar
thefazzer committed
3
import numpy as np
&'s avatar
& committed
4
5

from lm_eval.metrics import mean
Jason Phang's avatar
gpt3  
Jason Phang committed
6

Jason Phang's avatar
Jason Phang committed
7

Leo Gao's avatar
Leo Gao committed
8
9
class LM(abc.ABC):
    @abc.abstractmethod
Leo Gao's avatar
Leo Gao committed
10
    def loglikelihood(self, requests):
Leo Gao's avatar
Leo Gao committed
11
12
13
        """Compute log-likelihood of generating a continuation from a context.
        Downstream tasks should attempt to use loglikelihood instead of other 
        LM calls whenever possible.
Jason Phang's avatar
gpt3  
Jason Phang committed
14

Leo Gao's avatar
Leo Gao committed
15
16
17
        :param requests: list
            A list of pairs (context, continuation)
            context: str
Leo Gao's avatar
Leo Gao committed
18
19
                Context string. Implementations of LM must be able to handle an 
                empty context string.
Leo Gao's avatar
Leo Gao committed
20
21
22
23
24
25
26
27
28
29
30
31
32
            continuation: str
                The continuation over which log likelihood will be calculated. If 
                there is a word boundary, the space should be in the continuation. 
                For example, context="hello" continuation=" world" is correct.
        :return: list
            A list of pairs (logprob, isgreedy)
            logprob: float
                The log probability of `contination`
            isgreedy:
                Whether `contination` would be generated by greedy sampling from `context`
        """
        pass

&'s avatar
& committed
33
    # TODO: Add an optional max length
Leo Gao's avatar
Leo Gao committed
34
    @abc.abstractmethod
Leo Gao's avatar
Update  
Leo Gao committed
35
    def greedy_until(self, requests):
Leo Gao's avatar
Leo Gao committed
36
37
38
39
40
41
        """Generate greedily until a stopping sequence

        :param requests: list
            A list of pairs (context, until)
            context: str
                Context string
Leo Gao's avatar
Leo Gao committed
42
43
44
            until: [str]
                The string sequences to generate until. These string sequences 
                may each span across multiple tokens, or may be part of one token.
Leo Gao's avatar
Leo Gao committed
45
46
47
48
        :return: list
            A list of strings continuation
            continuation: str
                The generated continuation.
Jason Phang's avatar
gpt3  
Jason Phang committed
49
        """
Leo Gao's avatar
Leo Gao committed
50
51
        pass

Jason Phang's avatar
gpt3  
Jason Phang committed
52
53
54
55
56
57
58
59
60
61
62
    @classmethod
    def create_from_arg_string(cls, arg_string):
        """Constructor method, in case models need additional arguments
        e.g. OpenAI API engine, paths for loading, other params

        :param arg_string: str
            Left up to individual model class to handle

        """
        return cls()

Leo Gao's avatar
Leo Gao committed
63

64
class Task(abc.ABC):
&'s avatar
& committed
65
66
67
68
69
70
71
72
    """A task represents an entire benchmark including its dataset, problems,
    answers, and evaluation methods. See BoolQ for a simple example implementation

    A `doc` can be any python object which represents one instance of evaluation.
    This is usually a dictionary e.g.
        {"question": ..., "answer": ...} or
        {"question": ..., question, answer)
    """
Leo Gao's avatar
Leo Gao committed
73
74
    def __init__(self):
        self.download()
75
        self._training_docs = None
sdtblck's avatar
sdtblck committed
76
77
78
79
80

    def download(self):
        """Downloads the task dataset if necessary"""
        pass

81
82
    @abc.abstractmethod
    def has_training_docs(self):
Jason Phang's avatar
checkin  
Jason Phang committed
83
        """Whether the task has a training set"""
84
        pass
85

86
87
    @abc.abstractmethod
    def has_validation_docs(self):
Jason Phang's avatar
checkin  
Jason Phang committed
88
89
90
91
92
93
        """Whether the task has a validation set"""
        pass

    @abc.abstractmethod
    def has_test_docs(self):
        """Whether the task has a test set"""
94
95
        pass

Leo Gao's avatar
Leo Gao committed
96
    def training_docs(self):
Jason Phang's avatar
checkin  
Jason Phang committed
97
98
99
100
        """
        :return: Iterable[obj]
            A iterable of any object, that doc_to_text can handle
        """
Leo Gao's avatar
Leo Gao committed
101
        return []
102

Leo Gao's avatar
Leo Gao committed
103
    def validation_docs(self):
104
105
106
107
        """
        :return: Iterable[obj]
            A iterable of any object, that doc_to_text can handle
        """
Leo Gao's avatar
Leo Gao committed
108
        return []
109

Leo Gao's avatar
Leo Gao committed
110
    def test_docs(self):
111
112
113
114
        """
        :return: Iterable[obj]
            A iterable of any object, that doc_to_text can handle
        """
Leo Gao's avatar
Leo Gao committed
115
        return []
Leo Gao's avatar
Leo Gao committed
116

117
    def fewshot_examples(self, k):
118
119
        if self._training_docs is None:
            self._training_docs = list(self.training_docs())
Leo Gao's avatar
Leo Gao committed
120
121
122
        rnd = random.Random()
        rnd.seed(42)
        return rnd.sample(self._training_docs, k)
Leo Gao's avatar
Leo Gao committed
123
124

    @abc.abstractmethod
Leo Gao's avatar
Update  
Leo Gao committed
125
126
127
128
129
    def doc_to_text(self, doc):
        pass

    @abc.abstractmethod
    def doc_to_target(self, doc):
Leo Gao's avatar
Leo Gao committed
130
        pass
Leo Gao's avatar
Leo Gao committed
131
132

    @abc.abstractmethod
133
    def construct_requests(self, doc, ctx):
Leo Gao's avatar
Leo Gao committed
134
135
136
        """ Uses RequestFactory to construct Requests and returns an iterable of 
        Requests which will be sent to the LM.

137
138
        :param doc:
            The document as returned from training_docs, validation_docs, or test_docs.
Leo Gao's avatar
Leo Gao committed
139
        :param ctx: str
140
141
142
            The context string, generated by fewshot_context. This includes the natural 
            language description, as well as the few shot examples, and the question
            part of the document for `doc`. 
Leo Gao's avatar
Leo Gao committed
143
        """
Leo Gao's avatar
Leo Gao committed
144
        pass
145

Leo Gao's avatar
Leo Gao committed
146
    @abc.abstractmethod
Leo Gao's avatar
Leo Gao committed
147
    def process_results(self, doc, results):
Leo Gao's avatar
Update  
Leo Gao committed
148
        """Take a single document and the LM results and evaluates, returning a 
149
150
        dict where keys are the names of submetrics and values are the values of 
        the metric for that one document
Leo Gao's avatar
Leo Gao committed
151
152
153
154
155

        :param doc:
            The document as returned from training_docs, validation_docs, or test_docs.
        :param results:
            The results of the requests created in construct_requests.
Jason Phang's avatar
checkin  
Jason Phang committed
156
        """
Leo Gao's avatar
Leo Gao committed
157
        pass
Jason Phang's avatar
gpt3  
Jason Phang committed
158

159
160
161
    @abc.abstractmethod
    def aggregation(self):
        """
&'s avatar
& committed
162
        :returns: {str: [metric_score] -> float}
163
            A dictionary where keys are the names of submetrics and values are 
&'s avatar
& committed
164
            functions that aggregate a list of metric scores
165
166
167
168
169
170
171
172
173
174
175
176
        """
        pass

    @abc.abstractmethod
    def higher_is_better(self):
        """
        :returns: {str: bool}
            A dictionary where keys are the names of submetrics and values are 
            whether a higher value of the submetric is better
        """
        pass

Jason Phang's avatar
Jason Phang committed
177
    def fewshot_description(self):
Jason Phang's avatar
checkin  
Jason Phang committed
178
179
        return ""

Jason Phang's avatar
Jason Phang committed
180
    def fewshot_context(self, doc, num_fewshot, provide_description):
Jason Phang's avatar
Jason Phang committed
181
        raw_description = self.fewshot_description()
Jason Phang's avatar
Jason Phang committed
182
        description = (raw_description + "\n===\n\n") if provide_description and raw_description else ""
183

184
185
186
187
188
189
        if num_fewshot == 0:
            labeled_examples = ""
        else:
            labeled_examples = "\n\n".join(
                [self.doc_to_text(doc) + self.doc_to_target(doc) for doc in self.fewshot_examples(k=num_fewshot)]
            ) + "\n\n"
Leo Gao's avatar
Update  
Leo Gao committed
190

191
        example = self.doc_to_text(doc)
Leo Gao's avatar
Leo Gao committed
192
193
194
        return description + labeled_examples + example


Leo Gao's avatar
Leo Gao committed
195
class MultipleChoiceTask(Task):
Leo Gao's avatar
Leo Gao committed
196
197
198
    def doc_to_target(self, doc):
        return " " + doc['choices'][doc['gold']]

Leo Gao's avatar
Leo Gao committed
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
    def construct_requests(self, doc, ctx):
        lls = [
            rf.loglikelihood(ctx, " {}".format(choice))[0]
            for choice in doc['choices']
        ]

        return lls

    def process_results(self, doc, results):
        gold = doc["gold"]

        acc = 1. if np.argmax(results) == gold else 0.

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


227
req_ret_lens = {
Leo Gao's avatar
Leo Gao committed
228
    'loglikelihood': 2,
Leo Gao's avatar
Leo Gao committed
229
    'greedy_until': None,
230
231
}

Leo Gao's avatar
Leo Gao committed
232
233
234
235
236
import os
import json
import hashlib
from sqlitedict import SqliteDict

Leo Gao's avatar
Leo Gao committed
237
238
239
def hash_args(attr, args):
    dat = json.dumps([attr] + list(args))
    return hashlib.sha256(dat.encode('utf-8')).hexdigest()
Leo Gao's avatar
Leo Gao committed
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255


class CachingLM:
    def __init__(self, lm, cache_db):
        self.lm = lm
        self.cache_db = cache_db
        os.makedirs(os.path.dirname(cache_db), exist_ok=True)
        self.dbdict = SqliteDict(cache_db, autocommit=True)

    def __getattr__(self, attr):
        def fn(requests):
            res = []
            remaining_reqs = []
            
            # figure out which ones are cached and which ones are new
            for req in requests:
Leo Gao's avatar
Leo Gao committed
256
                hsh = hash_args(attr, req)
Leo Gao's avatar
Leo Gao committed
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
                if hsh in self.dbdict:
                    ob = self.dbdict[hsh]

                    assert ob is not None

                    res.append(ob)
                else:
                    res.append(None)
                    remaining_reqs.append(req)
            
            # actually run the LM
            rem_res = getattr(self.lm, attr)(remaining_reqs)

            # stick the new ones back into the list and also cache any of the new ones
            resptr = 0
            for req, r in zip(remaining_reqs, rem_res):
                while res[resptr] is not None: resptr += 1

                res[resptr] = r

                # caching
Leo Gao's avatar
Leo Gao committed
278
                hsh = hash_args(attr, req)
Leo Gao's avatar
Leo Gao committed
279
                self.dbdict[hsh] = r
Leo Gao's avatar
Leo Gao committed
280
            self.dbdict.commit()
Leo Gao's avatar
Leo Gao committed
281
282
283
284

            return res
        return fn

Jason Phang's avatar
Jason Phang committed
285

286
287
288
289
class Request:
    def __init__(self, type, args, index=None):
        if type not in req_ret_lens.keys():
            raise NotImplementedError('The request type {} is not implemented!'.format(type))
Leo Gao's avatar
Leo Gao committed
290

291
292
293
294
295
        self.type = type
        self.args = args
        self.index = index
    
    def __iter__(self):
Leo Gao's avatar
Leo Gao committed
296
297
        if req_ret_lens[self.type] is None:
            raise IndexError('This request type does not return multiple arguments!')
298
299
300
301
302
        i = 0
        for i in range(req_ret_lens[self.type]):
            yield Request(self.type, self.args, i)
    
    def __getitem__(self, i):
Leo Gao's avatar
Leo Gao committed
303
304
        if req_ret_lens[self.type] is None:
            raise IndexError('This request type does not return multiple arguments!')
305
        return Request(self.type, self.args, i)
Leo Gao's avatar
Leo Gao committed
306
307
308
    
    def __eq__(self, other):
        return self.type == other.type and self.args == other.args and self.index == other.index
Leo Gao's avatar
Leo Gao committed
309

Leo Gao's avatar
Leo Gao committed
310
311
    def __repr__(self):
        return f"Req_{self.type}{self.args}[{self.index}]\n"
Jason Phang's avatar
Jason Phang committed
312

Leo Gao's avatar
Leo Gao committed
313
314
class RequestFactory:
    def __getattr__(self, attr):
Leo Gao's avatar
Update  
Leo Gao committed
315
316
        def fn(*args):
            return Request(attr, args)
Leo Gao's avatar
Leo Gao committed
317
318
319
320
        return fn


rf = RequestFactory()