Add logs

model.py

@@ -146,12 +146,7 @@ def parse_output_watsonx(generated_tokens_list):
 @spaces.GPU
 def generate_text(messages, criteria_name):
     logger.debug(f"Messages used to create the prompt are: \n{messages}")
-
     start = time()
-
-    chat = get_prompt(messages, criteria_name)
-    logger.debug(f"Prompt is \n{chat}")
-
     if inference_engine == "MOCK":
         logger.debug("Returning mocked model result.")
         sleep(1)
@@ -159,12 +154,15 @@ def generate_text(messages, criteria_name):
 
     elif inference_engine == "WATSONX":
         chat = get_prompt(messages, criteria_name)
+        logger.debug(f"Prompt is \n{chat}")
         generated_tokens = generate_tokens(chat)
         label, prob_of_risk = parse_output_watsonx(generated_tokens)
 
     elif inference_engine == "VLLM":
         input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
+        logger.debug(f"input_ids are: {input_ids}")
         input_len = input_ids.shape[1]
+        logger.debug(f"input_len are: {input_len}")
 
         with torch.no_grad():
             # output = model.generate(chat, sampling_params, use_tqdm=False)
@@ -174,8 +172,11 @@ def generate_text(messages, criteria_name):
                 max_new_tokens=nlogprobs,
                 return_dict_in_generate=True,
                 output_scores=True,)
+            logger.debug(f"model output is are: {output}")
 
             label, prob_of_risk = parse_output(output, input_len)
+            logger.debug(f"label is are: {label}")
+            logger.debug(f"prob_of_risk is are: {prob_of_risk}")
     else:
         raise Exception("Environment variable 'INFERENCE_ENGINE' must be one of [WATSONX, MOCK, VLLM]")