Add batching support

modeling_tinyllava_phi.py

@@ -1,6 +1,7 @@
 # For licensing see accompanying LICENSE file.
 # Copyright (C) 2024 TinyLLaVA. All Rights Reserved.
 import time
+import numpy as np
 
 import dataclasses
 from enum import auto, Enum
@@ -160,26 +161,40 @@ def process_images(images, image_processor, model_cfg):
     return new_images
 
 
-def tokenizer_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):
-    prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split('<image>')]
+def tokenizer_image_token(prompts, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):
+    def process_single_prompt(prompt):
+        prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split('<image>')]
 
-    def insert_separator(X, sep):
-        return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]
+        def insert_separator(X, sep):
+            return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]
 
-    input_ids = []
-    offset = 0
-    if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:
-        offset = 1
-        input_ids.append(prompt_chunks[0][0])
+        input_ids = []
+        offset = 0
+        if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:
+            offset = 1
+            input_ids.append(prompt_chunks[0][0])
 
-    for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
-        input_ids.extend(x[offset:])
+        for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
+            input_ids.extend(x[offset:])
+
+        return input_ids
+
+    if isinstance(prompts, str):  # Handle single prompt
+        return process_single_prompt(prompts)
+
+    # Handle batch of prompts
+    batch_input_ids = [process_single_prompt(prompt) for prompt in prompts]
 
     if return_tensors is not None:
         if return_tensors == 'pt':
-            return torch.tensor(input_ids, dtype=torch.long)
+            max_length = max(len(ids) for ids in batch_input_ids)
+            padded_input_ids = [
+                ids + [tokenizer.pad_token_id] * (max_length - len(ids)) for ids in batch_input_ids
+            ]
+            return torch.tensor(padded_input_ids, dtype=torch.long)
         raise ValueError(f'Unsupported tensor type: {return_tensors}')
-    return input_ids
+
+    return batch_input_ids
 
 def load_image(image_file):
     if image_file.startswith("http") or image_file.startswith("https"):
@@ -204,9 +219,9 @@ class Connector(nn.Module):
         for _ in range(1, mlp_depth):
             modules.append(ACT_TYPE[act_type]())
             modules.append(nn.Linear(config.hidden_size, config.hidden_size))
-            
+
         self._connector = nn.Sequential(*modules)
-    
+
     def forward(self, x):
         return self._connector(x)
 
@@ -219,9 +234,9 @@ class VisionTower(nn.Module):
         else:
             self._vision_tower = SiglipVisionModel(cfg)
             self._image_processor = SiglipImageProcessor.from_pretrained(cfg.model_name_or_path)
-            
+
         self.config = cfg
-        
+
     def forward(self, x, **kwargs):
         image_features = self._vision_tower(x, output_hidden_states=True)
         image_features = image_features.hidden_states[kwargs.get('vision_feature_layer', -2)]
@@ -234,11 +249,11 @@ class VisionTower(nn.Module):
             raise ValueError(f"Unexpected select feature: {kwargs.get('vision_feature_select_strategy')}")
 
         return image_features
-        
+
     @property
     def vision_tower(self):
         return self._vision_tower
-        
+
     @vision_tower.setter
     def vision_tower(self, vision_tower):
         self._vision_tower = vision_tower
@@ -248,7 +263,7 @@ def get_value_from_kwargs(kwargs, name):
         return kwargs.pop(name)
     else:
         return None
-    
+
 
 class TinyLlavaPreTrainedModel(PreTrainedModel):
     config_class = TinyLlavaConfig
@@ -284,7 +299,7 @@ class TinyLlavaPreTrainedModel(PreTrainedModel):
 
 class TinyLlavaForConditionalGeneration(TinyLlavaPreTrainedModel):
     def __init__(self, config: TinyLlavaConfig):
-        
+
         super().__init__(config)
 
         self.language_model = PhiForCausalLM(config.text_config)
@@ -292,7 +307,7 @@ class TinyLlavaForConditionalGeneration(TinyLlavaPreTrainedModel):
         self.connector = Connector(config)
         self.post_init()
 
-    
+
     def get_input_embeddings(self):
         return self.language_model.get_input_embeddings()
 
@@ -322,7 +337,7 @@ class TinyLlavaForConditionalGeneration(TinyLlavaPreTrainedModel):
         self.vocab_size = model_embeds.num_embeddings
         return model_embeds
 
-    
+
     def forward(
         self,
         input_ids: torch.LongTensor = None,
@@ -368,7 +383,7 @@ class TinyLlavaForConditionalGeneration(TinyLlavaPreTrainedModel):
             output_hidden_states=output_hidden_states,
             return_dict=return_dict
         )
-    
+
     @torch.no_grad()
     def generate(
         self,
@@ -408,7 +423,7 @@ class TinyLlavaForConditionalGeneration(TinyLlavaPreTrainedModel):
             inputs_embeds=inputs_embeds,
             **kwargs
         )
-        
+
     def encode_images(self, images):
         kwargs = {}
         kwargs['vision_feature_layer'] = self.config.vision_feature_layer
@@ -417,9 +432,9 @@ class TinyLlavaForConditionalGeneration(TinyLlavaPreTrainedModel):
         image_features = self.vision_tower(images, **kwargs)
         image_features = self.connector(image_features)
         return image_features
-    
-    
-    
+
+
+
     def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
                                       inputs_embeds=None, **kwargs):
         images = kwargs.pop("images", None)
@@ -432,7 +447,7 @@ class TinyLlavaForConditionalGeneration(TinyLlavaPreTrainedModel):
         if image_sizes is not None:
             inputs['image_sizes'] = image_sizes
         return inputs
-        
+
     def prepare_inputs_labels_for_multimodal(
         self, input_ids, position_ids, attention_mask, past_key_values, labels,
         images, image_sizes=None
@@ -441,7 +456,7 @@ class TinyLlavaForConditionalGeneration(TinyLlavaPreTrainedModel):
         if vision_tower is None or images is None or input_ids.shape[1] == 1:
             return input_ids, position_ids, attention_mask, past_key_values, None, labels
 
-        
+
         image_features = self.encode_images(images)
 
         # TODO: image start / end is not implemented here to support pretraining.
@@ -565,40 +580,72 @@ class TinyLlavaForConditionalGeneration(TinyLlavaPreTrainedModel):
             position_ids = None
 
         return None, position_ids, attention_mask, past_key_values, new_input_embeds, new_labels
-    
+
     def chat(
         self,
-        prompt: str,
-        tokenizer = None,
-        image: str = None,
+        prompts: Union[list, str],
+        tokenizer=None,
+        images: Union[list, str] = None,
         max_new_tokens: int = 512,
-        num_beams = 1,
+        num_beams=1,
         top_p=None,
         temperature=0
     ):
+        """
+        Generate responses for a batch of prompts.
+
+        Args:
+            prompts (list): List of text prompts.
+            tokenizer: Tokenizer object.
+            images (list, optional): List of image file paths corresponding to the prompts. Defaults to None.
+            max_new_tokens (int): Maximum number of new tokens to generate. Defaults to 512.
+            num_beams (int): Number of beams for beam search. Defaults to 1.
+            top_p (float, optional): Nucleus sampling probability. Defaults to None.
+            temperature (float): Sampling temperature. Defaults to 0.
+
+        Returns:
+            list: List of generated outputs.
+            list: List of generation times for each batch.
+        """
+        if isinstance(prompts, list) and isinstance(images, str):
+            assert len(prompts) == len(images) or images is None, "Mismatch between prompts and images."
+        else:
+            prompts = [prompts]
+            images = [images]
         image_processor = self.vision_tower._image_processor
 
-        if image is not None:
-            prompt = DEFAULT_IMAGE_TOKEN + '\n' + prompt 
-        conv = conv_phi_v0.copy()
-        conv.append_message(conv.roles[0], prompt)
-        conv.append_message(conv.roles[1], None)
-        prompt = conv.get_prompt()
-        if image is not None:
-            image = load_image(image)
-            image_tensor = process_images(image, image_processor, self.config).to(self.device)
-
-        input_ids = (
-            tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt")
-            .unsqueeze(0).to(self.device)
-        )
-        # Generate
+        # Prepare inputs
+        input_texts = []
+        image_tensors = None
+
+        for i, prompt in enumerate(prompts):
+            if images and images[i] is not None:
+                prompt = DEFAULT_IMAGE_TOKEN + '\n' + prompt
+            conv = conv_phi_v0.copy()
+
+            conv.append_message(conv.roles[0], prompt)
+            conv.append_message(conv.roles[1], None)
+            input_texts.append(conv.get_prompt())
+
+        # Tokenize prompts
+        input_ids = tokenizer_image_token(
+            input_texts, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
+        ).to(self.device)
+
+        # Process images
+        if images:
+            processed_images = [
+                process_images(load_image(image), image_processor, self.config)
+                for image in images if image is not None
+            ]
+            image_tensors = torch.stack(processed_images).to(self.device).squeeze(1)
+
+        # Generate responses
         stime = time.time()
-
         with torch.inference_mode():
             output_ids = self.generate(
                 input_ids,
-                images=image_tensor,
+                images=image_tensors,
                 do_sample=True if temperature > 0 else False,
                 temperature=temperature,
                 top_p=top_p,
@@ -606,19 +653,15 @@ class TinyLlavaForConditionalGeneration(TinyLlavaPreTrainedModel):
                 pad_token_id=tokenizer.pad_token_id,
                 max_new_tokens=max_new_tokens,
                 use_cache=True,
-                # stopping_criteria=[stopping_criteria],
             )
-
-        # print('inference over')
         generation_time = time.time() - stime
-        outputs = tokenizer.batch_decode(
-            output_ids, skip_special_tokens=True
-        )[0]
 
-        outputs = outputs.strip()
+        # Decode outputs
+        outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+        outputs = [output.strip() for output in outputs]
 
         return outputs, generation_time
-    
 
-AutoConfig.register("tinyllava", TinyLlavaConfig)        
-AutoModelForCausalLM.register(TinyLlavaConfig, TinyLlavaForConditionalGeneration)
+
+AutoConfig.register("tinyllava", TinyLlavaConfig)
+AutoModelForCausalLM.register(TinyLlavaConfig, TinyLlavaForConditionalGeneration)