@TextGenerationDataset.register(name="flickr", alias="flickr30k")
class Flickr30K(TextGenerationDataset):
"""
:param dataset_args: configuration settings for dataset loading
:param split: split from dataset to load, for instance `test` or `train[:5%]`
:param processor: processor or tokenizer to use on dataset
"""
DEFAULT_CHAT_TEMPLATE = (
"{% for message in messages %}\n"
"{% if message['role'] == 'user' %}\n"
"{{ '<|user|>\n' + message['content'] + eos_token }}\n"
"{% elif message['role'] == 'system' %}\n"
"{{ '<|system|>\n' + message['content'] + eos_token }}\n"
"{% elif message['role'] == 'assistant' %}\n"
"{{ '<|assistant|>\n' + message['content'] + eos_token }}\n"
"{% endif %}\n"
"{% if loop.last and add_generation_prompt %}\n"
"{{ '<|assistant|>' }}\n{% endif %}\n{% endfor %}"
)
def __init__(
self, dataset_args: "DatasetArguments", split: str, processor: Processor
):
dataset_args = deepcopy(dataset_args)
dataset_args.dataset = "lmms-lab/flickr30k"
super().__init__(dataset_args=dataset_args, split=split, processor=processor)
if (
self.tokenizer is not None
and getattr(self.tokenizer, "chat_template", None) is None
):
# note that since tokenizer is a member of processor,
# this change affects processor.apply_chat_template
self.tokenizer.chat_template = self.DEFAULT_CHAT_TEMPLATE
logger.warning(
"tokenizer.chat_template is not set, using default chat template for "
f"{self.__class__.__name__}"
)
def dataset_template(self, sample):
messages = [
{
"role": "user",
"content": [
{"type": "image"},
{"type": "text", "text": "What does this image show?"},
],
},
{
"role": "assistant",
"content": " ".join(sample["caption"]),
},
]
return {
"text": self.processor.apply_chat_template(
messages,
add_generation_prompt=False,
),
"images": sample["image"],
}