From 86de7fbcc2da6e15f24dae80021b564ce9539aa0 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Tue, 7 Jan 2025 16:57:47 -0500 Subject: [PATCH] improve check for batched --- src/axolotl/prompt_strategies/chat_template.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/axolotl/prompt_strategies/chat_template.py b/src/axolotl/prompt_strategies/chat_template.py index 051460427..126c9d2f3 100644 --- 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]): """