fix for null past_key_values, getting supported tasks

mlfoundations · Oct 24, 2023 · e104a0a · e104a0a
1 parent 17cd546
commit e104a0a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 9 deletions.
diff --git a/open_flamingo/eval/eval_models/eval_model.py b/open_flamingo/eval/eval_models/eval_model.py
@@ -195,7 +195,7 @@ def supported_tasks(self):
         Parsed by checking whether the model has a method called `get_{task}_prompt`.
         """
         return [
-            task.split("_")[1]
+            "_".join(task.split("_")[1:-1])
             for task in dir(self)
             if task.startswith("get_") and task.endswith("_prompt")
         ]
@@ -207,4 +207,4 @@ def _validate_text(self, batch_text):
         if any([x.endswith(" ") for x in batch_text]):
             print(
                 "Warning: trailing whitespace detected in text. This can cause unexpected behavior."
-            )
+            )
diff --git a/open_flamingo/src/vlm.py b/open_flamingo/src/vlm.py
@@ -414,13 +414,14 @@ def _prepare_inputs_for_forward(
             past_vision_tokens=past_vision_tokens,
             num_beams=num_beams,
         )
-        past_key_values = [
-            (
-                k.repeat_interleave(num_beams, dim=0),
-                v.repeat_interleave(num_beams, dim=0)
-            )
-            for k, v in past_key_values
-        ]
+        if past_key_values is not None: 
+            past_key_values = [
+                (
+                    k.repeat_interleave(num_beams, dim=0),
+                    v.repeat_interleave(num_beams, dim=0)
+                )
+                for k, v in past_key_values
+            ]
         return {
             "input_ids": lang_x,
             "attention_mask": attention_mask,