change return type

handler.py

@@ -7,7 +7,7 @@ class EndpointHandler():
         self.tokenizer = AutoTokenizer.from_pretrained("verseAI/vai-GPT-NeoXT-Chat-Base-20B")
         self.model = AutoModelForCausalLM.from_pretrained("verseAI/vai-GPT-NeoXT-Chat-Base-20B", device_map="auto", load_in_8bit=True)        
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         """
        data args:
             inputs (:obj: `str`)
@@ -17,10 +17,16 @@ class EndpointHandler():
 
             from transformers import AutoTokenizer, AutoModelForCausalLM
         """
+
+        input = data.pop("inputs", data)
+        print(input)
+
         # infer
         inputs = self.tokenizer("<human>: Hello!\n<bot>:", return_tensors='pt').to(self.model.device)
         outputs = self.model.generate(**inputs, max_new_tokens=10, do_sample=True, temperature=0.8)
         output_str = self.tokenizer.decode(outputs[0])
         print(output_str)
-        return output_str
+        # return output_str
+        return {"generated_text": output_str}
+