lambada.py 1.75 KB
Newer Older
1
import json
&'s avatar
& committed
2
3
from lm_eval.base import Task, rf
from lm_eval.metrics import mean, perplexity
sdtblck's avatar
sdtblck committed
4
from lm_eval.utils import sh
Leo Gao's avatar
Leo Gao committed
5
from best_download import download_file
sdtblck's avatar
sdtblck committed
6
7


8
class LAMBADA(Task):
Leo Gao's avatar
Leo Gao committed
9
    VERSION = 0
sdtblck's avatar
sdtblck committed
10
11
    def download(self):
        sh("mkdir -p data/lambada")
12
13
14
15
16
17
        sh("wget http://eaidata.bmk.sh/data/lambada_test.jsonl -O data/lambada/lambada_test.jsonl")
#         download_file(
#             "http://eaidata.bmk.sh/data/lambada_test.jsonl", 
#             "data/lambada/lambada_test.jsonl", 
#             "4aa8d02cd17c719165fc8a7887fddd641f43fcafa4b1c806ca8abc31fabdb226"
#         )
sdtblck's avatar
sdtblck committed
18
19
20
21
22

    def has_training_docs(self):
        return False

    def has_validation_docs(self):
Leo Gao's avatar
Leo Gao committed
23
        return True
sdtblck's avatar
sdtblck committed
24
25

    def has_test_docs(self):
Leo Gao's avatar
Leo Gao committed
26
        return False
sdtblck's avatar
sdtblck committed
27
28
29
30
31

    def training_docs(self):
        pass

    def validation_docs(self):
Leo Gao's avatar
Leo Gao committed
32
33
34
35
        with open("data/lambada/lambada_test.jsonl") as fh:
            for line in fh:
                yield json.loads(line)

Leo Gao's avatar
Leo Gao committed
36
37
38
    def test_docs(self):
        pass

Leo Gao's avatar
Leo Gao committed
39
40
    def doc_to_text(self, doc):
        return doc['text'].rsplit(' ', 1)[0]
sdtblck's avatar
sdtblck committed
41

Leo Gao's avatar
Leo Gao committed
42
43
    def doc_to_target(self, doc):
        return " " + doc['text'].rsplit(' ', 1)[1]
Leo Gao's avatar
Leo Gao committed
44
45
46
47
    
    def fewshot_description(self):
        # TODO: figure out description
        return ""
sdtblck's avatar
sdtblck committed
48

Leo Gao's avatar
Leo Gao committed
49
    def construct_requests(self, doc, ctx):
Leo Gao's avatar
Leo Gao committed
50
        ll, is_greedy = rf.loglikelihood(ctx, self.doc_to_target(doc))
51

Leo Gao's avatar
Leo Gao committed
52
        return ll, is_greedy
Leo Gao's avatar
Leo Gao committed
53
54
    
    def process_results(self, doc, results):
Leo Gao's avatar
Leo Gao committed
55
        ll, is_greedy = results
Leo Gao's avatar
Leo Gao committed
56

Leo Gao's avatar
Leo Gao committed
57
        return {
Leo Gao's avatar
Leo Gao committed
58
59
            'ppl': ll,
            'acc': int(is_greedy)
Leo Gao's avatar
Leo Gao committed
60
61
        }
        
Leo Gao's avatar
Leo Gao committed
62
    def aggregation(self):
Leo Gao's avatar
Leo Gao committed
63
        return {
Leo Gao's avatar
Leo Gao committed
64
65
            'ppl': perplexity,
            'acc': mean
Leo Gao's avatar
Leo Gao committed
66
        }
Leo Gao's avatar
Leo Gao committed
67
68

    def higher_is_better(self):
Leo Gao's avatar
Leo Gao committed
69
        return {
Leo Gao's avatar
Leo Gao committed
70
71
            'ppl': False,
            'acc': True
Leo Gao's avatar
Leo Gao committed
72
        }