-
Notifications
You must be signed in to change notification settings - Fork 357
/
triviaqa.py
95 lines (77 loc) · 3.25 KB
/
triviaqa.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
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
89
90
91
92
93
94
95
import csv
import json
import os.path as osp
from datasets import Dataset, DatasetDict
from opencompass.openicl.icl_evaluator import BaseEvaluator
from opencompass.registry import ICL_EVALUATORS, LOAD_DATASET
from opencompass.utils.text_postprocessors import general_postprocess
from .base import BaseDataset
@LOAD_DATASET.register_module()
class TriviaQADataset(BaseDataset):
@staticmethod
def load(path: str):
dataset = DatasetDict()
for split in ['dev', 'test']:
filename = osp.join(path, f'trivia-{split}.qa.csv')
with open(filename, 'r', encoding='utf-8') as f:
reader = csv.reader(f, delimiter='\t')
raw_data = []
for row in reader:
assert len(row) == 2
question = row[0]
answers = eval(row[1])
if split == 'test':
answers = answers[0]
raw_data.append({'question': question, 'answer': answers})
dataset[split] = Dataset.from_list(raw_data)
return dataset
@LOAD_DATASET.register_module()
class TriviaQADataset_V2(BaseDataset):
@staticmethod
def load(path: str):
dataset = DatasetDict()
for split in ['validation', 'train']:
filename = osp.join(path, f'triviaqa-{split}.jsonl')
raw_data = []
with open(filename, 'r', encoding='utf-8') as f:
for doc in f:
doc = json.loads(doc)
raw_data.append(doc)
dataset[split] = Dataset.from_list(raw_data)
return dataset
@LOAD_DATASET.register_module()
class TriviaQADataset_V3(BaseDataset):
@staticmethod
def load(path: str):
data_list = []
with open(path, 'r', encoding='utf-8') as f:
for doc in f:
data_list.append(json.loads(doc))
return Dataset.from_list(data_list)
@ICL_EVALUATORS.register_module()
class TriviaQAEvaluator(BaseEvaluator):
def score(self, predictions, references):
if len(predictions) != len(references):
return {'error': 'preds and refrs have different length'}
processed_predictions = []
for prediction in predictions:
prediction = prediction.strip().split('\n')[0].lower()
prediction = prediction.split('answer is')[-1]
prediction = prediction.split('a:')[-1]
prediction = prediction.split('answer:')[-1]
prediction = prediction.strip()
prediction = general_postprocess(prediction)
processed_predictions.append(prediction)
processed_answers = [[general_postprocess(j).lower() for j in i]
for i in references]
details = []
cnt = 0
for pred, cand_ans in zip(processed_predictions, processed_answers):
detail = {'pred': pred, 'answer': cand_ans, 'correct': False}
# is_correct = any([cand == pred for cand in cand_ans])
is_correct = any([cand in pred for cand in cand_ans])
cnt += int(is_correct)
detail['correct'] = is_correct
details.append(detail)
score = cnt / len(predictions) * 100
return {'score': score, 'details': details}