pubmedqa.py 3.98 KB
Newer Older
jeffhsu3's avatar
jeffhsu3 committed
1
import numpy as np
jeffhsu3's avatar
jeffhsu3 committed
2
import json
3
4
5
import random
from .common import HFTask 
from lm_eval.base import rf, mean
jeffhsu3's avatar
jeffhsu3 committed
6
7
8
9
10
11
12
13
14
15
16
17


class Pubmed_QA(HFTask):
    DATASET_PATH = "pubmed_qa"
    DATASET_NAME = "pqa_labeled"

    def has_training_docs(self):
        return False

    def has_validation_docs(self):
        return False

jeffhsu3's avatar
jeffhsu3 committed
18
19
20
    def has_test_docs(self):
        return True

21
22
23
24
25
    def test_docs(self):
        if self.has_test_docs():
            # HF is labelled as train but its really just for testing
            return self.data["train"]

jeffhsu3's avatar
jeffhsu3 committed
26
27
    def fewshot_description(self):
        # Average ctx length in labelled dataset is 238.9
jeffhsu3's avatar
jeffhsu3 committed
28
        # 2 few-shot exmamples pushes it beyond context window
jeffhsu3's avatar
jeffhsu3 committed
29
30
31
        return ""

    def doc_to_text(self, doc):
jeffhsu3's avatar
jeffhsu3 committed
32
        ctxs = "\n".join(doc["context"]["contexts"])
jeffhsu3's avatar
jeffhsu3 committed
33
34
        return "abstract: {}\nquestion: {}\nanswer:".format(
            ctxs,
jeffhsu3's avatar
jeffhsu3 committed
35
36
            doc["question"],
            doc["final_decision"]
jeffhsu3's avatar
jeffhsu3 committed
37
38
39
        )

    def doc_to_target(self, doc):
jeffhsu3's avatar
jeffhsu3 committed
40
        return " {}".format(doc["final_decision"])
jeffhsu3's avatar
jeffhsu3 committed
41

42
43
44
45
46
47
    def fewshot_examples(self, k):
        # Since only test docs sample from test docs
        if self._training_docs is None:
            self._training_docs = list(self.test_docs())
        return random.sample(self._training_docs, k)

jeffhsu3's avatar
jeffhsu3 committed
48
49
50
51
52
53
54
55
56
57
    def construct_requests(self, doc, ctx):
        """ Uses RequestFactory to construct Requests and returns
        an iterable of Requests which will be sent to the LM.
        """
        ll_yes, _ = rf.loglikelihood(ctx, " yes")
        ll_no, _ = rf.loglikelihood(ctx, " no")
        ll_maybe, _ = rf.loglikelihood(ctx, " maybe")
        return ll_yes, ll_no, ll_maybe

    def process_results(self, doc, results):
jeffhsu3's avatar
jeffhsu3 committed
58
        gold = doc["final_decision"]
jeffhsu3's avatar
jeffhsu3 committed
59
60
61
        ll_yes, ll_no, ll_maybe = results
        pred = np.argmax(results)
        return {
jeffhsu3's avatar
jeffhsu3 committed
62
            "acc": ["yes", "no", "maybe"][pred] == gold, 
jeffhsu3's avatar
jeffhsu3 committed
63
64
65
66
        }

    def aggregation(self):
        return {
jeffhsu3's avatar
jeffhsu3 committed
67
            "acc" : mean
jeffhsu3's avatar
jeffhsu3 committed
68
69
70
71
        }

    def higher_is_better(self):
        return {
jeffhsu3's avatar
jeffhsu3 committed
72
            "acc" : True
jeffhsu3's avatar
jeffhsu3 committed
73
        }
jeffhsu3's avatar
merge  
jeffhsu3 committed
74
75
<<<<<<< HEAD
=======
jeffhsu3's avatar
jeffhsu3 committed
76

jeffhsu3's avatar
sciq  
jeffhsu3 committed
77
78
79
80
81
    def test_docs(self):
        if self.has_test_docs():
            # HF is labelled as train but its really just for testing
            return self.data["train"]

jeffhsu3's avatar
jeffhsu3 committed
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

class SciQ(MultipleChoiceTask):
    def download(self):
        if not os.path.exists('data/sciq'):
            os.mkdir('data/sciq')
            sh((
                "wget https://ai2-public-datasets.s3.amazonaws.com/sciq/SciQ.zip -O data/sciq/SciQ.zip"
            ))
            with zipfile.ZipFile("data/sciq/SciQ.zip", "r") as zf:
                zf.extractall("data/sciq/")

    def has_training_docs(self):
        return True

    def has_validation_docs(self):
        return True

    def has_test_docs(self):
        return True

jeffhsu3's avatar
sciq  
jeffhsu3 committed
102
    def _convert_standard(self, doc):
jeffhsu3's avatar
jeffhsu3 committed
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
        choices = [
            doc["distractor1"], 
            doc["distractor2"], 
            doc["distractor3"],
            doc["correct_answer"],
        ]
        src = doc['support']
        out_doc = {
            "source" : src,
            "query" : doc['question'],
            "choices" : choices,
            "gold" : 3,
        }
        return out_doc
    
    def load_docs(self, textfilename):
jeffhsu3's avatar
sciq  
jeffhsu3 committed
119
120
        with open(textfilename, 'r') as j:
            docs = json.loads(j.read()) 
jeffhsu3's avatar
jeffhsu3 committed
121
        for record in docs:
jeffhsu3's avatar
sciq  
jeffhsu3 committed
122
            yield self._convert_standard(record)
jeffhsu3's avatar
jeffhsu3 committed
123
124
125
126
127
128
129

    def fewshot_description(self):
        # Average ctx length in labelled dataset is 238.9
        # 2 few-shot exmamples pushes it beyond context window
        return ""

    def training_docs(self):
jeffhsu3's avatar
sciq  
jeffhsu3 committed
130
        return self.load_docs("data/sciq/SciQ dataset-2 3/train.json")
jeffhsu3's avatar
jeffhsu3 committed
131
132

    def validation_docs(self):
jeffhsu3's avatar
sciq  
jeffhsu3 committed
133
        return self.load_docs("data/sciq/SciQ dataset-2 3/valid.json")
jeffhsu3's avatar
jeffhsu3 committed
134
135

    def test_docs(self):
jeffhsu3's avatar
sciq  
jeffhsu3 committed
136
        return self.load_docs("data/sciq/SciQ dataset-2 3/test.json")
jeffhsu3's avatar
jeffhsu3 committed
137
138

    def doc_to_text(self, doc):
Leo Gao's avatar
Leo Gao committed
139
        return "{}\n{}".format(doc["source"], doc["query"])
jeffhsu3's avatar
jeffhsu3 committed
140
141
142
143

class EmrQA():
    def load_docs(self, textfilename):
        pass
jeffhsu3's avatar
merge  
jeffhsu3 committed
144
>>>>>>> 30c2fe23657981d8a6155da1bd8e8098487b771a