Update handler.py

handler.py

@@ -22,7 +22,7 @@ class EndpointHandler:
         self.model = AutoModelForCausalLM.from_pretrained(
             base_model_name,
             quantization_config=self.bnb_config,
-            device_map="auto",  # Auto selects device to put model on.
+            device_map="auto", 
         )
         self.model.config.use_cache = False
 
@@ -31,23 +31,14 @@ class EndpointHandler:
 
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
-        INTRO = "Below is a conversation between a user and you."
-        END = "Instruction: Write a response appropriate to the conversation."
-        prompt = "<user>:"
-
-        # process input
-        inputs = data.pop("inputs", data)
-        parameters = data.pop("parameters", None)
-
-        prompt = prompt+inputs
-        # preprocess 
-
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-
-        inputs = self.tokenizer(INTRO+'\n '+prompt+'\n '+END +'\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=100, do_sample=True, temperature=0.1, top_p=0.9, repetition_penalty=1.5)
-        reply = self.tokenizer.batch_decode(output.detach().cpu().numpy(), skip_special_tokens=True)
+        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('>>')
+        response_ = response_ls[1].split('<assistant>:')[1]
+        response_ = response_.split('<user>')[0]
+        response_ = response_.split('Instruction:')[0]
+        
+        response_ = response_.replace('\n','')
+        response = '<assistant>:' + response_.strip()
         
-        return [{"generated_reply": reply}]
+        return [{"generated_reply": response}]