Merge pull request #29 from philschmid/fix-input-fn

added fix

Author: philschmid

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,
             )