Update model files

modeling_llava.py

@@ -9,8 +9,7 @@ from torch import nn
 from transformers import PreTrainedModel
 from transformers.modeling_outputs import ModelOutput
 
-from modeling_phi import PhiForCausalLM, InferenceParams
-from processing_llava import OpenCLIPImageProcessor
+from modeling_phi import PhiForCausalLM
 from configuration_llava import LlavaConfig
 from open_clip import create_model
 
@@ -22,7 +21,7 @@ class LlavaCausalLMOutputWithPast(ModelOutput):
     past_key_values: Optional[List[torch.FloatTensor]] = None
     hidden_states: Optional[Tuple[torch.FloatTensor]] = None
     attentions: Optional[Tuple[torch.FloatTensor]] = None
-    image_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    image_features: Optional[torch.FloatTensor] = None
 
 
 class LlavaMultiModalProjector(nn.Module):
@@ -214,14 +213,11 @@ class LlavaForConditionalGeneration(LlavaPreTrainedModel):
     def forward(
         self,
         input_ids: torch.LongTensor = None,
-        pixel_values: torch.FloatTensor = None,
+        image_features: torch.FloatTensor = None,
         attention_mask: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.LongTensor] = None,
         past_key_values: Optional[List[torch.FloatTensor]] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
-        vision_feature_layer: Optional[int] = None,
-        vision_feature_select_strategy: Optional[str] = None,
-        labels: Optional[torch.LongTensor] = None,
         use_cache: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
@@ -242,14 +238,8 @@ class LlavaForConditionalGeneration(LlavaPreTrainedModel):
         )
 
         if inputs_embeds is None:
-            # 1. Extra the input embeddings
             inputs_embeds = self.get_input_embeddings()(input_ids)
-
-            # 2. Merge text and images
-            if pixel_values is not None and input_ids.shape[1] != 1:
-                image_outputs = self.vision_model(pixel_values)
-
-                image_features = self.multi_modal_projector(image_outputs)
+            if image_features is not None and input_ids.shape[1] != 1:
                 (
                     inputs_embeds,
                     attention_mask,
@@ -261,46 +251,6 @@ class LlavaForConditionalGeneration(LlavaPreTrainedModel):
                     attention_mask,
                     position_ids,
                 )
-                # if labels is None:
-                #     labels = torch.full_like(
-                #         attention_mask, self.config.ignore_index
-                #     ).to(torch.long)
-            else:
-                # In case input_ids.shape[1] == 1 & pixel_values==None & past_key_values != None, we are in the case of
-                # generation with cache
-                if (
-                    past_key_values is not None
-                    and pixel_values is not None
-                    and input_ids.shape[1] == 1
-                ):
-                    # Retrieve the first layer to inspect the logits and mask out the hidden states
-                    # that are set to 0
-                    first_layer_past_key_value = past_key_values[0][0][:, :, :, 0]
-
-                    # Sum all dimensions of head_dim (-2) to avoid random errors such as: https://github.com/huggingface/transformers/pull/28032#issuecomment-1863691941
-                    batch_index, non_attended_tokens = torch.where(
-                        first_layer_past_key_value.float().sum(-2) == 0
-                    )
-
-                    # Get the target length
-                    target_seqlen = first_layer_past_key_value.shape[-1] + 1
-
-                    extended_attention_mask = torch.ones(
-                        (
-                            attention_mask.shape[0],
-                            target_seqlen - attention_mask.shape[1],
-                        ),
-                        dtype=attention_mask.dtype,
-                        device=attention_mask.device,
-                    )
-
-                    # Zero-out the places where we don't need to attend
-                    extended_attention_mask[batch_index, non_attended_tokens] = 0
-
-                    attention_mask = torch.cat(
-                        (attention_mask, extended_attention_mask), dim=1
-                    )
-                    position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 1
 
         outputs = self.language_model(
             input_ids=None,
@@ -316,37 +266,17 @@ class LlavaForConditionalGeneration(LlavaPreTrainedModel):
 
         logits = outputs[0]
 
-        loss = None
-        if labels is not None:
-            # Shift so that tokens < n predict n
-            if attention_mask is not None:
-                shift_attention_mask = attention_mask[..., 1:]
-                shift_logits = logits[..., :-1, :][
-                    shift_attention_mask.to(logits.device) != 0
-                ].contiguous()
-                shift_labels = labels[..., 1:][
-                    shift_attention_mask.to(labels.device) != 0
-                ].contiguous()
-            else:
-                shift_logits = logits[..., :-1, :].contiguous()
-                shift_labels = labels[..., 1:].contiguous()
-            # Flatten the tokens
-            loss_fct = nn.CrossEntropyLoss()
-            loss = loss_fct(
-                shift_logits.view(-1, shift_logits.size(-1)),
-                shift_labels.view(-1).to(shift_logits.device),
-            )
 
         if not return_dict:
             output = (logits,) + outputs[1:]
-            return (loss,) + output if loss is not None else output
+            return output
 
         return LlavaCausalLMOutputWithPast(
-            loss=loss,
             logits=logits,
             past_key_values=outputs.past_key_values,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
+            image_features=image_features,
         )
 
     def prepare_inputs_for_generation(
@@ -354,49 +284,15 @@ class LlavaForConditionalGeneration(LlavaPreTrainedModel):
         input_ids,
         past_key_values=None,
         inputs_embeds=None,
-        pixel_values=None,
         attention_mask=None,
+        image_features=None,
         **kwargs,
     ):
-        if past_key_values is not None:
-            if isinstance(past_key_values, InferenceParams):
-                cache_length = past_key_values.max_seqlen
-                past_length = past_key_values.seqlen_offset
-            else:
-                cache_length = past_length = past_key_values[0][0].shape[2]
-
-            # Keep only the unprocessed tokens:
-            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
-            # some of the inputs are exclusivelly passed as part of the cache (e.g. when passing input_embeds as
-            # input)
-            if (
-                attention_mask is not None
-                and attention_mask.shape[1] > input_ids.shape[1]
-            ):
-                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
-            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
-            # input_ids based on the past_length.
-            elif past_length < input_ids.shape[1]:
-                input_ids = input_ids[:, past_length:]
-            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
-            elif self.config.image_token_index in input_ids:
-                input_ids = input_ids[:, input_ids.shape[1] - 1 :]
-            # If the cache has seen more tokens than it can hold, then the cache has a size limit. Let's discard the
-            # older attention values, as their corresponding values are not part of the input.
-            if cache_length < past_length and attention_mask is not None:
-                attention_mask = attention_mask[
-                    :, -(cache_length + input_ids.shape[1]) :
-                ]
-
-        position_ids = kwargs.get("position_ids", None)
-        if attention_mask is not None and position_ids is None:
-            # create position_ids on the fly for batch generation
-            position_ids = attention_mask.long().cumsum(-1) - 1
-            position_ids.masked_fill_(attention_mask == 0, 1)
-            if past_key_values:
-                position_ids = position_ids[:, -input_ids.shape[1] :]
-
-        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        res = self.language_model.prepare_inputs_for_generation(input_ids, past_key_values, attention_mask, **kwargs)
+        input_ids = res["input_ids"]
+        past_key_values = res["past_key_values"]
+        attention_mask = res["attention_mask"]
+
         if inputs_embeds is not None and past_key_values is None:
             model_inputs = {"inputs_embeds": inputs_embeds}
         else:
@@ -404,11 +300,10 @@ class LlavaForConditionalGeneration(LlavaPreTrainedModel):
 
         model_inputs.update(
             {
-                "position_ids": position_ids,
                 "past_key_values": past_key_values,
                 "use_cache": kwargs.get("use_cache"),
                 "attention_mask": attention_mask,
-                "pixel_values": pixel_values,
+                "image_features": image_features,
             }
         )
         return model_inputs