try to log

handler.py

@@ -10,8 +10,7 @@ from utils import MultiHeadModel, BertInputBuilder, get_num_words
 import transformers
 from transformers import BertTokenizer, BertForSequenceClassification
 import psutil
-import time
-
+from transformers.utils import logging
 
 transformers.logging.set_verbosity_debug()
 
@@ -244,33 +243,39 @@ class EndpointHandler():
 
         print("Running inference on %d examples..." % transcript.length())
         cpu_percent = psutil.cpu_percent()
-        print(f"CPU Usage before models loaded: {cpu_percent}%")
+        logging.set_verbosity_info()
+        logger = logging.get_logger("transformers")
+        logger.info(f"CPU Usage before models loaded: {cpu_percent}%")
         # Uptake
         uptake_model = UptakeModel(
             self.device, self.tokenizer, self.input_builder)
         uptake_model.run_inference(transcript, min_prev_words=params['uptake_min_num_words'],
                                    uptake_speaker=params.pop("uptake_speaker", None))
         cpu_percent = psutil.cpu_percent()
-        print(f"CPU Usage after model 1 loaded: {cpu_percent}%")
+        logger.info(f"CPU Usage after model 1 loaded: {cpu_percent}%")
         del uptake_model
         cpu_percent = psutil.cpu_percent()
-        print(f"CPU Usage after model 1 deleted: {cpu_percent}%")
+        logger.info(f"CPU Usage after model 1 deleted: {cpu_percent}%")
         # Reasoning
         reasoning_model = ReasoningModel(
             self.device, self.tokenizer, self.input_builder)
         reasoning_model.run_inference(transcript)
         cpu_percent = psutil.cpu_percent()
+        logger.info(f"CPU Usage after model 2 loaded: {cpu_percent}%")
         print(f"CPU Usage after model 2 loaded: {cpu_percent}%")
         del reasoning_model
         cpu_percent = psutil.cpu_percent()
+        logger.info(f"CPU Usage after model 2 deleted: {cpu_percent}%")
         print(f"CPU Usage after model 2 deleted: {cpu_percent}%")
         # Question
         question_model = QuestionModel(
             self.device, self.tokenizer, self.input_builder)
         question_model.run_inference(transcript)
         cpu_percent = psutil.cpu_percent()
+        logger.info(f"CPU Usage after model 3 loaded: {cpu_percent}%")
         print(f"CPU Usage after model 3 loaded: {cpu_percent}%")
         del question_model
         cpu_percent = psutil.cpu_percent()
+        logger.info(f"CPU Usage after model 3 deleted: {cpu_percent}%")
         print(f"CPU Usage after model 3 deleted: {cpu_percent}%")
         return transcript.to_dict()