|
from ..root import DATASETS |
|
from ..utils import MInstrDataset |
|
|
|
|
|
@DATASETS.register_module() |
|
class InstructDataset(MInstrDataset): |
|
def __init__(self, *args, add_coco_prefix=False, **kwargs): |
|
super().__init__(*args, **kwargs, placeholders=(), template_string='', template_file=None) |
|
self.add_coco_prefix = add_coco_prefix |
|
|
|
def __getitem__(self, index): |
|
item = self.get_raw_item(index) |
|
if self.add_coco_prefix: |
|
img_path = f"COCO_train2014_{item['image']}" |
|
else: |
|
img_path = item['image'] |
|
conversations = item['conversations'] |
|
|
|
image = self.get_image(img_path) |
|
ret = { |
|
'image': image, |
|
'conversations': conversations, |
|
} |
|
return ret |
|
|
