File size: 2,606 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 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 |
import base64
import io
from typing import List, Optional
import pandas as pd
from mmengine.dataset import Compose
from PIL import Image
from torch.utils.data import Dataset
from opencompass.registry import DATASETS
def decode_base64_to_image(base64_string) -> Image:
"""Convert raw data into Pillow image."""
image_data = base64.b64decode(base64_string)
image = Image.open(io.BytesIO(image_data))
return image
@DATASETS.register_module()
class MMBenchDataset(Dataset):
"""Dataset to load MMBench dataset.
Args:
data_file (str): The path of the dataset.
pipeline (dict): The data augmentation.
sys_prompt (str): The system prompt added to the head
of these options. Defaults to
There are several options:
"""
def __init__(self,
data_file: str,
pipeline: List[dict],
sys_prompt: str = 'There are several options:') -> None:
self.df = pd.read_csv(data_file, sep='\t')
self.pipeline = Compose(pipeline)
self.sys_prompt = sys_prompt
def __len__(self) -> None:
return len(self.df)
def __getitem__(self, idx: int) -> dict:
# Mandatory Fields Begin
index = self.df.iloc[idx]['index']
image = self.df.iloc[idx]['image']
image = decode_base64_to_image(image)
question = self.df.iloc[idx]['question']
option_candidate = ['A', 'B', 'C', 'D', 'E']
options = {
cand: self.load_from_df(idx, cand)
for cand in option_candidate
if self.load_from_df(idx, cand) is not None
}
options_prompt = f'{self.sys_prompt}\n'
for key, item in options.items():
options_prompt += f'{key}. {item}\n'
# Mandatory Fields End
# Optional Fields Begin
hint = self.load_from_df(idx, 'hint')
category = self.load_from_df(idx, 'category')
l2_catetory = self.load_from_df(idx, 'l2-category')
# Optional Fields End
data = {
'img': image,
'question': question,
'options': options_prompt,
'category': category,
'l2-category': l2_catetory,
'options_dict': options,
'index': index,
'context': hint,
}
data = self.pipeline(data)
return data
def load_from_df(self, idx: int, key: str) -> Optional[str]:
if key in self.df.iloc[idx] and not pd.isna(self.df.iloc[idx][key]):
return self.df.iloc[idx][key]
else:
return None
|