UnknownExampleType Error Message display only keys

mosaicml · May 30, 2024 · 8e76bee · 8e76bee
1 parent fb9a225
commit 8e76bee
Show file tree

Hide file tree

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

diff --git a/tests/data/test_dataloader.py b/tests/data/test_dataloader.py
@@ -915,7 +915,8 @@ def test_malformed_data(
     if add_unknown_example_type:
         error_context = pytest.raises(
             UnknownExampleTypeError,
-            match=r'.*Unknown example type',
+            match=
+            r'\[("|\')(\w+)("|\')(?:,\s*("|\')(\w+)("|\'))*\].*Unknown example type',
         )
     if add_too_many_example_keys:
         error_context = pytest.raises(