fix: guardian_config is not passed

model.py

@@ -105,7 +105,7 @@ def get_probablities_watsonx(top_tokens_list):
 def get_prompt(messages, criteria_name, return_tensors=None):
     guardian_config = {"risk_name": criteria_name if criteria_name != "general_harm" else "harm"}
     return tokenizer.apply_chat_template(
-        messages, guardian_config=guardian_config, tokenize=False, add_generation_prompt=True, return_tensors=return_tensors
+        messages, guardian_config=guardian_config, tokenize=False, add_generation_prompt=return_tensors is not None, return_tensors=return_tensors
     )
 
 
@@ -159,7 +159,7 @@ def generate_text(messages, criteria_name):
         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)
+        input_ids = get_prompt(messages=messages, criteria_name=criteria_name, 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}")
@@ -172,7 +172,7 @@ 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}")
+            logger.debug(f"model output is:\n{output}")
 
             label, prob_of_risk = parse_output(output, input_len)
             logger.debug(f"label is are: {label}")