Update handler.py

handler.py

@@ -31,6 +31,20 @@ class EndpointHandler:
 
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
+        INTRO = "A chat between a curious user and a human like artificial intelligence assistant. The assistant gives helpful, intelligent, detailed, and polite answers to the user's questions."
+        prompt = ""
+
+        # process input
+        inputs = data.pop("inputs", data)
+        parameters = data.pop("parameters", None)
+        chat_history = ' \n '.join(str(x) for x in inputs)
+        prompt = INTRO+'\n ' + chat_history
+
+        # preprocess 
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        inputs = self.tokenizer(prompt+' \n  <assistant>:', return_tensors="pt").to(device)
+        inputs = {k: v.to('cuda') for k, v in inputs.items()}
+
         output = self.inference_model.generate(input_ids=inputs["input_ids"],pad_token_id=self.tokenizer.pad_token_id, max_new_tokens=256, do_sample=True, temperature=0.9, top_p=0.9, repetition_penalty=1.5, early_stopping=True, length_penalty = -0.3, num_beams=5, num_return_sequences=1)
         response_raw = self.tokenizer.batch_decode(output.detach().cpu().numpy(), skip_special_tokens=True)
         response_ls = response_raw[0].split('>>')