improve check for batched

axolotl-ai-cloud · Jan 7, 2025 · 86de7fb · 86de7fb
1 parent df9abb3
commit 86de7fb
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/src/axolotl/prompt_strategies/chat_template.py b/src/axolotl/prompt_strategies/chat_template.py
@@ -227,9 +227,12 @@ def supports_batched(self) -> bool:
         return True
 
     def is_prompt_batched(self, prompt: dict[str, Any]) -> bool:
-        return all(isinstance(v, list) for v in prompt.values()) and all(
-            isinstance(v, list) for v in prompt[self.messages]
-        )
+        try:
+            return all(isinstance(v, list) for v in prompt.values()) and all(
+                isinstance(v, list) for v in prompt[self.messages]
+            )
+        except KeyError:
+            return False
 
     def tokenize_prompt(self, prompt: dict[str, Any]):
         """