base.py 32.3 KB
Newer Older
Leo Gao's avatar
Leo Gao committed
1
import abc
2
from typing import Iterable
thefazzer's avatar
thefazzer committed
3
import numpy as np
4
import random
Leo Gao's avatar
Leo Gao committed
5
import re
6
7
8
import os
import json
import hashlib
Jonathan Tow's avatar
Jonathan Tow committed
9
import datasets
10
from sqlitedict import SqliteDict
11
from tqdm import tqdm
12
import torch
Leo Gao's avatar
Leo Gao committed
13
import torch.nn.functional as F
14
from accelerate import find_executable_batch_size
15
import gc
&'s avatar
& committed
16

17
from lm_eval.metrics import mean, weighted_perplexity, weighted_mean, bits_per_byte
18
from lm_eval import utils
19
from abc import abstractmethod
Jason Phang's avatar
gpt3  
Jason Phang committed
20

Jason Phang's avatar
Jason Phang committed
21

Leo Gao's avatar
Leo Gao committed
22
class LM(abc.ABC):
Leo Gao's avatar
Leo Gao committed
23
24
25
    def __init__(self):
        self.cache_hook = CacheHook(None)

26
    @abstractmethod
Leo Gao's avatar
Leo Gao committed
27
    def loglikelihood(self, requests):
Leo Gao's avatar
Leo Gao committed
28
        """Compute log-likelihood of generating a continuation from a context.
Fabrizio Milo's avatar
Fabrizio Milo committed
29
        Downstream tasks should attempt to use loglikelihood instead of other
Leo Gao's avatar
Leo Gao committed
30
        LM calls whenever possible.
Jason Phang's avatar
gpt3  
Jason Phang committed
31

Leo Gao's avatar
Leo Gao committed
32
33
34
        :param requests: list
            A list of pairs (context, continuation)
            context: str
Fabrizio Milo's avatar
Fabrizio Milo committed
35
                Context string. Implementations of LM must be able to handle an
Leo Gao's avatar
Leo Gao committed
36
                empty context string.
Leo Gao's avatar
Leo Gao committed
37
            continuation: str
Fabrizio Milo's avatar
Fabrizio Milo committed
38
39
                The continuation over which log likelihood will be calculated. If
                there is a word boundary, the space should be in the continuation.
Leo Gao's avatar
Leo Gao committed
40
41
42
43
                For example, context="hello" continuation=" world" is correct.
        :return: list
            A list of pairs (logprob, isgreedy)
            logprob: float
Jason Phang's avatar
Jason Phang committed
44
                The log probability of `continuation`
Leo Gao's avatar
Leo Gao committed
45
            isgreedy:
Jason Phang's avatar
Jason Phang committed
46
47
48
49
                Whether `continuation` would be generated by greedy sampling from `context`
        """
        pass

50
    @abstractmethod
Leo Gao's avatar
Leo Gao committed
51
    def loglikelihood_rolling(self, requests):
Jason Phang's avatar
Jason Phang committed
52
53
54
55
        """Compute full log-likelihood of a string, with no truncation, for perplexity computation
        - We will use the full max context length of the model.
        - For inputs that exceed the max context length, we divide the tokenized string into chunks of up to
        the max context length.
Fabrizio Milo's avatar
Fabrizio Milo committed
56
        - IMPORTANT: Each document's loglikelihood/perplexity is computed *separately*, unlike other implementations
Jason Phang's avatar
Jason Phang committed
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
          which may simply concatenate multiple documents together.
        - IMPORTANT: We maximize the amount of context for each prediction. Specifically, for inputs that we break into
          multiple chunks, the last input will still a full-sized context.
          Example:
            Input tokens: [ 0 1 2 3 4 5 6 7 8 9 ]
            Prefix: EOT
            Max context length: 4
            Resulting input/prediction pairs:

                INPUT:  EOT   0   1   2
                PRED:     0   1   2   3

                INPUT:    3   4   5   6
                PRED:     4   5   6   7

                INPUT:    5   6   7   8
                PRED:             8   9

          Observe that:
            1. Each token is predicted exactly once
            2. For the last pair, we provide the full context, but only score the last two tokens

        :param requests: list
            A list of strings
            string: str
                String for which we are computing per-toke  loglikelihood
        :return: list
            A list of pairs (logprob, isgreedy)
            logprob: float
                The log probability of `continuation`
            isgreedy:
                Whether `continuation` would be generated by greedy sampling from `context`
Leo Gao's avatar
Leo Gao committed
89
90
91
        """
        pass

&'s avatar
& committed
92
    # TODO: Add an optional max length
93
    @abstractmethod
Leo Gao's avatar
Update  
Leo Gao committed
94
    def greedy_until(self, requests):
Leo Gao's avatar
Leo Gao committed
95
96
97
98
99
100
        """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
101
            until: [str]
Fabrizio Milo's avatar
Fabrizio Milo committed
102
                The string sequences to generate until. These string sequences
Leo Gao's avatar
Leo Gao committed
103
                may each span across multiple tokens, or may be part of one token.
Leo Gao's avatar
Leo Gao committed
104
105
106
107
        :return: list
            A list of strings continuation
            continuation: str
                The generated continuation.
Jason Phang's avatar
gpt3  
Jason Phang committed
108
        """
Leo Gao's avatar
Leo Gao committed
109
110
        pass

Jason Phang's avatar
gpt3  
Jason Phang committed
111
    @classmethod
112
113
    def create_from_arg_string(cls, arg_string, additional_config=None):
        additional_config = {} if additional_config is None else additional_config
114
115
116
        args = utils.simple_parse_args_string(arg_string)
        args2 = {k: v for k, v in additional_config.items() if v is not None}
        return cls(**args, **args2)
Jason Phang's avatar
gpt3  
Jason Phang committed
117

Leo Gao's avatar
Leo Gao committed
118
119
120
    def set_cache_hook(self, cache_hook):
        self.cache_hook = cache_hook

Leo Gao's avatar
Leo Gao committed
121

122
class BaseLM(LM):
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
    @property
    @abstractmethod
    def eot_token_id(self):
        pass

    @property
    @abstractmethod
    def max_length(self):
        pass

    @property
    @abstractmethod
    def max_gen_toks(self):
        pass

    @property
    @abstractmethod
    def batch_size(self):
        pass

    @property
    @abstractmethod
    def device(self):
        pass

148
    @abstractmethod
Fabrizio Milo's avatar
Fabrizio Milo committed
149
150
151
    def tok_encode(self, string: str):
        pass

152
    @abstractmethod
Fabrizio Milo's avatar
Fabrizio Milo committed
153
154
    def tok_decode(self, tokens: Iterable[int]):
        pass
Jason Phang's avatar
gpt3  
Jason Phang committed
155

156
    @abstractmethod
Fabrizio Milo's avatar
Fabrizio Milo committed
157
158
    def _model_generate(self, context, max_length, eos_token_id):
        pass
Jason Phang's avatar
gpt3  
Jason Phang committed
159

160
161
    @abstractmethod
    def _model_call(self, inps):
Jason Phang's avatar
gpt3  
Jason Phang committed
162
        """
163
164
        inps: a torch tensor of shape [batch, sequence]
        the size of sequence may vary from call to call
Jason Phang's avatar
gpt3  
Jason Phang committed
165

166
        returns: a torch tensor of shape [batch, sequence, vocab] with the
167
        logits returned from the model
168
169
        """
        pass
170

Leo Gao's avatar
Leo Gao committed
171
    # subclass must implement properties vocab_size, eot_token_id, max_gen_toks, batch_size, device, max_length.
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
    # TODO: enforce this somehow

    def loglikelihood(self, requests):
        new_reqs = []
        for context, continuation in requests:
            if context == "":
                # end of text as context
                context_enc = [self.eot_token_id]
            else:
                context_enc = self.tok_encode(context)

            continuation_enc = self.tok_encode(continuation)

            new_reqs.append(((context, continuation), context_enc, continuation_enc))

        return self._loglikelihood_tokens(new_reqs)

    def loglikelihood_rolling(self, requests):
        # TODO: Implement caching once we've confirmed the perplexity implementation
191
192
193
194
195
196
197
198
199
        
        # automatic batch size detection for vectorization
        adaptive_batch_size = None
        if self.batch_size == 'auto': 
            # using rolling window with maximum context
            print('Passed argument batch_size = auto. Detecting largest batch size')
            @find_executable_batch_size(starting_batch_size=512) # if OOM, then halves batch_size and tries again
            def forward_batch(batch_size):
                test_batch = torch.ones((batch_size, self.max_length), device=self.device).long()
200
                out = F.log_softmax(self._model_call(test_batch), dim = -1)
201
202
203
204
205
                return batch_size
            
            batch_size = forward_batch() 
            print(f"Determined Largest batch size: {batch_size}")
            adaptive_batch_size = batch_size
206
207
            torch.cuda.empty_cache()
            gc.collect()
208
209

        loglikelihoods = []
Fabrizio Milo's avatar
Fabrizio Milo committed
210
211
212
213
214
215
216
217
218
219
220
221
        for (string,) in tqdm(requests):
            rolling_token_windows = list(
                map(
                    utils.make_disjoint_window,
                    utils.get_rolling_token_windows(
                        token_list=self.tok_encode(string),
                        prefix_token=self.eot_token_id,
                        max_seq_len=self.max_length,
                        context_len=1,
                    ),
                )
            )
222
223
224

            rolling_token_windows = [(None,) + x for x in rolling_token_windows]

225
226
            # TODO: extract out this call so it only gets called once and also somehow figure out partial caching for
            # that
Fabrizio Milo's avatar
Fabrizio Milo committed
227
            string_nll = self._loglikelihood_tokens(
228
                rolling_token_windows, disable_tqdm=True, override_bs = adaptive_batch_size
Fabrizio Milo's avatar
Fabrizio Milo committed
229
230
            )

231
232
            # discard is_greedy
            string_nll = [x[0] for x in string_nll]
Fabrizio Milo's avatar
Fabrizio Milo committed
233

234
235
236
237
238
            string_nll = sum(string_nll)
            loglikelihoods.append(string_nll)

        return loglikelihoods

239
    def _loglikelihood_tokens(self, requests, disable_tqdm=False, override_bs = None):
240
241
242
243
244
245
        # TODO: implement some kind of efficient-request-middleware that lumps together requests with the same context
        res = []

        def _collate(x):
            # the negative sign on len(toks) sorts descending - this has a few advantages:
            # - time estimates will always be over not underestimates, which is more useful for planning
246
247
248
            # - to know the size of a batch when going through the list, you know the first one is always the batch
            #   padded context length. this is useful to simplify the batching logic and more importantly to make
            #   automatic adaptive batches much much easier to implement
249
250
251
            # - any OOMs will happen right away rather than near the end

            toks = x[1] + x[2]
252
            return -len(toks), tuple(toks)
Fabrizio Milo's avatar
Fabrizio Milo committed
253

254
        
Fabrizio Milo's avatar
Fabrizio Milo committed
255
        re_ord = utils.Reorderer(requests, _collate)
256
257
258

        # automatic (variable) batch size detection for vectorization
        # pull longest context sample from request
259
        _, context_enc, continuation_enc = re_ord.get_reordered()[0]
260
261
        max_context = len((context_enc + continuation_enc)[-(self.max_length + 1) :][:-1])
        if (self.batch_size == 'auto'):
262
            
263
264
265
266
            if override_bs is None:
                print('Passed argument batch_size = auto. Detecting largest batch size')
                @find_executable_batch_size(starting_batch_size=512) # if OOM, then halves batch_size and tries again
                def forward_batch(batch_size):
267
                    test_batch = torch.ones((batch_size, max_context), device=self.device).long()
268
                    out = F.log_softmax(self._model_call(test_batch), dim = -1)
269
270
271
                    return batch_size
                
                batch_size = forward_batch() 
272
                print(f"Determined largest batch size: {batch_size}")
273
                adaptive_batch_size = batch_size
274
275
276
                torch.cuda.empty_cache()
                gc.collect()
                    
277
278
            else:
                adaptive_batch_size = override_bs
279

Fabrizio Milo's avatar
Fabrizio Milo committed
280
        for chunk in utils.chunks(
281
            tqdm(re_ord.get_reordered(), disable=disable_tqdm), self.batch_size if self.batch_size != "auto" else adaptive_batch_size
Fabrizio Milo's avatar
Fabrizio Milo committed
282
        ):
283
            inps = []
284
            cont_toks_list = []
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
            inplens = []

            padding_length = None

            # because vectorizing is annoying, we first convert each (context, continuation) pair to padded
            # tensors, then we pack them together into a batch, call the model, and then pick it all apart
            # again because vectorizing is annoying

            for _, context_enc, continuation_enc in chunk:
                # sanity check
                assert len(context_enc) > 0
                assert len(continuation_enc) > 0
                assert len(continuation_enc) <= self.max_length

                # how this all works:
                #          CTX      CONT
301
                # inp    0 1 2 3|4 5 6 7 8 9   <- last token is deleted by inp[:, :-1]
302
                # gpt2    \               \
303
304
                # logits   1 2 3|4 5 6 7 8 9   <- the ctx half gets tossed out by the
                # cont_toks      4 5 6 7 8 9      [:, -len(continuation_enc):, :self.vocab_size] slice
305
306
307

                # when too long to fit in context, truncate from the left
                inp = torch.tensor(
Fabrizio Milo's avatar
Fabrizio Milo committed
308
309
                    (context_enc + continuation_enc)[-(self.max_length + 1) :][:-1],
                    dtype=torch.long,
310
                ).to(self.device)
Fabrizio Milo's avatar
Fabrizio Milo committed
311
                (inplen,) = inp.shape
312
313
314
315

                cont = continuation_enc

                # since in _collate we make sure length is descending, the longest is always the first one.
Fabrizio Milo's avatar
Fabrizio Milo committed
316
317
318
                padding_length = (
                    padding_length if padding_length is not None else inplen
                )
319

320
                # pad length from seq to padding_length
Fabrizio Milo's avatar
Fabrizio Milo committed
321
322
323
324
325
326
327
328
329
                inp = torch.cat(
                    [
                        inp,  # [seq]
                        torch.zeros(padding_length - inplen, dtype=torch.long).to(
                            inp.device
                        ),  # [padding_length - seq]
                    ],
                    dim=0,
                )
330

331
332
                inps.append(inp.unsqueeze(0))  # [1, padding_length]
                cont_toks_list.append(cont)
333
334
                inplens.append(inplen)

335
            batched_inps = torch.cat(inps, dim=0)  # [batch, padding_length
Fabrizio Milo's avatar
Fabrizio Milo committed
336
337
338
            multi_logits = F.log_softmax(
                self._model_call(batched_inps), dim=-1
            ).cpu()  # [batch, padding_length, vocab]
339

Fabrizio Milo's avatar
Fabrizio Milo committed
340
341
342
            for (cache_key, _, _), logits, inp, inplen, cont_toks in zip(
                chunk, multi_logits, inps, inplens, cont_toks_list
            ):
343

344
345
                # Slice to original seq length
                contlen = len(cont_toks)
Fabrizio Milo's avatar
Fabrizio Milo committed
346
347
348
                logits = logits[inplen - contlen : inplen].unsqueeze(
                    0
                )  # [1, seq, vocab]
349

350
                # Check if per-token argmax is exactly equal to continuation
351
                greedy_tokens = logits.argmax(dim=-1)
Fabrizio Milo's avatar
Fabrizio Milo committed
352
353
354
                cont_toks = torch.tensor(cont_toks, dtype=torch.long).unsqueeze(
                    0
                )  # [1, seq]
355
356
                max_equal = (greedy_tokens == cont_toks).all()

357
358
                # Obtain log-probs at the corresponding continuation token indices
                # last_token_slice = logits[:, -1, :].squeeze(0).tolist()
Fabrizio Milo's avatar
Fabrizio Milo committed
359
360
361
                logits = torch.gather(logits, 2, cont_toks.unsqueeze(-1)).squeeze(
                    -1
                )  # [1, seq]
362

363
                # Answer: (log prob, is-exact-match)
364
365
366
367
368
369
370
371
                answer = (float(logits.sum()), bool(max_equal))

                # partial caching
                if cache_key is not None:
                    self.cache_hook.add_partial("loglikelihood", cache_key, answer)

                res.append(answer)

Fabrizio Milo's avatar
Fabrizio Milo committed
372
        return re_ord.get_original(res)
Fabrizio Milo's avatar
Fabrizio Milo committed
373

374
    def greedy_until(self, requests):
Fabrizio Milo's avatar
Fabrizio Milo committed
375
        # TODO: implement fully general `until` that handles until that are
376
        #       multiple tokens or that span multiple tokens correctly
377
378
379
380
381
382

        # TODO: extract to TokenizedLM?
        res = []

        def _collate(x):
            toks = self.tok_encode(x[0])
383
            return len(toks), x[0]
Fabrizio Milo's avatar
Fabrizio Milo committed
384

Fabrizio Milo's avatar
Fabrizio Milo committed
385
        re_ord = utils.Reorderer(requests, _collate)
386

Fabrizio Milo's avatar
Fabrizio Milo committed
387
        for context, until in tqdm(re_ord.get_reordered()):
388
389
            if isinstance(until, str):
                until = [until]
390

Fabrizio Milo's avatar
Fabrizio Milo committed
391
            (primary_until,) = self.tok_encode(until[0])
392

Fabrizio Milo's avatar
Fabrizio Milo committed
393
394
395
            context_enc = torch.tensor(
                [self.tok_encode(context)[self.max_gen_toks - self.max_length :]]
            ).to(self.device)
396

Fabrizio Milo's avatar
Fabrizio Milo committed
397
398
399
400
401
            cont = self._model_generate(
                context_enc, context_enc.shape[1] + self.max_gen_toks, primary_until
            )

            s = self.tok_decode(cont[0].tolist()[context_enc.shape[1] :])
402
403
404

            for term in until:
                s = s.split(term)[0]
Fabrizio Milo's avatar
Fabrizio Milo committed
405

406
407
            # partial caching
            self.cache_hook.add_partial("greedy_until", (context, until), s)
Fabrizio Milo's avatar
Fabrizio Milo committed
408

409
            res.append(s)
Fabrizio Milo's avatar
Fabrizio Milo committed
410

Fabrizio Milo's avatar
Fabrizio Milo committed
411
        return re_ord.get_original(res)
Leo Gao's avatar
Leo Gao committed
412

Leo Gao's avatar
Leo Gao committed
413

414
class Task(abc.ABC):
&'s avatar
&amp; committed
415
416
417
418
419
420
421
422
    """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)
    """
Jonathan Tow's avatar
Jonathan Tow committed
423

Jon Tow's avatar
Jon Tow committed
424
425
    # The name of the `Task` benchmark as denoted in the HuggingFace datasets Hub
    # or a path to a custom `datasets` loading script.
Jonathan Tow's avatar
Jonathan Tow committed
426
427
428
429
430
    DATASET_PATH: str = None

    # The name of a subset within `DATASET_PATH`.
    DATASET_NAME: str = None

Jon Tow's avatar
Jon Tow committed
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
    def __init__(self, data_dir=None, cache_dir=None, download_mode=None):
        """
        :param data_dir: str
            Stores the path to a local folder containing the `Task`'s data files.
            Use this to specify the path to manually downloaded data (usually when
            the dataset is not publicly accessible).
        :param cache_dir: str
            The directory to read/write the `Task` dataset. This follows the
            HuggingFace `datasets` API with the default cache directory located at:
                `~/.cache/huggingface/datasets`
            NOTE: You can change the cache location globally for a given process
            by setting the shell environment variable, `HF_DATASETS_CACHE`,
            to another directory:
                `export HF_DATASETS_CACHE="/path/to/another/directory"`
        :param download_mode: datasets.DownloadMode
            How to treat pre-existing `Task` downloads and data.
            - `datasets.DownloadMode.REUSE_DATASET_IF_EXISTS`
                Reuse download and reuse dataset.
            - `datasets.DownloadMode.REUSE_CACHE_IF_EXISTS`
                Reuse download with fresh dataset.
            - `datasets.DownloadMode.FORCE_REDOWNLOAD`
                Fresh download and fresh dataset.
        """
        self.download(data_dir, cache_dir, download_mode)
455
        self._training_docs = None
456
        self._fewshot_docs = None
sdtblck's avatar
sdtblck committed
457

Jon Tow's avatar
Jon Tow committed
458
    def download(self, data_dir=None, cache_dir=None, download_mode=None):
Fabrizio Milo's avatar
Fabrizio Milo committed
459
        """Downloads and returns the task dataset.
Jonathan Tow's avatar
Jonathan Tow committed
460
461
        Override this method to download the dataset from a custom API.

Jon Tow's avatar
Jon Tow committed
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
        :param data_dir: str
            Stores the path to a local folder containing the `Task`'s data files.
            Use this to specify the path to manually downloaded data (usually when
            the dataset is not publicly accessible).
        :param cache_dir: str
            The directory to read/write the `Task` dataset. This follows the
            HuggingFace `datasets` API with the default cache directory located at:
                `~/.cache/huggingface/datasets`
            NOTE: You can change the cache location globally for a given process
            by setting the shell environment variable, `HF_DATASETS_CACHE`,
            to another directory:
                `export HF_DATASETS_CACHE="/path/to/another/directory"`
        :param download_mode: datasets.DownloadMode
            How to treat pre-existing `Task` downloads and data.
            - `datasets.DownloadMode.REUSE_DATASET_IF_EXISTS`
                Reuse download and reuse dataset.
            - `datasets.DownloadMode.REUSE_CACHE_IF_EXISTS`
                Reuse download with fresh dataset.
            - `datasets.DownloadMode.FORCE_REDOWNLOAD`
                Fresh download and fresh dataset.
Jonathan Tow's avatar
Jonathan Tow committed
482
483
484
485
        """
        self.dataset = datasets.load_dataset(
            path=self.DATASET_PATH,
            name=self.DATASET_NAME,
Jon Tow's avatar
Jon Tow committed
486
487
            data_dir=data_dir,
            cache_dir=cache_dir,
Fabrizio Milo's avatar
Fabrizio Milo committed
488
            download_mode=download_mode,
Jonathan Tow's avatar
Jonathan Tow committed
489
        )
sdtblck's avatar
sdtblck committed
490

491
492
493
494
    def should_decontaminate(self):
        """Whether this task supports decontamination against model training set."""
        return False

495
    @abstractmethod
496
    def has_training_docs(self):
Jason Phang's avatar
checkin  
Jason Phang committed
497
        """Whether the task has a training set"""
498
        pass
499

500
    @abstractmethod
501
    def has_validation_docs(self):
Jason Phang's avatar
checkin  
Jason Phang committed
502
503
504
        """Whether the task has a validation set"""
        pass

505
    @abstractmethod
Jason Phang's avatar
checkin  
Jason Phang committed
506
507
    def has_test_docs(self):
        """Whether the task has a test set"""
508
509
        pass

Leo Gao's avatar
Leo Gao committed
510
    def training_docs(self):
Jason Phang's avatar
checkin  
Jason Phang committed
511
512
513
514
        """
        :return: Iterable[obj]
            A iterable of any object, that doc_to_text can handle
        """
Leo Gao's avatar
Leo Gao committed
515
        return []
516

Leo Gao's avatar
Leo Gao committed
517
    def validation_docs(self):
518
519
520
521
        """
        :return: Iterable[obj]
            A iterable of any object, that doc_to_text can handle
        """
Leo Gao's avatar
Leo Gao committed
522
        return []
523

Leo Gao's avatar
Leo Gao committed
524
    def test_docs(self):
525
526
527
528
        """
        :return: Iterable[obj]
            A iterable of any object, that doc_to_text can handle
        """
Leo Gao's avatar
Leo Gao committed
529
        return []
Leo Gao's avatar
Leo Gao committed
530

Jon Tow's avatar
Jon Tow committed
531
532
533
534
    def _process_doc(self, doc):
        """
        Override this to process (detokenize, strip, replace, etc.) individual
        documents. This can be used in a map over documents of a data split.
Jon Tow's avatar
Jon Tow committed
535
        E.g. `map(self._process_doc, self.dataset["validation"])`
Jon Tow's avatar
Jon Tow committed
536
537
538
539
540
541

        :return: dict
            The processed version of the specified `doc`.
        """
        return doc

542
    def fewshot_examples(self, k, rnd):
543
544
        if self._training_docs is None:
            self._training_docs = list(self.training_docs())
545

Leo Gao's avatar
Leo Gao committed
546
        return rnd.sample(self._training_docs, k)
Leo Gao's avatar
Leo Gao committed
547

548
    def doc_to_decontamination_query(self, doc):
Fabrizio Milo's avatar
Fabrizio Milo committed
549
550
551
552
        print(
            "Override doc_to_decontamination_query with document specific decontamination query."
        )
        assert False
553

554
    @abstractmethod
Leo Gao's avatar
Update  
Leo Gao committed
555
556
557
    def doc_to_text(self, doc):
        pass

558
    @abstractmethod
Leo Gao's avatar
Update  
Leo Gao committed
559
    def doc_to_target(self, doc):
Leo Gao's avatar
Leo Gao committed
560
        pass
Leo Gao's avatar
Leo Gao committed
561

562
    @abstractmethod
563
    def construct_requests(self, doc, ctx):
Fabrizio Milo's avatar
Fabrizio Milo committed
564
        """Uses RequestFactory to construct Requests and returns an iterable of
Leo Gao's avatar
Leo Gao committed
565
566
        Requests which will be sent to the LM.

567
568
        :param doc:
            The document as returned from training_docs, validation_docs, or test_docs.
Leo Gao's avatar
Leo Gao committed
569
        :param ctx: str
Fabrizio Milo's avatar
Fabrizio Milo committed
570
            The context string, generated by fewshot_context. This includes the natural
571
            language description, as well as the few shot examples, and the question
Fabrizio Milo's avatar
Fabrizio Milo committed
572
            part of the document for `doc`.
Leo Gao's avatar
Leo Gao committed
573
        """
Leo Gao's avatar
Leo Gao committed
574
        pass
575

576
    @abstractmethod
Leo Gao's avatar
Leo Gao committed
577
    def process_results(self, doc, results):
Fabrizio Milo's avatar
Fabrizio Milo committed
578
579
        """Take a single document and the LM results and evaluates, returning a
        dict where keys are the names of submetrics and values are the values of
580
        the metric for that one document
Leo Gao's avatar
Leo Gao committed
581
582
583
584
585

        :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
586
        """
Leo Gao's avatar
Leo Gao committed
587
        pass
Jason Phang's avatar
gpt3  
Jason Phang committed
588

589
    @abstractmethod
590
591
    def aggregation(self):
        """
&'s avatar
&amp; committed
592
        :returns: {str: [metric_score] -> float}
Fabrizio Milo's avatar
Fabrizio Milo committed
593
            A dictionary where keys are the names of submetrics and values are
&'s avatar
&amp; committed
594
            functions that aggregate a list of metric scores
595
596
597
        """
        pass

598
    @abstractmethod
599
600
601
    def higher_is_better(self):
        """
        :returns: {str: bool}
Fabrizio Milo's avatar
Fabrizio Milo committed
602
            A dictionary where keys are the names of submetrics and values are
603
604
605
606
            whether a higher value of the submetric is better
        """
        pass

Jason Phang's avatar
Jason Phang committed
607
    def fewshot_description(self):
608
        import warnings
Fabrizio Milo's avatar
Fabrizio Milo committed
609

610
        warnings.warn(
Jonathan Tow's avatar
Jonathan Tow committed
611
            "`fewshot_description` will be removed in futures versions. Pass "
612
            "any custom descriptions to the `evaluate` function instead.",
Fabrizio Milo's avatar
Fabrizio Milo committed
613
614
            DeprecationWarning,
        )
Jason Phang's avatar
checkin  
Jason Phang committed
615
616
        return ""

617
    @utils.positional_deprecated
Fabrizio Milo's avatar
Fabrizio Milo committed
618
619
620
621
    def fewshot_context(
        self, doc, num_fewshot, provide_description=None, rnd=None, description=None
    ):
        """Returns a fewshot context string that is made up of a prepended description
622
623
624
625
626
627
628
629
630
631
        (if provided), the `num_fewshot` number of examples, and an appended prompt example.

        :param doc: str
            The document as returned from training_docs, validation_docs, or test_docs.
        :param num_fewshot: int
            The number of fewshot examples to provide in the returned context string.
        :param provide_description: bool
            Not implemented, and this option is deprecated and will be removed in a future version in favor of a different description providing method
        :param rnd: random.Random
            The pseudo-random number generator used to randomly sample examples.
632
            WARNING: This is currently a required arg although it's optionalized with a default `None`.
633
634
635
636
637
        :param description: str
            The task's description that will be prepended to the fewshot examples.
        :returns: str
            The fewshot context.
        """
Fabrizio Milo's avatar
Fabrizio Milo committed
638
639
640
        assert (
            rnd is not None
        ), "A `random.Random` generator argument must be provided to `rnd`"
Jonathan Tow's avatar
Merge  
Jonathan Tow committed
641
        assert not provide_description, (
Jonathan Tow's avatar
Jonathan Tow committed
642
            "The `provide_description` arg will be removed in future versions. To prepend "
643
            "a custom description to the context, supply the corresponding string via the "
Jonathan Tow's avatar
Jonathan Tow committed
644
            "`description` arg."
Jonathan Tow's avatar
Merge  
Jonathan Tow committed
645
        )
646
647
        if provide_description is not None:
            # nudge people to not specify it at all
Fabrizio Milo's avatar
Fabrizio Milo committed
648
649
650
            print(
                "WARNING: provide_description is deprecated and will be removed in a future version in favor of description_dict"
            )
651

652
        description = description + "\n\n" if description else ""
653

654
655
        if num_fewshot == 0:
            labeled_examples = ""
656
        else:
657
658
659
660
661
            # 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:
662
                    self._fewshot_docs = list(
Fabrizio Milo's avatar
Fabrizio Milo committed
663
664
665
                        self.validation_docs()
                        if self.has_validation_docs()
                        else self.test_docs()
666
                    )
667

668
                fewshotex = rnd.sample(self._fewshot_docs, num_fewshot + 1)
669

670
671
                # 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]
672

Fabrizio Milo's avatar
Fabrizio Milo committed
673
674
675
676
677
678
679
680
681
            labeled_examples = (
                "\n\n".join(
                    [
                        self.doc_to_text(doc) + self.doc_to_target(doc)
                        for doc in fewshotex
                    ]
                )
                + "\n\n"
            )
Leo Gao's avatar
Update  
Leo Gao committed
682

683
        example = self.doc_to_text(doc)
Leo Gao's avatar
Leo Gao committed
684
685
686
        return description + labeled_examples + example


Jon Tow's avatar
Jon Tow committed
687
class MultipleChoiceTask(Task):
Leo Gao's avatar
Leo Gao committed
688
    def doc_to_target(self, doc):
Fabrizio Milo's avatar
Fabrizio Milo committed
689
        return " " + doc["choices"][doc["gold"]]
Leo Gao's avatar
Leo Gao committed
690

Leo Gao's avatar
Leo Gao committed
691
692
    def construct_requests(self, doc, ctx):
        lls = [
Fabrizio Milo's avatar
Fabrizio Milo committed
693
            rf.loglikelihood(ctx, " {}".format(choice))[0] for choice in doc["choices"]
Leo Gao's avatar
Leo Gao committed
694
695
696
697
698
699
700
        ]

        return lls

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

Fabrizio Milo's avatar
Fabrizio Milo committed
701
        acc = 1.0 if np.argmax(results) == gold else 0.0
702
        completion_len = np.array([float(len(i)) for i in doc["choices"]])
Fabrizio Milo's avatar
Fabrizio Milo committed
703
        acc_norm = 1.0 if np.argmax(results / completion_len) == gold else 0.0
Leo Gao's avatar
Leo Gao committed
704
705

        return {
Leo Gao's avatar
Leo Gao committed
706
707
            "acc": acc,
            "acc_norm": acc_norm,
Leo Gao's avatar
Leo Gao committed
708
        }
Fabrizio Milo's avatar
Fabrizio Milo committed
709

Leo Gao's avatar
Leo Gao committed
710
711
    def higher_is_better(self):
        return {
Leo Gao's avatar
Leo Gao committed
712
713
            "acc": True,
            "acc_norm": True,
Leo Gao's avatar
Leo Gao committed
714
        }
Fabrizio Milo's avatar
Fabrizio Milo committed
715

Leo Gao's avatar
Leo Gao committed
716
717
    def aggregation(self):
        return {
Leo Gao's avatar
Leo Gao committed
718
719
            "acc": mean,
            "acc_norm": mean,
Leo Gao's avatar
Leo Gao committed
720
721
722
        }


Jason Phang's avatar
Jason Phang committed
723
class PerplexityTask(Task, abc.ABC):
724
725
726
727
    def should_decontaminate(self):
        """Whether this task supports decontamination against model training set."""
        return True

Jason Phang's avatar
Jason Phang committed
728
729
730
731
732
733
734
    def has_training_docs(self):
        return False

    def fewshot_examples(self, k, rnd):
        assert k == 0
        return []

Fabrizio Milo's avatar
Fabrizio Milo committed
735
736
737
738
739
740
741
742
743
    def fewshot_context(
        self, doc, num_fewshot, provide_description=None, rnd=None, description=None
    ):
        assert (
            num_fewshot == 0
        ), "The number of fewshot examples must be 0 for perplexity tasks."
        assert (
            rnd is not None
        ), "A `random.Random` generator argument must be provided to `rnd`."
Jonathan Tow's avatar
Merge  
Jonathan Tow committed
744
        assert not provide_description, (
Jonathan Tow's avatar
Jonathan Tow committed
745
            "The `provide_description` arg will be removed in future versions. To prepend "
Jonathan Tow's avatar
Jonathan Tow committed
746
            "a custom description to the context, supply the corresponding string via the "
Jonathan Tow's avatar
Jonathan Tow committed
747
            "`description` arg."
Jonathan Tow's avatar
Merge  
Jonathan Tow committed
748
        )
749
750
        if provide_description is not None:
            # nudge people to not specify it at all
Fabrizio Milo's avatar
Fabrizio Milo committed
751
752
753
            print(
                "WARNING: provide_description is deprecated and will be removed in a future version in favor of description_dict"
            )
754

Jason Phang's avatar
Jason Phang committed
755
756
757
        return ""

    def higher_is_better(self):
Leo Gao's avatar
Leo Gao committed
758
759
760
761
762
        return {
            "word_perplexity": False,
            "byte_perplexity": False,
            "bits_per_byte": False,
        }
Jason Phang's avatar
Jason Phang committed
763

764
765
766
    def doc_to_decontamination_query(self, doc):
        return doc

Jason Phang's avatar
Jason Phang committed
767
    def doc_to_text(self, doc):
768
        return ""
Jason Phang's avatar
Jason Phang committed
769
770

    def doc_to_target(self, doc):
771
        return doc
Jason Phang's avatar
Jason Phang committed
772
773
774

    def construct_requests(self, doc, ctx):
        assert not ctx
Leo Gao's avatar
Leo Gao committed
775
        req = rf.loglikelihood_rolling(self.doc_to_target(doc))
Jason Phang's avatar
Jason Phang committed
776
777
778
        return req

    def process_results(self, doc, results):
Fabrizio Milo's avatar
Fabrizio Milo committed
779
        (loglikelihood,) = results
Leo Gao's avatar
Leo Gao committed
780
        words = self.count_words(doc)
781
        bytes_ = self.count_bytes(doc)
Jason Phang's avatar
Jason Phang committed
782
        return {
Leo Gao's avatar
Leo Gao committed
783
            "word_perplexity": (loglikelihood, words),
784
            "byte_perplexity": (loglikelihood, bytes_),
785
            "bits_per_byte": (loglikelihood, bytes_),
Jason Phang's avatar
Jason Phang committed
786
787
788
789
        }

    def aggregation(self):
        return {
Leo Gao's avatar
Leo Gao committed
790
791
            "word_perplexity": weighted_perplexity,
            "byte_perplexity": weighted_perplexity,
792
            "bits_per_byte": bits_per_byte,
Jason Phang's avatar
Jason Phang committed
793
794
        }

795
796
    @classmethod
    def count_bytes(cls, doc):
Leo Gao's avatar
Leo Gao committed
797
        return len(doc.encode("utf-8"))
798
799
800

    @classmethod
    def count_words(cls, doc):
Fabrizio Milo's avatar
Fabrizio Milo committed
801
        """Downstream tasks with custom word boundaries should override this!"""
Leo Gao's avatar
Leo Gao committed
802
        return len(re.split(r"\s+", doc))
Leo Gao's avatar
Leo Gao committed
803

Jason Phang's avatar
Jason Phang committed
804

Leo Gao's avatar
Leo Gao committed
805
806
def hash_args(attr, args):
    dat = json.dumps([attr] + list(args))
Fabrizio Milo's avatar
Fabrizio Milo committed
807
    return hashlib.sha256(dat.encode("utf-8")).hexdigest()
Leo Gao's avatar
Leo Gao committed
808
809


Leo Gao's avatar
Leo Gao committed
810
811
class CacheHook:
    def __init__(self, cachinglm):
Fabrizio Milo's avatar
Fabrizio Milo committed
812
        if cachinglm is None:
Leo Gao's avatar
Leo Gao committed
813
814
815
816
            self.dbdict = None
            return

        self.dbdict = cachinglm.dbdict
Fabrizio Milo's avatar
Fabrizio Milo committed
817

Leo Gao's avatar
Leo Gao committed
818
819
820
821
822
823
824
    def add_partial(self, attr, req, res):
        if self.dbdict is None:
            return
        hsh = hash_args(attr, req)
        self.dbdict[hsh] = res


Leo Gao's avatar
Leo Gao committed
825
826
class CachingLM:
    def __init__(self, lm, cache_db):
827
828
829
830
831
832
833
        """LM wrapper that returns cached results if they exist, and uses the underlying LM if not.

        :param lm: LM
            Underlying LM
        :param cache_db: str
            Path to cache db
        """
Leo Gao's avatar
Leo Gao committed
834
835
        self.lm = lm
        self.cache_db = cache_db
836
837
        if os.path.dirname(cache_db):
            os.makedirs(os.path.dirname(cache_db), exist_ok=True)
Leo Gao's avatar
Leo Gao committed
838
839
        self.dbdict = SqliteDict(cache_db, autocommit=True)

Leo Gao's avatar
Leo Gao committed
840
841
842
        # add hook to lm
        lm.set_cache_hook(self.get_cache_hook())

Leo Gao's avatar
Leo Gao committed
843
844
845
846
    def __getattr__(self, attr):
        def fn(requests):
            res = []
            remaining_reqs = []
Fabrizio Milo's avatar
Fabrizio Milo committed
847

Leo Gao's avatar
Leo Gao committed
848
849
            # figure out which ones are cached and which ones are new
            for req in requests:
Leo Gao's avatar
Leo Gao committed
850
                hsh = hash_args(attr, req)
Leo Gao's avatar
Leo Gao committed
851
852
853
854
855
856
857
858
859
                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)
Fabrizio Milo's avatar
Fabrizio Milo committed
860

861
            # actually run the LM on the requests that do not have cached results
Leo Gao's avatar
Leo Gao committed
862
863
864
865
866
            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):
867
868
                while res[resptr] is not None:
                    resptr += 1
Leo Gao's avatar
Leo Gao committed
869
870
871
872

                res[resptr] = r

                # caching
Leo Gao's avatar
Leo Gao committed
873
                hsh = hash_args(attr, req)
Leo Gao's avatar
Leo Gao committed
874
                self.dbdict[hsh] = r
Leo Gao's avatar
Leo Gao committed
875
            self.dbdict.commit()
Leo Gao's avatar
Leo Gao committed
876
877

            return res
Fabrizio Milo's avatar
Fabrizio Milo committed
878

Leo Gao's avatar
Leo Gao committed
879
        return fn
Fabrizio Milo's avatar
Fabrizio Milo committed
880

Leo Gao's avatar
Leo Gao committed
881
882
    def get_cache_hook(self):
        return CacheHook(self)
Leo Gao's avatar
Leo Gao committed
883

Jason Phang's avatar
Jason Phang committed
884

885
REQUEST_RETURN_LENGTHS = {
Fabrizio Milo's avatar
Fabrizio Milo committed
886
887
888
    "loglikelihood": 2,
    "greedy_until": None,
    "loglikelihood_rolling": None,
889
890
891
}


892
class Request:
Leo Gao's avatar
Leo Gao committed
893
894
    def __init__(self, request_type, args, index=None):
        if request_type not in REQUEST_RETURN_LENGTHS.keys():
Fabrizio Milo's avatar
Fabrizio Milo committed
895
896
897
            raise NotImplementedError(
                "The request type {} is not implemented!".format(request_type)
            )
Leo Gao's avatar
Leo Gao committed
898

Leo Gao's avatar
Leo Gao committed
899
        self.request_type = request_type
900
901
        self.args = args
        self.index = index
Fabrizio Milo's avatar
Fabrizio Milo committed
902

903
    def __iter__(self):
Leo Gao's avatar
Leo Gao committed
904
        if REQUEST_RETURN_LENGTHS[self.request_type] is None:
Fabrizio Milo's avatar
Fabrizio Milo committed
905
            raise IndexError("This request type does not return multiple arguments!")
Leo Gao's avatar
Leo Gao committed
906
907
        for i in range(REQUEST_RETURN_LENGTHS[self.request_type]):
            yield Request(self.request_type, self.args, i)
Fabrizio Milo's avatar
Fabrizio Milo committed
908

909
    def __getitem__(self, i):
Leo Gao's avatar
Leo Gao committed
910
        if REQUEST_RETURN_LENGTHS[self.request_type] is None:
Fabrizio Milo's avatar
Fabrizio Milo committed
911
            raise IndexError("This request type does not return multiple arguments!")
Leo Gao's avatar
Leo Gao committed
912
        return Request(self.request_type, self.args, i)
Fabrizio Milo's avatar
Fabrizio Milo committed
913

Leo Gao's avatar
Leo Gao committed
914
    def __eq__(self, other):
Fabrizio Milo's avatar
Fabrizio Milo committed
915
916
917
918
919
        return (
            self.request_type == other.request_type
            and self.args == other.args
            and self.index == other.index
        )
Leo Gao's avatar
Leo Gao committed
920

Leo Gao's avatar
Leo Gao committed
921
    def __repr__(self):
Leo Gao's avatar
Leo Gao committed
922
        return f"Req_{self.request_type}{self.args}[{self.index}]\n"
923

Jason Phang's avatar
Jason Phang committed
924

Leo Gao's avatar
Leo Gao committed
925
926
class RequestFactory:
    def __getattr__(self, attr):
Leo Gao's avatar
Update  
Leo Gao committed
927
928
        def fn(*args):
            return Request(attr, args)
Fabrizio Milo's avatar
Fabrizio Milo committed
929

Leo Gao's avatar
Leo Gao committed
930
931
932
933
        return fn


rf = RequestFactory()