diff --git a/src/sagemaker_huggingface_inference_toolkit/handler_service.py b/src/sagemaker_huggingface_inference_toolkit/handler_service.py index b7d839a..34ccf23 100644 --- a/src/sagemaker_huggingface_inference_toolkit/handler_service.py +++ b/src/sagemaker_huggingface_inference_toolkit/handler_service.py @@ -126,10 +126,11 @@ def preprocess(self, input_data, content_type): """ # raises en error when using zero-shot-classification or table-question-answering, not possible due to nested properties if ( - os.environ["HF_TASK"] == "zero-shot-classification" or os.environ["HF_TASK"] == "table-question-answering" + os.environ.get("HF_TASK", None) == "zero-shot-classification" + or os.environ.get("HF_TASK", None) == "table-question-answering" ) and content_type == content_types.CSV: raise PredictionException( - f"content type {content_type} not support with {os.environ['HF_TASK']}, use different content_type", + f"content type {content_type} not support with {os.environ.get('HF_TASK', 'unknown task')}, use different content_type", 400, )