Skip to content

Commit

Permalink
Update mm_chat.py
Browse files Browse the repository at this point in the history
Handle string image (paths)
  • Loading branch information
afrizalhasbi authored Oct 11, 2024
1 parent 2fbc6b0 commit b1d2a3c
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions src/axolotl/utils/collators/mm_chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from transformers import PreTrainedTokenizerBase, ProcessorMixin
from transformers.data.data_collator import DataCollatorMixin
from transformers.utils import PaddingStrategy

from PIL import Image

@dataclass
class MultiModalChatDataCollator(DataCollatorMixin):
Expand Down Expand Up @@ -52,7 +52,7 @@ def process_rows(examples, processor, chat_template, max_images, length_only=Fal
)
for example in examples
]
images = [example["images"] for example in examples]
images = [Image.open(example["images"]) if type(example["images"])==str else example["images"] for example in examples]

if max_images > 0:
images = [img_batch[:max_images] for img_batch in images]
Expand Down

0 comments on commit b1d2a3c

Please sign in to comment.