feat: temporarily switch out to 2b model

app.py

@@ -205,7 +205,7 @@ with gr.Blocks(
             gr.HTML("<h2>IBM Granite Guardian 3.0</h2>", elem_classes="title")
             gr.HTML(
                 elem_classes="system-description",
-                value="<p>Granite Guardian models are specialized language models in the Granite family that can detect harms and risks in generative AI systems. They can be used with any large language model to make interactions with generative AI systems safe. Select an example in the left panel to see how the Granite Guardian model evaluates harms and risks in user prompts, assistant responses, and for hallucinations in retrieval-augmented generation. In this demo, we use granite-guardian-3.0-8b.</p>",
+                value="<p>Granite Guardian models are specialized language models in the Granite family that can detect harms and risks in generative AI systems. They can be used with any large language model to make interactions with generative AI systems safe. Select an example in the left panel to see how the Granite Guardian model evaluates harms and risks in user prompts, assistant responses, and for hallucinations in retrieval-augmented generation. In this demo, we use granite-guardian-3.0-2b.</p>",
             )
     with gr.Row(elem_classes="column-gap"):
         with gr.Column(scale=0, elem_classes="no-gap"):

model.py

@@ -23,7 +23,7 @@ logger.debug(f"Inference engine is: '{inference_engine}'")
 if inference_engine == "VLLM":
     device = torch.device("cuda")
 
-    model_path = os.getenv("MODEL_PATH", "ibm-granite/granite-guardian-3.0-8b")
+    model_path = os.getenv("MODEL_PATH", "ibm-granite/granite-guardian-3.0-2b")
     logger.debug(f"model_path is {model_path}")
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     # sampling_params = SamplingParams(temperature=0.0, logprobs=nlogprobs)
@@ -37,10 +37,10 @@ elif inference_engine == "WATSONX":
     )
 
     client.set.default_project(os.getenv("WATSONX_PROJECT_ID"))
-    hf_model_path = "ibm-granite/granite-guardian-3.0-8b"
+    hf_model_path = "ibm-granite/granite-guardian-3.0-2b"
     tokenizer = AutoTokenizer.from_pretrained(hf_model_path)
 
-    model_id = "ibm/granite-guardian-3-8b"  # 8B Model: "ibm/granite-guardian-3-8b"
+    model_id = "ibm/granite-guardian-3-2b"  # 2b Model: "ibm/granite-guardian-3-2b"
     model = ModelInference(model_id=model_id, api_client=client)
 
 
@@ -48,13 +48,14 @@ def parse_output(output, input_len):
     label, prob_of_risk = None, None
     if nlogprobs > 0:
 
-        list_index_logprobs_i = [torch.topk(token_i, k=nlogprobs, largest=True, sorted=True)
-                                 for token_i in list(output.scores)[:-1]]
+        list_index_logprobs_i = [
+            torch.topk(token_i, k=nlogprobs, largest=True, sorted=True) for token_i in list(output.scores)[:-1]
+        ]
         if list_index_logprobs_i is not None:
             prob = get_probablities(list_index_logprobs_i)
             prob_of_risk = prob[1]
 
-    res = tokenizer.decode(output.sequences[:,input_len:][0],skip_special_tokens=True).strip()
+    res = tokenizer.decode(output.sequences[:, input_len:][0], skip_special_tokens=True).strip()
     if risky_token.lower() == res.lower():
         label = risky_token
     elif safe_token.lower() == res.lower():
@@ -64,6 +65,7 @@ def parse_output(output, input_len):
 
     return label, prob_of_risk.item()
 
+
 def get_probablities(logprobs):
     safe_token_prob = 1e-50
     unsafe_token_prob = 1e-50
@@ -75,9 +77,7 @@ def get_probablities(logprobs):
             if decoded_token.strip().lower() == risky_token.lower():
                 unsafe_token_prob += math.exp(logprob)
 
-    probabilities = torch.softmax(
-        torch.tensor([math.log(safe_token_prob), math.log(unsafe_token_prob)]), dim=0
-    )
+    probabilities = torch.softmax(torch.tensor([math.log(safe_token_prob), math.log(unsafe_token_prob)]), dim=0)
 
     return probabilities
 
@@ -87,6 +87,7 @@ def softmax(values):
     total = sum(exp_values)
     return [v / total for v in exp_values]
 
+
 def get_probablities_watsonx(top_tokens_list):
     safe_token_prob = 1e-50
     risky_token_prob = 1e-50
@@ -109,9 +110,9 @@ def get_prompt(messages, criteria_name, tokenize=False, add_generation_prompt=Fa
         guardian_config=guardian_config,
         tokenize=tokenize,
         add_generation_prompt=add_generation_prompt,
-        return_tensors=return_tensors
+        return_tensors=return_tensors,
     )
-    logger.debug(f'prompt is\n{prompt}')
+    logger.debug(f"prompt is\n{prompt}")
     return prompt
 
 
@@ -166,18 +167,15 @@ def generate_text(messages, criteria_name):
 
     elif inference_engine == "VLLM":
         # input_ids = get_prompt(
-        #     messages=messages, 
-        #     criteria_name=criteria_name, 
+        #     messages=messages,
+        #     criteria_name=criteria_name,
         #     tokenize=True,
         #     add_generation_prompt=True,
         #     return_tensors="pt").to(model.device)
         guardian_config = {"risk_name": criteria_name if criteria_name != "general_harm" else "harm"}
-        logger.debug(f'guardian_config is: {guardian_config}')
+        logger.debug(f"guardian_config is: {guardian_config}")
         input_ids = tokenizer.apply_chat_template(
-            messages,
-            guardian_config=guardian_config,
-            add_generation_prompt=True,
-            return_tensors='pt'
+            messages, guardian_config=guardian_config, add_generation_prompt=True, return_tensors="pt"
         ).to(model.device)
         logger.debug(f"input_ids are: {input_ids}")
         input_len = input_ids.shape[1]
@@ -190,7 +188,8 @@ def generate_text(messages, criteria_name):
                 do_sample=False,
                 max_new_tokens=nlogprobs,
                 return_dict_in_generate=True,
-                output_scores=True,)
+                output_scores=True,
+            )
             logger.debug(f"model output is:\n{output}")
 
             label, prob_of_risk = parse_output(output, input_len)