Handle Batch Sizes

I hope I'm not overstepping boundaries. I tried the fix on #16 to handle batch sizes larger than 1 but stumbled upon a few issues. For example, issues processing single strings as the text input or not handling padding when the input texts are of different size.

I tried to address all the issues I had. Hope this can be of use.



@gugarosa

feel free to have a look and let me know if I'm overlooking something.

Best,
William

processing_phi3_v.py

@@ -73,9 +73,7 @@ class Phi3VProcessor(ProcessorMixin):
 
         Args:
             text (`str`, `List[str]`, `List[List[str]]`):
-                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
-                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
-                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+                The sequence or batch of sequences to be encoded. Each sequence must be a string.
             images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
                 The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
                 tensor. Both channels-first and channels-last formats are supported.
@@ -150,7 +148,15 @@ class Phi3VProcessor(ProcessorMixin):
             return BatchFeature(data={**model_inputs})
 
         pattern = r"<\|image_\d+\|>"
-        prompt_chunks = [self.tokenizer(chunk).input_ids for chunk in re.split(pattern, texts)] 
+        
+        if isinstance(texts, str):
+           texts = [texts]
+
+        prompt_chunks = []
+        image_tags = []
+        for text in texts:
+            prompt_chunks.append([self.tokenizer(chunk).input_ids for chunk in re.split(pattern, text)])
+            image_tags.append(re.findall(pattern, text))
 
         if 'num_img_tokens' in images:
             num_img_tokens = images['num_img_tokens']
@@ -162,30 +168,38 @@ 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_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_ids = [int(s.split("|")[1].split("_")[-1]) for s in image_tags]
-        unique_image_ids = sorted(list(set(image_ids)))
+        image_ids = [[int(s.split("|")[1].split("_")[-1]) for s in tags] for tags in image_tags]
+        unique_image_ids = sorted(list(set([iid for ids in image_ids for iid in ids])))
         # image_ids must start from 1, and must be continuous int, e.g. [1, 2, 3], cannot be [1, 4, 5]
         # check the condition
         assert unique_image_ids == list(range(1, len(unique_image_ids)+1)), f"image_ids must start from 1, and must be continuous int, e.g. [1, 2, 3], cannot be {unique_image_ids}"
         # total images must be the same as the number of image tags
         assert len(unique_image_ids) == len(images), f"total images must be the same as the number of image tags, got {len(unique_image_ids)} image tags and {len(images)} images"
 
-        image_ids_pad = [[-iid]*num_img_tokens[iid-1] for iid in image_ids]
+        image_ids_pad = [[[-iid]*num_img_tokens[iid-1] for iid in ids] for ids in image_ids]
 
         def insert_separator(X, sep_list):
             if len(X) > len(sep_list):
                 sep_list.append([])
             return [ele for sublist in zip(X, sep_list) for ele in sublist]
         input_ids = []
-        offset = 0                
-        for x in insert_separator(prompt_chunks, image_ids_pad):
-            input_ids.extend(x[offset:])
+        for sub_prompt_chunks, sub_image_ids_pad in zip(prompt_chunks, image_ids_pad):
+            input_ids.append([])
+            offset = 0
+            for x in insert_separator(sub_prompt_chunks, sub_image_ids_pad):
+                input_ids[-1].extend(x[offset:])
+
+        max_length = max(len(ids) for ids in input_ids)
+        for i in range(len(input_ids)):
+            while len(input_ids[i]) < max_length:
+                input_ids[i] = [self.tokenizer.pad_token_id]+input_ids[i]
 
         input_ids = torch.tensor(input_ids, dtype=torch.long).unsqueeze(0)
         attention_mask = (input_ids > -1000000).to(torch.long)
+        attention_mask[input_ids == self.tokenizer.pad_token_id] = 0
 
         return BatchFeature(data={"input_ids": input_ids,
                                   "attention_mask": attention_mask,