Log model runtime in seconds

model.py

@@ -3,13 +3,20 @@ from time import time, sleep
 from logger import logger
 import math
 
-
-
-
 safe_token = "No"
 unsafe_token = "Yes"
 nlogprobs = 5
 
+mock_model_call = os.getenv('MOCK_MODEL_CALL') == 'true'
+if not mock_model_call:
+    import torch
+    from vllm import LLM, SamplingParams
+    from transformers import AutoTokenizer
+    model_path = os.getenv('MODEL_PATH')#"granite-guardian-3b-pipecleaner-r241024a"
+    sampling_params = SamplingParams(temperature=0.0, logprobs=nlogprobs)
+    model = LLM(model=model_path, tensor_parallel_size=1)
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+
 def parse_output(output):
     label, prob = None, None
 
@@ -46,17 +53,6 @@ def get_probablities(logprobs):
 
     return probabilities
 
-
-mock_model_call = os.getenv('MOCK_MODEL_CALL') == 'true'
-if not mock_model_call:
-    import torch
-    from vllm import LLM, SamplingParams
-    from transformers import AutoTokenizer
-    model_path = os.getenv('MODEL_PATH')#"granite-guardian-3b-pipecleaner-r241024a"
-    sampling_params = SamplingParams(temperature=0.0, logprobs=nlogprobs)
-    model = LLM(model=model_path, tensor_parallel_size=1)
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
-
 def generate_text(prompt):
     logger.debug(f'Prompts content is: \n{prompt["content"]}')
     mock_model_call = os.getenv('MOCK_MODEL_CALL') == 'true'
@@ -72,8 +68,7 @@ def generate_text(prompt):
         with torch.no_grad():
             output = model.generate(tokenized_chat, sampling_params, use_tqdm=False)
 
-
-        predicted_label = output[0].outputs[0].text.strip()
+        # predicted_label = output[0].outputs[0].text.strip()
 
         label, prob_of_risk = parse_output(output[0])
 
@@ -82,6 +77,6 @@ def generate_text(prompt):
         
         end = time()
         total = end - start
-        logger.debug(f'it took {round(total/60, 2)} mins')
+        logger.debug(f'The evaluation took {total} secs')
 
         return {'assessment': label, 'certainty': prob_of_risk}