Update handler.py

handler.py

@@ -1,11 +1,11 @@
 import torch
 from typing import Dict, List, Any
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import logging
 
 # get dtype
 dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
 
-
 class EndpointHandler:
     def __init__(self, path=""):
         # load the model
@@ -23,15 +23,15 @@ class EndpointHandler:
             prediction = self.pipeline(inputs, **parameters)
         else:
             prediction = self.pipeline(inputs)
-        print("---start---")
-        print(prediction)
-        print("---end---")
+        logging.warn("---start---")
+        logging.warn(prediction)
+        logging.warn("---end---")
 
         # ignoring parameters! Default to configs in generation_config.json.
         messages = [{"role": "user", "content": data}]
         response = self.model.chat(self.tokenizer, messages)
-        print("---start chat response---")
-        print(response)
-        print("---end chat response---")
+        logging.warn("---start chat response---")
+        logging.warn(response)
+        logging.warn("---end chat response---")
 
-        return [{'generated_text': response}]
+        return [[{response: 1.0}]]