base.py 9.97 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
76
        self._fewshot_docs = None
sdtblck's avatar
sdtblck committed
77
78
79
80
81

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

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

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

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

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

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

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

118
    def fewshot_examples(self, k, rnd):
119
120
        if self._training_docs is None:
            self._training_docs = list(self.training_docs())
121

Leo Gao's avatar
Leo Gao committed
122
        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 ""

180
    def fewshot_context(self, doc, num_fewshot, provide_description, rnd):
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
        if num_fewshot == 0:
            labeled_examples = ""
186
        else:
187
188
189
190
191
192
            # for sets with no training docs, draw from other set *but ensure no overlap with current doc*
            if self.has_training_docs():
                fewshotex = self.fewshot_examples(k=num_fewshot, rnd=rnd)
            else:
                if self._fewshot_docs is None:
                    self._fewshot_docs = list(self.validation_docs() if self.has_validation_docs else self.test_docs())
193

194
                fewshotex = rnd.sample(self._fewshot_docs, num_fewshot + 1)
195

196
197
                # get rid of the doc that's the one we're evaluating, if it's in the fewshot
                fewshotex = [x for x in fewshotex if x != doc][:num_fewshot]
198

199
            labeled_examples = "\n\n".join(
200
                [self.doc_to_text(doc) + self.doc_to_target(doc) for doc in fewshotex]
201
            ) + "\n\n"
Leo Gao's avatar
Update  
Leo Gao committed
202

203
        example = self.doc_to_text(doc)
Leo Gao's avatar
Leo Gao committed
204
205
206
        return description + labeled_examples + example


Leo Gao's avatar
Leo Gao committed
207
class MultipleChoiceTask(Task):
Leo Gao's avatar
Leo Gao committed
208
209
210
    def doc_to_target(self, doc):
        return " " + doc['choices'][doc['gold']]

Leo Gao's avatar
Leo Gao committed
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
    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
        }


239
req_ret_lens = {
Leo Gao's avatar
Leo Gao committed
240
    'loglikelihood': 2,
Leo Gao's avatar
Leo Gao committed
241
    'greedy_until': None,
242
243
}

Leo Gao's avatar
Leo Gao committed
244
245
246
247
248
import os
import json
import hashlib
from sqlitedict import SqliteDict

Leo Gao's avatar
Leo Gao committed
249
250
251
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
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267


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
268
                hsh = hash_args(attr, req)
Leo Gao's avatar
Leo Gao committed
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
                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
290
                hsh = hash_args(attr, req)
Leo Gao's avatar
Leo Gao committed
291
                self.dbdict[hsh] = r
Leo Gao's avatar
Leo Gao committed
292
            self.dbdict.commit()
Leo Gao's avatar
Leo Gao committed
293
294
295
296

            return res
        return fn

Jason Phang's avatar
Jason Phang committed
297

298
299
300
301
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
302

303
304
305
306
307
        self.type = type
        self.args = args
        self.index = index
    
    def __iter__(self):
Leo Gao's avatar
Leo Gao committed
308
309
        if req_ret_lens[self.type] is None:
            raise IndexError('This request type does not return multiple arguments!')
310
311
312
313
314
        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
315
316
        if req_ret_lens[self.type] is None:
            raise IndexError('This request type does not return multiple arguments!')
317
        return Request(self.type, self.args, i)
Leo Gao's avatar
Leo Gao committed
318
319
320
    
    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
321

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

Leo Gao's avatar
Leo Gao committed
325
326
class RequestFactory:
    def __getattr__(self, attr):
Leo Gao's avatar
Update  
Leo Gao committed
327
328
        def fn(*args):
            return Request(attr, args)
Leo Gao's avatar
Leo Gao committed
329
330
331
332
        return fn


rf = RequestFactory()