Upload folder using huggingface_hub

modelling_medicap.py

@@ -160,6 +160,9 @@ class MedICapEncoderDecoderModel(VisionEncoderDecoderModel):
             argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
         }
 
+        if decoder_inputs_embeds is None:
+            decoder_inputs_embeds = self.decoder.transformer.wte(decoder_input_ids)
+
         if encoder_outputs is None:
             if pixel_values is None:
                 raise ValueError("You have to specify pixel_values")
@@ -170,23 +173,18 @@ class MedICapEncoderDecoderModel(VisionEncoderDecoderModel):
                 return_dict=return_dict,
                 **kwargs_encoder,
             )  # CvT does not support output_attentions.
-        elif isinstance(encoder_outputs, tuple):
-            encoder_outputs = BaseModelOutput(*encoder_outputs)
-
-        embeddings = self.decoder.transformer.wte(decoder_input_ids)
-        embeddings = torch.cat([encoder_outputs[0], embeddings], dim=1)
-
-        if torch.is_tensor(decoder_attention_mask):
-            decoder_attention_mask = torch.cat(
-                [
-                    torch.ones(encoder_outputs[0].shape[:-1], dtype=decoder_attention_mask.dtype, device=self.device), 
-                    decoder_attention_mask
-                ], 
-                dim=1,
-            )
+            assert decoder_inputs_embeds.shape[1] == 1
+            decoder_inputs_embeds = torch.cat([encoder_outputs[0], decoder_inputs_embeds], dim=1)
+            if decoder_attention_mask is not None:
+                decoder_attention_mask = torch.cat(
+                    [
+                        torch.ones(encoder_outputs[0].shape[:-1], dtype=decoder_attention_mask.dtype, device=self.device), 
+                        decoder_attention_mask
+                    ], 
+                    dim=1,
+                )
         
         decoder_outputs = self.decoder(
-            input_ids=decoder_input_ids,
             attention_mask=decoder_attention_mask,
             inputs_embeds=decoder_inputs_embeds,
             output_attentions=output_attentions,
@@ -220,6 +218,48 @@ class MedICapEncoderDecoderModel(VisionEncoderDecoderModel):
             encoder_last_hidden_state=encoder_outputs.last_hidden_state,
         )
 
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        use_cache=None,
+        encoder_outputs=None,
+        **kwargs,
+    ):
+        """
+        Modification of: 
+            https://github.com/huggingface/transformers/blob/main/src/transformers/models/encoder_decoder/modeling_encoder_decoder.py#L660
+
+        This can help with managing input_embeds and input_ids: 
+            https://github.com/huggingface/transformers/issues/6535
+        """
+        input_dict = {'use_cache': use_cache, 'encoder_outputs': encoder_outputs, 'attention_mask': attention_mask}
+    
+        if past_key_values is None:
+            decoder_inputs = self.decoder.prepare_inputs_for_generation(
+                input_ids, inputs_embeds=encoder_outputs[0], past_key_values=past_key_values,
+            )
+            input_dict['decoder_inputs_embeds'] = decoder_inputs['inputs_embeds']
+        else:
+            decoder_inputs = self.decoder.prepare_inputs_for_generation(
+                input_ids, past_key_values=past_key_values,
+            )
+            input_dict['decoder_input_ids'] = decoder_inputs['input_ids']
+        input_dict['past_key_values'] = decoder_inputs['past_key_values']
+        input_dict['decoder_attention_mask'] = decoder_inputs['attention_mask'] if 'attention_mask' in decoder_inputs else None
+
+        # if torch.is_tensor(decoder_attention_mask):
+        #     decoder_attention_mask = torch.cat(
+        #         [
+        #             torch.ones(encoder_outputs[0].shape[:-1], dtype=decoder_attention_mask.dtype, device=self.device), 
+        #             decoder_attention_mask
+        #         ], 
+        #         dim=1,
+        #     )
+
+        return input_dict
+
     def tokenize_captions_teacher_forcing(
         self, 
         captions: str,