|
import json |
|
import os.path as osp |
|
from typing import Dict, Optional |
|
|
|
import mmengine |
|
from datasets import Dataset, DatasetDict |
|
|
|
from opencompass.registry import TEXT_POSTPROCESSORS |
|
|
|
from ..base import BaseDataset |
|
|
|
|
|
class TEvalDataset(BaseDataset): |
|
|
|
def __init__(self, reader_cfg: Optional[Dict] = {}, **kwargs): |
|
super().__init__(reader_cfg=reader_cfg, **kwargs) |
|
|
|
def load(self, path: str, name: str): |
|
|
|
dataset = DatasetDict() |
|
data = mmengine.load(osp.join(path, f'{name}.json')) |
|
raw_data = [] |
|
for i in data.keys(): |
|
origin_prompt = data[i]['origin_prompt'] |
|
if isinstance(origin_prompt, str): |
|
origin_prompt = json.loads(origin_prompt) |
|
|
|
prompt = origin_prompt + [ |
|
dict(role='assistant', |
|
content=str(data[i].get('ground_truth'))) |
|
] |
|
raw_data.append({ |
|
'prompt': prompt, |
|
'ground_truth': json.dumps(data[i]) |
|
}) |
|
dataset['test'] = Dataset.from_list(raw_data) |
|
dataset['train'] = Dataset.from_list(raw_data) |
|
return dataset |
|
|
|
|
|
|
|
|
|
@TEXT_POSTPROCESSORS.register_module('teval') |
|
def teval_postprocess(text: str) -> str: |
|
if isinstance(text, str): |
|
text = text.split('<eoa>')[0] |
|
text = text.split('<TOKENS_UNUSED_1>')[0] |
|
text = text.split('<|im_end|>')[0] |
|
text = text.split('\nuser')[0] |
|
text = text.split('\nUSER')[0] |
|
text = text.split('[INST]')[0] |
|
text = text.strip() |
|
if text.startswith('```json'): |
|
text = text[len('```json'):] |
|
text = text.strip('`').strip() |
|
if text[:2] == '{{' and text[-2:] == '}}': |
|
text = text[1:-1] |
|
return str(text) |
|
|
