File size: 1,081 Bytes
bbfa6f6 |
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 |
import os
import random
from llava.datasets.builder import DATASETS
from typing import Dict, Optional, Sequence, List
from llava.datasets.data_cfgs import data_configs
from llava.datasets.base_dataset import ImageTaskDataset
from llava.datasets.prompts import cc_sbu_prompt
from llava.constants import DEFAULT_IMAGE_TOKEN
class CCSBUDataset(ImageTaskDataset):
def __init__(self, anno_path, data_args=None, name='cc_sbu'):
super().__init__(anno_path=anno_path,
data_args=data_args,
name=name)
def text_preprocess(self, item) -> List[Dict[str, str]]:
caption = item['caption']
conversations = [
{
'from': 'human',
'value': DEFAULT_IMAGE_TOKEN + random.choice(cc_sbu_prompt)
},
{
'from': 'model',
'value': caption
}
]
return conversations
@DATASETS.register_obj
def cc_sbu(data_args):
return CCSBUDataset(data_configs["cc_sbu"]['train_data_path'], data_args)
|