From ca89825773b0a77f2d7dafe6af9b6d0721138e38 Mon Sep 17 00:00:00 2001 From: Milo Cress Date: Mon, 30 Sep 2024 17:29:47 -0400 Subject: [PATCH] Add error to catch more unknown example types --- llmfoundry/data/finetuning/tasks.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/llmfoundry/data/finetuning/tasks.py b/llmfoundry/data/finetuning/tasks.py index e8f6484ef2..ec170873d2 100644 --- a/llmfoundry/data/finetuning/tasks.py +++ b/llmfoundry/data/finetuning/tasks.py @@ -170,10 +170,12 @@ def _is_empty_or_nonexistent(dirpath: str) -> bool: return not os.path.isdir(dirpath) or len(os.listdir(dirpath)) == 0 -def _get_key(dictionary: Mapping[str, Any], allowed_keys: set[str]): +def _get_key(dictionary: Mapping[str, Any], allowed_keys: set[str], error_type): if not isinstance(dictionary, Mapping): raise InvalidExampleTypeError(str(type(dictionary))) desired_keys = allowed_keys.intersection(dictionary.keys()) + if len(desired_keys) == 0: + raise UnknownExampleTypeError(str(set(dictionary.keys()))) return list(desired_keys)[0]