squad.py 4.58 KB
Newer Older
1
2
3
import datasets
from lm_eval.base import rf, f1_score, mean
from . common import HFTask
Charles Foster's avatar
Charles Foster committed
4
5
6
7
8
9
10
11
12
13
14

class SQuAD(HFTask):
    DATASET_PATH = "squad_v2"
    DATASET_NAME = None

    def has_training_docs(self):
        return True

    def has_validation_docs(self):
        return True

15
16
17
    def has_test_docs(self):
        return False

Charles Foster's avatar
Charles Foster committed
18
    def training_docs(self):
19
        return self.data["train"]
Charles Foster's avatar
Charles Foster committed
20
21

    def validation_docs(self):
22
        return self.data["validation"]
Charles Foster's avatar
Charles Foster committed
23
24

    def fewshot_description(self):
25
26
        # TODO: figure out description
        return ""
Charles Foster's avatar
Charles Foster committed
27

28
    def doc_to_text(self, doc):
Leo Gao's avatar
Leo Gao committed
29
        return 'Title: ' + doc['title'] + '\n\n' + 'Background: ' + doc['context'] + '\n\n' + 'Question: ' + doc['question'] + '\n\n' + 'Answer:'
30
31
32
33
34
35
36

    def doc_to_target(self, doc):
        answer_list = doc['answers']['text']
        if len(answer_list) > 0:
            answer = answer_list[0]
        else:
            answer = 'unanswerable'
37
        return " " + answer
Charles Foster's avatar
Charles Foster committed
38

Leo Gao's avatar
Leo Gao committed
39
40
41
    def construct_requests(self, doc, ctx):
        """ Uses RequestFactory to construct Requests and returns an iterable of 
        Requests which will be sent to the LM.
42

Leo Gao's avatar
Leo Gao committed
43
44
45
46
47
48
49
        :param doc:
            The document as returned from training_docs, validation_docs, or test_docs.
        :param ctx: str
            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`. 
        """
50
51
        continuation = rf.greedy_until(ctx, ['\n'])
        return continuation
Leo Gao's avatar
Leo Gao committed
52
53
54
55
56
57
58
59
60
61
62
    
    def process_results(self, doc, results):
        """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 
        the metric for that one document

        :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.
        """
63
64
        squad_metric = datasets.load_metric("squad_v2")

Leo Gao's avatar
Leo Gao committed
65
66
        continuation, = results

67
68
        predictions = {
            'id': doc['id'],
Leo Gao's avatar
Leo Gao committed
69
            'prediction_text': continuation,
70
71
72
73
74
75
76
77
78
        }

        references = {
            'id': doc['id'],
            'answers': doc['answers'],
        }

        metrics = squad_metric.compute(predictions=predictions, references=references)

79
80
81
82
83
84
        metrics.pop('total', None)
        metrics.pop('HasAns_total', None)
        metrics.pop('NoAns_total', None)
        metrics.pop('best_exact_thresh', None)
        metrics.pop('best_f1_thresh', None)

85
        return metrics
Leo Gao's avatar
Leo Gao committed
86
87
88
89
90
91
92

    def aggregation(self):
        """
        :returns: {str: [float] -> float}
            A dictionary where keys are the names of submetrics and values are 
            functions that aggregate a list of metrics
        """
93
94
95
96
97
98
99
100
101
102
        return { 
            'exact': mean, # Exact match (the normalized answer exactly match the gold answer)
            'f1': mean, #  The F-score of predicted tokens versus the gold answer
            'HasAns_exact': mean, # Exact match (the normalized answer exactly match the gold answer)
            'HasAns_f1': mean, # The F-score of predicted tokens versus the gold answer
            'NoAns_exact': mean, # Exact match (the normalized answer exactly match the gold answer)
            'NoAns_f1': mean, # The F-score of predicted tokens versus the gold answer
            'best_exact': mean, # Best exact match (with varying threshold)
            'best_f1': mean, # Best F1 (with varying threshold)
        }
Leo Gao's avatar
Leo Gao committed
103
104
105
106
107
108
109

    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
        """
110
111
112
113
114
115
116
117
118
119
        return { 
            'exact': True, # Exact match (the normalized answer exactly match the gold answer)
            'f1': True, #  The F-score of predicted tokens versus the gold answer
            'HasAns_exact': True, # Exact match (the normalized answer exactly match the gold answer)
            'HasAns_f1': True, # The F-score of predicted tokens versus the gold answer
            'NoAns_exact': True, # Exact match (the normalized answer exactly match the gold answer)
            'NoAns_f1': True, # The F-score of predicted tokens versus the gold answer
            'best_exact': True, # Best exact match (with varying threshold)
            'best_f1': True, # Best F1 (with varying threshold)
        }