File size: 1,006 Bytes
256a159 |
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 |
import re
from datasets import load_dataset
from opencompass.registry import LOAD_DATASET, TEXT_POSTPROCESSORS
from .base import BaseDataset
@LOAD_DATASET.register_module()
class TheoremQADataset(BaseDataset):
@staticmethod
def load(path: str):
return load_dataset('csv', data_files={'test': path})
@TEXT_POSTPROCESSORS.register_module('TheoremQA')
def TheoremQA_postprocess(text: str) -> str:
text = text.strip()
matches = re.findall(r'answer is ([^\s]+)', text)
if len(matches) == 0:
return text
else:
text = matches[0].strip().strip('.,?!\"\';:')
return text
def TheoremQA_postprocess_v2(text: str) -> str:
prediction = text.strip().strip('\n').split('\n')[-1]
tmp = ''
for entry in prediction.split(' ')[::-1]:
if entry == 'is' or entry == 'be' or entry == 'are' or entry.endswith(
':'):
break
tmp = entry + ' ' + tmp
prediction = tmp.strip().strip('.')
return prediction
|