Clearer error message for unknown example type (#1202)

* error when example is a string * make error message the same * strings only --------- Co-authored-by: Daniel King <[email protected]>
mosaicml · May 14, 2024 · 8274c6c · 8274c6c
1 parent 0449b60
commit 8274c6c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/llmfoundry/data/finetuning/tasks.py b/llmfoundry/data/finetuning/tasks.py
@@ -147,7 +147,7 @@ def _get_example_type(example: Example) -> ExampleType:
     ):
         return 'prompt_response'
     else:
-        raise UnknownExampleTypeError(example)
+        raise UnknownExampleTypeError(str(example.keys()))
 
 
 def _is_empty_or_nonexistent(dirpath: str) -> bool:

diff --git a/llmfoundry/utils/exceptions.py b/llmfoundry/utils/exceptions.py
@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 """Custom exceptions for the LLMFoundry."""
-from collections.abc import Mapping
 from typing import Any, Dict, List
 
 __all__ = [
@@ -79,13 +78,14 @@ def __init__(
 class UnknownExampleTypeError(KeyError):
     """Error thrown when an unknown example type is used in a task."""
 
-    def __init__(self, example: Mapping) -> None:
-        self.example = example
+    def __init__(self, example_keys: str) -> None:
+        self.example = example_keys
         message = (
-            f'Found keys {example.keys()} in dataset. Unknown example type. For prompt and response '
+            f'Found keys {example_keys} in dataset. Unknown example type. For prompt and response '
             f'finetuning, the valid prompt keys are {ALLOWED_PROMPT_KEYS} and the valid response keys are '
             f'{ALLOWED_RESPONSE_KEYS}. For chat finetuning, the allowed keys are {ALLOWED_MESSAGES_KEYS}'
         )
+
         super().__init__(message)