use pipeline

handler.py

@@ -1,13 +1,16 @@
 from typing import Dict, List, Any
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 class EndpointHandler():
     def __init__(self, path=""):
         # init
-        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)        
+        # load the model
+        tokenizer = AutoTokenizer.from_pretrained("verseAI/vai-GPT-NeoXT-Chat-Base-20B")
+        model = AutoModelForCausalLM.from_pretrained("verseAI/vai-GPT-NeoXT-Chat-Base-20B", device_map="auto", load_in_8bit=True)        
+        # create inference pipeline
+        self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+    def __call__(self, data: Dict[str, Any]) -> List[List[Dict[str, float]]]:
         """
        data args:
             inputs (:obj: `str`)
@@ -18,15 +21,25 @@ class EndpointHandler():
             from transformers import AutoTokenizer, AutoModelForCausalLM
         """
 
-        input = data.pop("inputs", data)
-        print(input)
+        inputs = data.pop("inputs", data)
+        parameters = data.pop("parameters", None)
+        # print(input)
 
-        # infer
+        # pass inputs with all kwargs in data
+        if parameters is not None:
+            prediction = self.pipeline(inputs, **parameters)
+        else:
+            prediction = self.pipeline(inputs)
+        # postprocess the prediction
+        return prediction
+
+        """
         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 {"generated_text": output_str}
+        """