Handle image input as string paths for MMLMs (#1958)

* Update mm_chat.py

Handle string image (paths)

* chore: lint

---------

Co-authored-by: Wing Lian <wing.lian@gmail.com>
This commit is contained in:
Afrizal Hasbi Azizy
2024-10-12 00:34:13 +07:00
committed by GitHub
parent 76883851d2
commit df359c8a6e

View File

@@ -4,6 +4,7 @@ Collators for multi-modal chat messages and packing
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Union from typing import Any, Dict, List, Optional, Union
from PIL import Image
from transformers import PreTrainedTokenizerBase, ProcessorMixin from transformers import PreTrainedTokenizerBase, ProcessorMixin
from transformers.data.data_collator import DataCollatorMixin from transformers.data.data_collator import DataCollatorMixin
from transformers.utils import PaddingStrategy from transformers.utils import PaddingStrategy
@@ -52,7 +53,12 @@ class MultiModalChatDataCollator(DataCollatorMixin):
) )
for example in examples for example in examples
] ]
images = [example["images"] for example in examples] images = [
Image.open(example["images"])
if isinstance(example["images"], str)
else example["images"]
for example in examples
]
if max_images > 0: if max_images > 0:
images = [img_batch[:max_images] for img_batch in images] images = [img_batch[:max_images] for img_batch in images]