Update handler.py

handler.py

@@ -2,31 +2,20 @@ from typing import Dict, List, Any
 from peft import PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-class EndpointHandler:
-    def __init__(self, path=""):
-    
+class EndpointHandler():
+    def __init__(self, path="."):
+        self.model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
+        self.model.eval()
         self.device_map = "cuda" # the device to load the model onto
+        self.tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
         
-        llm_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
-        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
 
-        # llm_model = PeftModel.from_pretrained(llm_model, ".")
-        llm_model.eval()
-        self.llm_model = llm_model
-        self.tokenizer = tokenizer
-
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        """
-         data args:
-              inputs (:obj: `str`)
-              date (:obj: `str`)
-        Return:
-              A :obj:`list` | `dict`: will be serialized and returned
-        """
-        # get inputs
-        prompt = data.pop("prompt", "")
-        model_input = self.tokenizer(prompt, return_tensors="pt").to(self.device_map)
+    def __call__(self, inputs: str):
+        if len(inputs) == 0:
+            raise ValueError("prompt cannot be empty")
+        model_input = self.tokenizer(inputs, return_tensors="pt").to(self.device_map)
         output = self.llm_model.generate(input_ids=model_input["input_ids"].to(self.device_map),
                                   use_cache=False,
                                   temperature=0.1, top_k=1, top_p=1.0, repetition_penalty=1.4,