52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
import base64
|
|
from io import BytesIO
|
|
from PIL import Image
|
|
from typing import Any, Dict, Iterator, List
|
|
|
|
from evalscope.perf.arguments import Arguments
|
|
from evalscope.perf.plugin.datasets.base import DatasetPluginBase
|
|
from evalscope.perf.plugin.registry import register_dataset
|
|
|
|
|
|
def PIL_to_base64(image: Image.Image) -> str:
|
|
buffered = BytesIO()
|
|
image.save(buffered, format='JPEG')
|
|
img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
|
|
return img_str
|
|
|
|
|
|
@register_dataset('flickr8k')
|
|
class FlickrDatasetPlugin(DatasetPluginBase):
|
|
"""Read dataset and return prompt.
|
|
Datasets: https://www.modelscope.cn/datasets/clip-benchmark/wds_flickr8k/files
|
|
"""
|
|
|
|
def __init__(self, query_parameters: Arguments):
|
|
super().__init__(query_parameters)
|
|
|
|
def build_messages(self) -> Iterator[List[Dict]]:
|
|
from modelscope.msdatasets import MsDataset
|
|
dataset = MsDataset.load('clip-benchmark/wds_flickr8k', split='test')
|
|
|
|
for item in dataset:
|
|
pil_image = item['jpg']
|
|
text = item['txt']
|
|
base64_iamge = PIL_to_base64(pil_image)
|
|
|
|
yield [{
|
|
'role':
|
|
'user',
|
|
'content': [
|
|
{
|
|
'type': 'text',
|
|
'text': text,
|
|
},
|
|
{
|
|
'type': 'image_url',
|
|
'image_url': {
|
|
'url': f'data:image/jpeg;base64,{base64_iamge}',
|
|
}
|
|
},
|
|
],
|
|
}]
|