Update handler.py

handler.py

@@ -31,33 +31,37 @@ class EndpointHandler():
       Return:
             A :obj:`list` | `dict`: will be serialized and returned
         """
-        context = data.pop("context",None)
-        question = data.pop("question",None)
-        messages = generate_rag_prompt_message(context, question)
-
-        inputs = self.tokenizer(
-        [
-            self.tokenizer.apply_chat_template(messages, tokenize=False)
-        ]*1, return_tensors = "pt").to("cuda")
-        
-        
-        with torch.no_grad():
-            generated_ids = self.model.generate(
-            input_ids=inputs["input_ids"], 
-            attention_mask=inputs["attention_mask"], 
-            generation_config=self.generation_config, **{
-                "temperature": 0.1,
-                "penalty_alpha": 0.6,
-                "min_p": 0.5,
-                "do_sample": True,
-                "repetition_penalty": 1.28,
-                "min_length": 10,
-                "max_new_tokens": 250
-            })
-        
-        generated_text = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True)[0]
+        print(data)
         try:
-            generated_answer = generated_text.split('[/INST]')[1].strip()
-            return json.dumps({"answer": generated_answer})
+            context = data.pop("context",None)
+            question = data.pop("question",None)
+            messages = generate_rag_prompt_message(context, question)
+    
+            inputs = self.tokenizer(
+            [
+                self.tokenizer.apply_chat_template(messages, tokenize=False)
+            ]*1, return_tensors = "pt").to("cuda")
+            
+            
+            with torch.no_grad():
+                generated_ids = self.model.generate(
+                input_ids=inputs["input_ids"], 
+                attention_mask=inputs["attention_mask"], 
+                generation_config=self.generation_config, **{
+                    "temperature": 0.1,
+                    "penalty_alpha": 0.6,
+                    "min_p": 0.5,
+                    "do_sample": True,
+                    "repetition_penalty": 1.28,
+                    "min_length": 10,
+                    "max_new_tokens": 250
+                })
+            
+            generated_text = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True)[0]
+            try:
+                generated_answer = generated_text.split('[/INST]')[1].strip()
+                return json.dumps({"answer": generated_answer})
+            except Exception as e:
+                return json.dumps({"answer": str(e)})
         except Exception as e:
-            return json.dumps({"answer": str(e)})
+            print(e)