blip.py import and output truncation fix

mlfoundations · Mar 28, 2024 · a5378a8 · a5378a8
1 parent c7a5ae5
commit a5378a8
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/open_flamingo/eval/eval_models/blip.py b/open_flamingo/eval/eval_models/blip.py
@@ -4,7 +4,7 @@
 import torch
 
 from transformers import Blip2Processor, Blip2ForConditionalGeneration
-from models.eval_model import BaseEvalModel
+from eval_models.eval_model import BaseEvalModel
 from utils import unwrap_model
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
@@ -92,7 +92,8 @@ def get_outputs(
                 )
 
         # Extract only the new gnerated tokens
-        outputs = outputs[:, len(input_ids[0]) :]
+        if not unwrap_model(self.model).config.is_encoder_decoder:
+            outputs = outputs[:, len(input_ids[0]) :]
 
         return self.tokenizer.batch_decode(outputs, skip_special_tokens=True)