OpenCompass/opencompass/datasets/TheoremQA.py
2023-07-04 22:11:33 +08:00

28 lines
631 B
Python

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().split('\n')[0].strip()
matches = re.findall(r'answer is (.*)', text)
if len(matches) == 0:
return text
else:
text = matches[0].strip()[:-1]
return text