Update handler.py

handler.py

@@ -3,8 +3,18 @@ import torch
 
 class EndpointHandler:
     def __init__(self, path=""):
+        if not path:
+            raise ValueError("A valid model path or name must be provided.")
+        
+        # Load tokenizer and model
         self.tokenizer = AutoTokenizer.from_pretrained(path)
-        self.model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.float16, device_map="auto")
+        self.model = AutoModelForCausalLM.from_pretrained(
+            path, 
+            torch_dtype=torch.float16, 
+            device_map="auto"
+        )
+        
+        # Set up text-generation pipeline
         self.pipe = pipeline(
             "text-generation",
             model=self.model,
@@ -16,14 +26,23 @@ class EndpointHandler:
         )
 
     def __call__(self, data):
+        # Validate input data
+        if not isinstance(data, dict):
+            return {"error": "Input must be a JSON object."}
+        
         prompt = data.get("inputs", "")
         if not prompt:
-            return {"error": "No input provided"}
-
-        # Generate response
-        response = self.pipe(prompt)[0]['generated_text']
-
-        # Remove the original prompt from the response
-        response = response[len(prompt):].strip()
-
-        return {"generated_text": response}
+            return {"error": "No input provided."}
+        
+        try:
+            # Generate response
+            outputs = self.pipe(prompt)
+            if outputs:
+                response = outputs[0]['generated_text']
+                # Remove the original prompt from the response
+                response = response[len(prompt):].strip()
+                return {"generated_text": response}
+            else:
+                return {"error": "No output generated."}
+        except Exception as e:
+            return {"error": str(e)}