Update handler.py

handler.py

@@ -5,9 +5,9 @@ import torch
 
 dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
 
-
 class EndpointHandler:
     def __init__(self, path=""):
+
         tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
         model = AutoModelForCausalLM.from_pretrained(
             path,
@@ -27,10 +27,12 @@ class EndpointHandler:
         self.generation_config = generation_config
 
         self.pipeline = transformers.pipeline(
-            "text-generation", model=model, tokenizer=tokenizer
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer
         )
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        prompt = data.pop("prompt", data)
+        prompt = data.pop("inputs", data)
         result = self.pipeline(prompt, generation_config=self.generation_config)
-        return result
+        return result