Continue adaption

.env.example

@@ -1,4 +1,4 @@
 MODEL_PATH='ibm-granite/granite-guardian-3.1-8b'
-INFERENCE_ENGINE='VLLM' # one of [WATSONX, MOCK, VLLM]
+INFERENCE_ENGINE='VLLM' # one of [WATSONX, MOCK, TORCH]
 WATSONX_API_KEY=''
 WATSONX_PROJECT_ID=''

src/logger.py

@@ -1,14 +1,15 @@
 import logging
 
-logger = logging.getLogger("demo")
+logger = logging.getLogger("guardian-demo")
 logger.setLevel(logging.DEBUG)
 
+formatter = logging.Formatter("%(asctime)s - %(filename)s:%(lineno)d - %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
+
 stream_handler = logging.StreamHandler()
 stream_handler.setLevel(logging.DEBUG)
+stream_handler.setFormatter(formatter)
 logger.addHandler(stream_handler)
 
 file_handler = logging.FileHandler("logs.txt")
-file_handler.setFormatter(
-    logging.Formatter("%(asctime)s - %(filename)s:%(lineno)d - %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
-)
+file_handler.setFormatter(formatter)
 logger.addHandler(file_handler)

src/model.py

@@ -13,15 +13,15 @@ safe_token = "No"
 risky_token = "Yes"
 nlogprobs = 20
 
-inference_engine = os.getenv("INFERENCE_ENGINE", "VLLM")
+inference_engine = os.getenv("INFERENCE_ENGINE", "TORCH")
 logger.debug(f"Inference engine is: '{inference_engine}'")
 
-if inference_engine == "VLLM":
+if inference_engine == "TORCH":
     import torch
 
     device = torch.device("cpu")
 
-    model_path = os.getenv("MODEL_PATH", "ibm-granite/granite-guardian-3.1-2b")
+    model_path = os.getenv("MODEL_PATH", "ibm-granite/granite-guardian-3.1-8b")
     logger.debug(f"model_path is {model_path}")
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     model = AutoModelForCausalLM.from_pretrained(model_path)
@@ -160,7 +160,6 @@ def get_prompt(messages, criteria_name, tokenize=False, add_generation_prompt=Fa
 
 @spaces.GPU
 def get_guardian_response(messages, criteria_name):
-    logger.debug(f"Messages used to create the prompt are: \n{messages}")
     start = time()
     if inference_engine == "MOCK":
         logger.debug("Returning mocked model result.")
@@ -173,7 +172,7 @@ def get_guardian_response(messages, criteria_name):
         generated_tokens = generate_tokens_watsonx(chat)
         label, prob_of_risk = parse_output_watsonx(generated_tokens)
 
-    elif inference_engine == "VLLM":
+    elif inference_engine == "TORCH":
         input_ids = get_prompt(
             messages=messages,
             criteria_name=criteria_name,
@@ -181,7 +180,7 @@ def get_guardian_response(messages, criteria_name):
             add_generation_prompt=True,
             return_tensors="pt",
         ).to(model.device)
-        logger.debug(f"input_ids are: {input_ids}")
+        # logger.debug(f"input_ids are: {input_ids}")
         input_len = input_ids.shape[1]
         logger.debug(f"input_len is: {input_len}")
 
@@ -194,16 +193,16 @@ def get_guardian_response(messages, criteria_name):
                 return_dict_in_generate=True,
                 output_scores=True,
             )
-            logger.debug(f"model output is:\n{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}")
-            logger.debug(f"prob_of_risk is are: {prob_of_risk}")
+            logger.debug(f"Label is: {label}")
+            logger.debug(f"Prob_of_risk is: {prob_of_risk}")
     else:
-        raise Exception("Environment variable 'INFERENCE_ENGINE' must be one of [WATSONX, MOCK, VLLM]")
+        raise Exception("Environment variable 'INFERENCE_ENGINE' must be one of [WATSONX, MOCK, TORCH]")
 
-    logger.debug(f"Model generated label: \n{label}")
-    logger.debug(f"Model prob_of_risk: \n{prob_of_risk}")
+    logger.debug(f"Model generated label: {label}")
+    logger.debug(f"Model prob_of_risk: {prob_of_risk}")
 
     end = time()
     total = end - start

src/utils.py

@@ -1,4 +1,5 @@
 import argparse
+import json
 import os
 
 
@@ -25,9 +26,9 @@ def get_messages(test_case, sub_catalog_name) -> list[dict[str, str]]:
             messages += create_message("user", test_case["user_message"])
             messages += create_message("assistant", test_case["assistant_message"])
     elif sub_catalog_name == "risks_in_agentic_workflows":
-        messages += create_message("tools", test_case["tools"])
+        messages += create_message("tools", json.loads(test_case["tools"]))
         messages += create_message("user", test_case["user_message"])
-        messages += create_message("assistant", test_case["assistant_message"])
+        messages += create_message("assistant", json.loads(test_case["assistant_message"]))
     return messages