Fix(message): Improve error message for bad format (axolotl-ai-cloud#365

)
mkeoliya · Aug 12, 2023 · 1a7e2a0 · 1a7e2a0
1 parent b500e0c
commit 1a7e2a0
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/src/axolotl/prompt_strategies/llama2_chat.py b/src/axolotl/prompt_strategies/llama2_chat.py
@@ -29,7 +29,7 @@
 from typing import Generator, List, Sequence
 
 from axolotl.prompt_tokenizers import PromptTokenizingStrategy
-from axolotl.prompters import IGNORE_TOKEN_ID
+from axolotl.prompters import IGNORE_TOKEN_ID, SHAREGPT_ASSERTION_FAILED_ROLE
 
 
 @dataclass
@@ -190,7 +190,7 @@ def build_prompt(self, source) -> Generator[Llama2ChatConversation, None, None]:
         conv.messages = []  # pylint: disable=R0801
         for j, sentence in enumerate(source):
             role = roles[sentence["from"]]
-            assert role == conv.roles[j % 2]
+            assert role == conv.roles[j % 2], SHAREGPT_ASSERTION_FAILED_ROLE
             if sentence["value"]:
                 conv.append_message(role, sentence["value"])
         yield conv

diff --git a/src/axolotl/prompters.py b/src/axolotl/prompters.py
@@ -260,6 +260,11 @@ def append_message(self, role, message):
         self.messages.append([role, message])
 
 
+SHAREGPT_ASSERTION_FAILED_ROLE = (
+    "Role did not alternate between turns (gpt and human). Please check your data."
+)
+
+
 class ShareGPTPrompter:  # pylint: disable=too-few-public-methods
     """
     A prompter that generates prompts for the ShareGPT
@@ -316,7 +321,7 @@ def build_prompt(self, source) -> Generator[str, None, None]:
         conv.messages = []
         for j, sentence in enumerate(source):
             role = roles[sentence["from"]]
-            assert role == conv.roles[j % 2]
+            assert role == conv.roles[j % 2], SHAREGPT_ASSERTION_FAILED_ROLE
             conv.append_message(role, sentence["value"])
 
         for part in conv.get_prompt():