Allow for multiple images per batch sample.

processing_phi3_v.py

@@ -20,14 +20,19 @@ import re
 from typing import List, Optional, Union
 
 import torch
-
 import transformers
 from transformers.feature_extraction_utils import BatchFeature
 from transformers.image_utils import ImageInput
 from transformers.processing_utils import ProcessorMixin
-from transformers.tokenization_utils_base import PaddingStrategy, TextInput, TruncationStrategy
+from transformers.tokenization_utils_base import (
+    PaddingStrategy,
+    TextInput,
+    TruncationStrategy,
+)
 from transformers.utils import TensorType
-from .image_processing_phi3_v import Phi3VImageProcessor 
+
+from .image_processing_phi3_v import Phi3VImageProcessor
+
 transformers.Phi3VImageProcessor = Phi3VImageProcessor 
 
 class Phi3VProcessor(ProcessorMixin):
@@ -144,13 +149,25 @@ class Phi3VProcessor(ProcessorMixin):
         return self.tokenizer.convert_tokens_to_ids(self.special_image_token)
     
     def _convert_images_texts_to_inputs(self, images, texts, padding=False, truncation=None, max_length=None, return_tensors=None):
-
         if not len(images):
             model_inputs = self.tokenizer(texts, return_tensors=return_tensors, padding=padding, truncation=truncation, max_length=max_length)
             return BatchFeature(data={**model_inputs})
 
         pattern = r"<\|image_\d+\|>"
-        prompt_chunks = [self.tokenizer(chunk).input_ids for chunk in re.split(pattern, texts)] 
+        
+        # Don't over list-comprehend this, it's already hard to read.
+        prompt_chunks = []  
+        image_tags = []
+        for text in texts:
+            chunks = re.split(pattern, text)
+            chunk_image_tags = re.findall(pattern, text)
+            for chunk in chunks:
+                tokenized_chunk = self.tokenizer(chunk).input_ids
+                prompt_chunks.append(tokenized_chunk)
+            for tag in chunk_image_tags:
+                image_tags.append(tag)
+
+                
 
         if 'num_img_tokens' in images:
             num_img_tokens = images['num_img_tokens']
@@ -161,10 +178,8 @@ class Phi3VProcessor(ProcessorMixin):
 
         images, image_sizes = images['pixel_values'], images['image_sizes']
 
-        # image_tags needs to start from 1 to n
-        image_tags = re.findall(pattern, texts) 
-        # image_ids = [int(s.split("|")[1].split("_")[-1]) * -1 for s in image_tags]
-        # image_ids_pad = [[iid]*num_img_tokens[i] for i, iid in enumerate(image_ids)]
+        
+         # image_tags needs to start from 1 to num_images
         image_ids = [int(s.split("|")[1].split("_")[-1]) for s in image_tags]
         unique_image_ids = sorted(list(set(image_ids)))
         # image_ids must start from 1, and must be continuous int, e.g. [1, 2, 3], cannot be [1, 4, 5]