fix: deployment

.flake8

@@ -0,0 +1,5 @@
+[flake8]
+max-line-length = 120
+
+select = C,E,F,W,B,B950
+extend-ignore = E501,E203,W503

requirements.txt

@@ -1,10 +1,6 @@
-gradio>=4,<5
+gradio_modal
 python-dotenv
-tqdm
-jinja2
-ibm_watsonx_ai
 transformers
-gradio_modal
-spaces
-torch
-vllm
+accelerate
+ibm_watsonx_ai
+vllm

src/app.py

@@ -1,121 +1,150 @@
+import json
+import os
+
 import gradio as gr
 from dotenv import load_dotenv
+from gradio_modal import Modal
+
+from logger import logger
+from model import generate_text, get_prompt
+from utils import (
+    get_messages,
+    get_result_description,
+    load_command_line_args,
+    to_snake_case,
+    to_title_case,
+)
 
-from utils import get_result_description, to_title_case, to_snake_case, load_command_line_args, get_messages
 load_command_line_args()
 load_dotenv()
-import json
-from model import generate_text, get_prompt
-from logger import logger
-import os 
-from gradio_modal import Modal
 
 catalog = {}
 
-with open('catalog.json') as f:
-    logger.debug('Loading catalog from json.')
+with open("catalog.json") as f:
+    logger.debug("Loading catalog from json.")
     catalog = json.load(f)
 
+
 def update_selected_test_case(button_name, state: gr.State, event: gr.EventData):
-    target_sub_catalog_name, target_test_case_name = event.target.elem_id.split('---')
-    state['selected_sub_catalog'] = target_sub_catalog_name
-    state['selected_criteria_name'] = target_test_case_name
-    state['selected_test_case'] = [t for sub_catalog_name, sub_catalog in catalog.items() for t in sub_catalog if t['name'] == to_snake_case(button_name) and to_snake_case(sub_catalog_name) == target_sub_catalog_name][0]
+    target_sub_catalog_name, target_test_case_name = event.target.elem_id.split("---")
+    state["selected_sub_catalog"] = target_sub_catalog_name
+    state["selected_criteria_name"] = target_test_case_name
+    state["selected_test_case"] = [
+        t
+        for sub_catalog_name, sub_catalog in catalog.items()
+        for t in sub_catalog
+        if t["name"] == to_snake_case(button_name) and to_snake_case(sub_catalog_name) == target_sub_catalog_name
+    ][0]
     return state
 
+
 def on_test_case_click(state: gr.State):
-    selected_sub_catalog = state['selected_sub_catalog']
-    selected_criteria_name = state['selected_criteria_name']
-    selected_test_case = state['selected_test_case']
+    selected_sub_catalog = state["selected_sub_catalog"]
+    selected_criteria_name = state["selected_criteria_name"]
+    selected_test_case = state["selected_test_case"]
 
     logger.debug(f'Changing to test case "{selected_criteria_name}" from catalog "{selected_sub_catalog}".')
 
-    is_context_iditable = selected_criteria_name == 'context_relevance'
-    is_user_message_editable = selected_sub_catalog == 'harmful_content_in_user_prompt'
-    is_assistant_message_editable = selected_sub_catalog == 'harmful_content_in_assistant_response' or \
-                                    selected_criteria_name == 'groundedness' or \
-                                    selected_criteria_name == 'answer_relevance'
+    is_context_iditable = selected_criteria_name == "context_relevance"
+    is_user_message_editable = selected_sub_catalog == "harmful_content_in_user_prompt"
+    is_assistant_message_editable = (
+        selected_sub_catalog == "harmful_content_in_assistant_response"
+        or selected_criteria_name == "groundedness"
+        or selected_criteria_name == "answer_relevance"
+    )
     return {
         test_case_name: f'<h2>{to_title_case(selected_test_case["name"])}</h2>',
-        criteria: selected_test_case['criteria'],
-        context: gr.update(
-                value=selected_test_case['context'],
-                interactive=True,
-                visible=True,
-                elem_classes=['input-box']
-            ) if is_context_iditable else gr.update(
-                visible=selected_test_case['context'] is not None,
-                value=selected_test_case['context'],
-                interactive=False,
-                elem_classes=['read-only', 'input-box']
-            ),
-        user_message: gr.update(
-                value=selected_test_case['user_message'],
-                visible=True,
-                interactive=True,
-                elem_classes=['input-box']
-            ) if is_user_message_editable else gr.update(
-                value=selected_test_case['user_message'],
+        criteria: selected_test_case["criteria"],
+        context: (
+            gr.update(value=selected_test_case["context"], interactive=True, visible=True, elem_classes=["input-box"])
+            if is_context_iditable
+            else gr.update(
+                visible=selected_test_case["context"] is not None,
+                value=selected_test_case["context"],
                 interactive=False,
-                elem_classes=['read-only', 'input-box']
-            ),
-        assistant_message: gr.update(
-                value=selected_test_case['assistant_message'],
+                elem_classes=["read-only", "input-box"],
+            )
+        ),
+        user_message: (
+            gr.update(
+                value=selected_test_case["user_message"], visible=True, interactive=True, elem_classes=["input-box"]
+            )
+            if is_user_message_editable
+            else gr.update(
+                value=selected_test_case["user_message"], interactive=False, elem_classes=["read-only", "input-box"]
+            )
+        ),
+        assistant_message: (
+            gr.update(
+                value=selected_test_case["assistant_message"],
                 visible=True,
                 interactive=True,
-                elem_classes=['input-box']
-            ) if is_assistant_message_editable else gr.update(
-                visible=selected_test_case['assistant_message'] is not None,
-                value=selected_test_case['assistant_message'],
+                elem_classes=["input-box"],
+            )
+            if is_assistant_message_editable
+            else gr.update(
+                visible=selected_test_case["assistant_message"] is not None,
+                value=selected_test_case["assistant_message"],
                 interactive=False,
-                elem_classes=['read-only', 'input-box']
-            ),
-        result_text: gr.update(visible=False, value='')
+                elem_classes=["read-only", "input-box"],
+            )
+        ),
+        result_text: gr.update(visible=False, value=""),
     }
 
+
 def change_button_color(event: gr.EventData):
-    return [gr.update(
-        elem_classes=['catalog-button', 'selected']
-    ) if v.elem_id == event.target.elem_id else gr.update(
-        elem_classes=['catalog-button']
-    ) for c in catalog_buttons.values() for v in c.values()]
+    return [
+        (
+            gr.update(elem_classes=["catalog-button", "selected"])
+            if v.elem_id == event.target.elem_id
+            else gr.update(elem_classes=["catalog-button"])
+        )
+        for c in catalog_buttons.values()
+        for v in c.values()
+    ]
+
 
 def on_submit(criteria, context, user_message, assistant_message, state):
-    criteria_name = state['selected_criteria_name']
+    criteria_name = state["selected_criteria_name"]
     test_case = {
-        'name': criteria_name,
-        'criteria': criteria,
-        'context': context,
-        'user_message': user_message,
-        'assistant_message': assistant_message
+        "name": criteria_name,
+        "criteria": criteria,
+        "context": context,
+        "user_message": user_message,
+        "assistant_message": assistant_message,
     }
 
-    messages = get_messages(test_case=test_case, sub_catalog_name=state['selected_sub_catalog'])
-    
-    logger.debug(f"Starting evaluation for subcatelog {state['selected_sub_catalog']} and criteria name {state['selected_criteria_name']}")
-    
-    result_label = generate_text(messages=messages, criteria_name=criteria_name)['assessment'] # Yes or No
+    messages = get_messages(test_case=test_case, sub_catalog_name=state["selected_sub_catalog"])
+
+    logger.debug(
+        f"Starting evaluation for subcatelog {state['selected_sub_catalog']} and criteria name {state['selected_criteria_name']}"
+    )
+
+    result_label = generate_text(messages=messages, criteria_name=criteria_name)["assessment"]  # Yes or No
 
     html_str = f"<p>{get_result_description(state['selected_sub_catalog'], state['selected_criteria_name'])} <strong>{result_label}</strong></p>"
     # html_str = f"{get_result_description(state['selected_sub_catalog'], state['selected_criteria_name'])} {result_label}"
     return gr.update(value=html_str)
 
+
 def on_show_prompt_click(criteria, context, user_message, assistant_message, state):
-    criteria_name = state['selected_criteria_name']
+    criteria_name = state["selected_criteria_name"]
     test_case = {
-        'name': criteria_name,
-        'criteria': criteria,
-        'context': context,
-        'user_message': user_message,
-        'assistant_message': assistant_message
+        "name": criteria_name,
+        "criteria": criteria,
+        "context": context,
+        "user_message": user_message,
+        "assistant_message": assistant_message,
     }
 
-    messages = get_messages(test_case=test_case, sub_catalog_name=state['selected_sub_catalog'])
+    messages = get_messages(test_case=test_case, sub_catalog_name=state["selected_sub_catalog"])
     prompt = get_prompt(messages, criteria_name)
     print(prompt)
-    prompt = prompt.replace('<', '&lt;').replace('>', '&gt;').replace('\\n', '<br>')
+    prompt = prompt.replace("<", "&lt;").replace(">", "&gt;").replace("\\n", "<br>")
     return gr.Markdown(prompt)
 
+
 ibm_blue = gr.themes.Color(
     name="ibm-blue",
     c50="#eff6ff",
@@ -128,7 +157,7 @@ ibm_blue = gr.themes.Color(
     c700="#1d4ed8",
     c800="#1e40af",
     c900="#1e3a8a",
-    c950="#1d3660"
+    c950="#1d3660",
 )
 
 head_style = """
@@ -149,107 +178,139 @@ head_style = """
 """
 
 with gr.Blocks(
-        title='Granite Guardian',
-        theme=gr.themes.Soft(
-            primary_hue=ibm_blue,
-        font=[gr.themes.GoogleFont("IBM Plex Sans"), gr.themes.GoogleFont('Source Sans 3')],
-        ),
-        head=head_style,
-        fill_width=False,
-        css=os.path.join(os.path.dirname(os.path.abspath(__file__)), './styles.css')
-    ) as demo:
-
-    state = gr.State(value={
-        'selected_sub_catalog': 'harmful_content_in_user_prompt',
-        'selected_criteria_name': 'general_harm'
-    })
-
-    starting_test_case = [t for sub_catalog_name, sub_catalog in catalog.items() for t in sub_catalog if t['name'] == state.value['selected_criteria_name'] and sub_catalog_name == state.value['selected_sub_catalog']][0]
-    
-    with gr.Row(elem_classes='header-row'):
+    title="Granite Guardian",
+    theme=gr.themes.Soft(
+        primary_hue=ibm_blue,
+        font=[gr.themes.GoogleFont("IBM Plex Sans"), gr.themes.GoogleFont("Source Sans 3")],
+    ),
+    head=head_style,
+    fill_width=False,
+    css=os.path.join(os.path.dirname(os.path.abspath(__file__)), "./styles.css"),
+) as demo:
+
+    state = gr.State(
+        value={"selected_sub_catalog": "harmful_content_in_user_prompt", "selected_criteria_name": "general_harm"}
+    )
+
+    starting_test_case = [
+        t
+        for sub_catalog_name, sub_catalog in catalog.items()
+        for t in sub_catalog
+        if t["name"] == state.value["selected_criteria_name"]
+        and sub_catalog_name == state.value["selected_sub_catalog"]
+    ][0]
+
+    with gr.Row(elem_classes="header-row"):
         with gr.Column(scale=4):
-            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>')
-    with gr.Row(elem_classes='column-gap'):
-        with gr.Column(scale=0, elem_classes='no-gap'):
-            title_display_left = gr.HTML("<h2>Harms & Risks</h2>", elem_classes=['subtitle', 'subtitle-harms'])
+            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>",
+            )
+    with gr.Row(elem_classes="column-gap"):
+        with gr.Column(scale=0, elem_classes="no-gap"):
+            title_display_left = gr.HTML("<h2>Harms & Risks</h2>", elem_classes=["subtitle", "subtitle-harms"])
             accordions = []
-            catalog_buttons: dict[str,dict[str,gr.Button]] = {}
+            catalog_buttons: dict[str, dict[str, gr.Button]] = {}
             for i, (sub_catalog_name, sub_catalog) in enumerate(catalog.items()):
-                with gr.Accordion(to_title_case(sub_catalog_name), open=(i==0), elem_classes='accordion') as accordion:
+                with gr.Accordion(
+                    to_title_case(sub_catalog_name), open=(i == 0), elem_classes="accordion"
+                ) as accordion:
                     for test_case in sub_catalog:
-                        elem_classes=['catalog-button']
-                        elem_id=f"{sub_catalog_name}---{test_case['name']}"
+                        elem_classes = ["catalog-button"]
+                        elem_id = f"{sub_catalog_name}---{test_case['name']}"
                         if starting_test_case == test_case:
-                            elem_classes.append('selected')
+                            elem_classes.append("selected")
 
-                        if not sub_catalog_name in catalog_buttons:
+                        if sub_catalog_name not in catalog_buttons:
                             catalog_buttons[sub_catalog_name] = {}
 
-                        catalog_buttons[sub_catalog_name][test_case['name']] = \
-                            gr.Button(to_title_case(test_case['name']), elem_classes=elem_classes, variant='secondary', size='sm', elem_id=elem_id)
-                        
+                        catalog_buttons[sub_catalog_name][test_case["name"]] = gr.Button(
+                            to_title_case(test_case["name"]),
+                            elem_classes=elem_classes,
+                            variant="secondary",
+                            size="sm",
+                            elem_id=elem_id,
+                        )
+
                     accordions.append(accordion)
 
         with gr.Column(visible=True, scale=1) as test_case_content:
-            with gr.Row(elem_classes='no-stretch'):
-                test_case_name = gr.HTML(f'<h2>{to_title_case(starting_test_case["name"])}</h2>', elem_classes='subtitle')
-                show_propt_button = gr.Button('Show prompt', size='sm', scale=0, min_width=110)
+            with gr.Row(elem_classes="no-stretch"):
+                test_case_name = gr.HTML(
+                    f'<h2>{to_title_case(starting_test_case["name"])}</h2>', elem_classes="subtitle"
+                )
+                show_propt_button = gr.Button("Show prompt", size="sm", scale=0, min_width=110)
 
-            criteria = gr.Textbox(label="Evaluation Criteria", lines=3, interactive=False, value=starting_test_case['criteria'], elem_classes=['read-only', 'input-box', 'margin-bottom'])
-            gr.HTML(elem_classes=['block', 'content-gap'])
-            context = gr.Textbox(label="Context", lines=3, interactive=True, value=starting_test_case['context'], visible=False, elem_classes=['input-box'])
-            user_message = gr.Textbox(label="User Prompt", lines=3, interactive=True, value=starting_test_case['user_message'], elem_classes=['input-box'])
-            assistant_message = gr.Textbox(label="Assistant Response", lines=3, interactive=True, visible=False, value=starting_test_case['assistant_message'], elem_classes=['input-box'])
+            criteria = gr.Textbox(
+                label="Evaluation Criteria",
+                lines=3,
+                interactive=False,
+                value=starting_test_case["criteria"],
+                elem_classes=["read-only", "input-box", "margin-bottom"],
+            )
+            gr.HTML(elem_classes=["block", "content-gap"])
+            context = gr.Textbox(
+                label="Context",
+                lines=3,
+                interactive=True,
+                value=starting_test_case["context"],
+                visible=False,
+                elem_classes=["input-box"],
+            )
+            user_message = gr.Textbox(
+                label="User Prompt",
+                lines=3,
+                interactive=True,
+                value=starting_test_case["user_message"],
+                elem_classes=["input-box"],
+            )
+            assistant_message = gr.Textbox(
+                label="Assistant Response",
+                lines=3,
+                interactive=True,
+                visible=False,
+                value=starting_test_case["assistant_message"],
+                elem_classes=["input-box"],
+            )
 
             submit_button = gr.Button(
-                "Evaluate", 
-                variant='primary',
-                icon=os.path.join(os.path.dirname(os.path.abspath(__file__)), 'send-white.png'),
-                elem_classes='submit-button')
-            
+                "Evaluate",
+                variant="primary",
+                icon=os.path.join(os.path.dirname(os.path.abspath(__file__)), "send-white.png"),
+                elem_classes="submit-button",
+            )
+
             # result_text = gr.HTML(label='Result', elem_classes=['result-text', 'read-only', 'input-box'], visible=False, value='')
             result_text = gr.HTML(
-                label='Result',
-                elem_classes=['result-root'],
-                show_label=True,
-                visible=False,
-                value='')
+                label="Result", elem_classes=["result-root"], show_label=True, visible=False, value=""
+            )
 
-            with Modal(visible=False, elem_classes='modal') as modal:
-                prompt = gr.Markdown('')
+            with Modal(visible=False, elem_classes="modal") as modal:
+                prompt = gr.Markdown("")
 
-    
     ### events
 
     show_propt_button.click(
-        on_show_prompt_click,
-        inputs=[criteria, context, user_message, assistant_message, state],
-        outputs=prompt
+        on_show_prompt_click, inputs=[criteria, context, user_message, assistant_message, state], outputs=prompt
     ).then(lambda: gr.update(visible=True), None, modal)
 
-    submit_button \
-        .click(lambda: gr.update(visible=True, value=''), None, result_text) \
-        .then(
-            on_submit,
-            inputs=[criteria, context, user_message, assistant_message, state],
-            outputs=[result_text],
-            scroll_to_output=True
+    submit_button.click(lambda: gr.update(visible=True, value=""), None, result_text).then(
+        on_submit,
+        inputs=[criteria, context, user_message, assistant_message, state],
+        outputs=[result_text],
+        scroll_to_output=True,
+    )
+
+    for button in [
+        t for sub_catalog_name, sub_catalog_buttons in catalog_buttons.items() for t in sub_catalog_buttons.values()
+    ]:
+        button.click(
+            change_button_color, inputs=None, outputs=[v for c in catalog_buttons.values() for v in c.values()]
+        ).then(update_selected_test_case, inputs=[button, state], outputs=[state]).then(
+            on_test_case_click,
+            inputs=state,
+            outputs={test_case_name, criteria, context, user_message, assistant_message, result_text},
         )
-      
-    for button in [t for sub_catalog_name, sub_catalog_buttons in catalog_buttons.items() for t in sub_catalog_buttons.values()]:
-        button \
-            .click(
-                change_button_color,
-                inputs=None,
-                outputs=[v for c in catalog_buttons.values() for v in c.values()]) \
-            .then(
-                update_selected_test_case,
-                inputs=[button, state],
-                outputs=[state]) \
-            .then(
-                on_test_case_click,
-                inputs=state,
-                outputs={test_case_name, criteria, context, user_message, assistant_message, result_text})
-    
-demo.launch(server_name='0.0.0.0')
+
+demo.launch(server_name="0.0.0.0")

src/logger.py

@@ -1,12 +1,14 @@
 import logging
 
-logger = logging.getLogger('demo')
+logger = logging.getLogger("demo")
 logger.setLevel(logging.DEBUG)
 
 stream_handler = logging.StreamHandler()
 stream_handler.setLevel(logging.DEBUG)
 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 = logging.FileHandler("logs.txt")
+file_handler.setFormatter(
+    logging.Formatter("%(asctime)s - %(filename)s:%(lineno)d - %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
+)
 logger.addHandler(file_handler)

src/model.py

@@ -1,46 +1,44 @@
-import os
-from time import time, sleep
-from logger import logger
 import math
 import os
+from time import sleep, time
+
+import spaces
+import torch
 from ibm_watsonx_ai.client import APIClient
 from ibm_watsonx_ai.foundation_models import ModelInference
 from transformers import AutoTokenizer
-import math
-import spaces
+from vllm import LLM, SamplingParams
+
+from logger import logger
 
 safe_token = "No"
 risky_token = "Yes"
 nlogprobs = 5
 
-inference_engine = os.getenv('INFERENCE_ENGINE', 'VLLM')
+inference_engine = os.getenv("INFERENCE_ENGINE", "VLLM")
 logger.debug(f"Inference engine is: '{inference_engine}'")
 
-if inference_engine == 'VLLM':
-    import torch
-    from vllm import LLM, SamplingParams
-    from transformers import AutoTokenizer
-    model_path = os.getenv('MODEL_PATH', 'ibm-granite/granite-guardian-3.0-8b') 
+if inference_engine == "VLLM":
+
+    model_path = os.getenv("MODEL_PATH", "ibm-granite/granite-guardian-3.0-8b")
     logger.debug(f"model_path is {model_path}")
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     sampling_params = SamplingParams(temperature=0.0, logprobs=nlogprobs)
     model = LLM(model=model_path, tensor_parallel_size=1)
 
 elif inference_engine == "WATSONX":
-    client = APIClient(credentials={
-        'api_key': os.getenv('WATSONX_API_KEY'), 
-        'url': 'https://us-south.ml.cloud.ibm.com'})
-    
-    client.set.default_project(os.getenv('WATSONX_PROJECT_ID'))
+    client = APIClient(
+        credentials={"api_key": os.getenv("WATSONX_API_KEY"), "url": "https://us-south.ml.cloud.ibm.com"}
+    )
+
+    client.set.default_project(os.getenv("WATSONX_PROJECT_ID"))
     hf_model_path = "ibm-granite/granite-guardian-3.0-8b"
     tokenizer = AutoTokenizer.from_pretrained(hf_model_path)
 
-    model_id = "ibm/granite-guardian-3-8b" # 8B Model: "ibm/granite-guardian-3-8b"
-    model = ModelInference(
-        model_id=model_id,
-        api_client=client
-    )
-    
+    model_id = "ibm/granite-guardian-3-8b"  # 8B Model: "ibm/granite-guardian-3-8b"
+    model = ModelInference(model_id=model_id, api_client=client)
+
+
 def parse_output(output):
     label, prob = None, None
 
@@ -60,11 +58,13 @@ def parse_output(output):
 
     return label, prob_of_risk.item()
 
+
 def softmax(values):
     exp_values = [math.exp(v) for v in values]
     total = sum(exp_values)
     return [v / total for v in exp_values]
 
+
 def get_probablities(logprobs):
     safe_token_prob = 1e-50
     unsafe_token_prob = 1e-50
@@ -76,59 +76,55 @@ def get_probablities(logprobs):
             if decoded_token.strip().lower() == risky_token.lower():
                 unsafe_token_prob += math.exp(token_prob.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
 
+
 def get_probablities_watsonx(top_tokens_list):
     safe_token_prob = 1e-50
     risky_token_prob = 1e-50
     for top_tokens in top_tokens_list:
         for token in top_tokens:
-            if token['text'].strip().lower() == safe_token.lower():
-                safe_token_prob += math.exp(token['logprob'])
-            if token['text'].strip().lower() == risky_token.lower():
-                risky_token_prob += math.exp(token['logprob'])
+            if token["text"].strip().lower() == safe_token.lower():
+                safe_token_prob += math.exp(token["logprob"])
+            if token["text"].strip().lower() == risky_token.lower():
+                risky_token_prob += math.exp(token["logprob"])
 
     probabilities = softmax([math.log(safe_token_prob), math.log(risky_token_prob)])
 
     return probabilities
 
+
 def get_prompt(messages, criteria_name):
-    guardian_config = {"risk_name": criteria_name if criteria_name != 'general_harm' else 'harm'}
+    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)
+        messages, guardian_config=guardian_config, tokenize=False, add_generation_prompt=True
+    )
+
 
 def generate_tokens(prompt):
     result = model.generate(
         prompt=[prompt],
         params={
-            'decoding_method':'greedy',
-            'max_new_tokens': 20, 
+            "decoding_method": "greedy",
+            "max_new_tokens": 20,
             "temperature": 0,
-            "return_options": {
-                "token_logprobs": True,
-                "generated_tokens": True,
-                "input_text": True,
-                "top_n_tokens": 5
-            }
-        })
-    return result[0]['results'][0]['generated_tokens']
+            "return_options": {"token_logprobs": True, "generated_tokens": True, "input_text": True, "top_n_tokens": 5},
+        },
+    )
+    return result[0]["results"][0]["generated_tokens"]
+
 
 def parse_output_watsonx(generated_tokens_list):
     label, prob_of_risk = None, None
 
     if nlogprobs > 0:
-        top_tokens_list = [generated_tokens['top_tokens'] for generated_tokens in generated_tokens_list]
+        top_tokens_list = [generated_tokens["top_tokens"] for generated_tokens in generated_tokens_list]
         prob = get_probablities_watsonx(top_tokens_list)
         prob_of_risk = prob[1]
 
-    res = next(iter(generated_tokens_list))['text'].strip()
+    res = next(iter(generated_tokens_list))["text"].strip()
 
     if risky_token.lower() == res.lower():
         label = risky_token
@@ -139,25 +135,26 @@ def parse_output_watsonx(generated_tokens_list):
 
     return label, prob_of_risk
 
+
 @spaces.GPU
 def generate_text(messages, criteria_name):
-    logger.debug(f'Messages used to create the prompt are: \n{messages}')
-    
+    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}')
+    logger.debug(f"Prompt is \n{chat}")
 
-    if inference_engine=="MOCK":
-        logger.debug('Returning mocked model result.')
+    if inference_engine == "MOCK":
+        logger.debug("Returning mocked model result.")
         sleep(1)
-        label, prob_of_risk = 'Yes', 0.97
-    
-    elif inference_engine=="WATSONX":
+        label, prob_of_risk = "Yes", 0.97
+
+    elif inference_engine == "WATSONX":
         generated_tokens = generate_tokens(chat)
         label, prob_of_risk = parse_output_watsonx(generated_tokens)
 
-    elif inference_engine=="VLLM":
+    elif inference_engine == "VLLM":
         with torch.no_grad():
             output = model.generate(chat, sampling_params, use_tqdm=False)
 
@@ -165,11 +162,11 @@ def generate_text(messages, criteria_name):
     else:
         raise Exception("Environment variable 'INFERENCE_ENGINE' must be one of [WATSONX, MOCK, VLLM]")
 
-    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: \n{label}")
+    logger.debug(f"Model prob_of_risk: \n{prob_of_risk}")
+
     end = time()
     total = end - start
-    logger.debug(f'The evaluation took {total} secs')
+    logger.debug(f"The evaluation took {total} secs")
 
-    return {'assessment': label, 'certainty': prob_of_risk}
+    return {"assessment": label, "certainty": prob_of_risk}

src/utils.py

@@ -1,27 +1,29 @@
 import argparse
 import os
 
-def get_messages(test_case, sub_catalog_name) -> list[dict[str,str]]:
+
+def get_messages(test_case, sub_catalog_name) -> list[dict[str, str]]:
     messages = []
 
-    if sub_catalog_name == 'harmful_content_in_user_prompt':
-        messages.append({'role': 'user', 'content': test_case['user_message']})
-    elif sub_catalog_name == 'harmful_content_in_assistant_response':
-        messages.append({'role': 'user', 'content': test_case['user_message']})
-        messages.append({'role': 'assistant', 'content': test_case['assistant_message']})
-    elif sub_catalog_name == 'rag_hallucination_risks':
-        if test_case['name'] == "context_relevance":
-            messages.append({'role': 'user', 'content': test_case['user_message']})
-            messages.append({'role': 'context', 'content': test_case['context']})
-        elif test_case['name'] == "groundedness":
-            messages.append({'role': 'context', 'content': test_case['context']})
-            messages.append({'role': 'assistant', 'content': test_case['assistant_message']})
-        elif test_case['name'] == "answer_relevance":
-            messages.append({'role': 'user', 'content': test_case['user_message']})
-            messages.append({'role': 'assistant', 'content': test_case['assistant_message']})
-    
+    if sub_catalog_name == "harmful_content_in_user_prompt":
+        messages.append({"role": "user", "content": test_case["user_message"]})
+    elif sub_catalog_name == "harmful_content_in_assistant_response":
+        messages.append({"role": "user", "content": test_case["user_message"]})
+        messages.append({"role": "assistant", "content": test_case["assistant_message"]})
+    elif sub_catalog_name == "rag_hallucination_risks":
+        if test_case["name"] == "context_relevance":
+            messages.append({"role": "user", "content": test_case["user_message"]})
+            messages.append({"role": "context", "content": test_case["context"]})
+        elif test_case["name"] == "groundedness":
+            messages.append({"role": "context", "content": test_case["context"]})
+            messages.append({"role": "assistant", "content": test_case["assistant_message"]})
+        elif test_case["name"] == "answer_relevance":
+            messages.append({"role": "user", "content": test_case["user_message"]})
+            messages.append({"role": "assistant", "content": test_case["assistant_message"]})
+
     return messages
 
+
 def get_result_description(sub_catalog_name, criteria_name):
     evaluated_component = get_evaluated_component(sub_catalog_name, criteria_name)
     messages = {
@@ -33,17 +35,18 @@ def get_result_description(sub_catalog_name, criteria_name):
         "unethical_behavior": f"Does the {evaluated_component} message include actions that violate moral or legal standards?",
         "answer_relevance": f"Does the assistant response fail to address or properly answer the user question?",
         "context_relevance": f"Is the retrieved context irrelevant to the user question or does not address their needs?",
-        "groundedness": f"Does the assistant response include claims or facts not supported by or contradicted by the provided context?"
+        "groundedness": f"Does the assistant response include claims or facts not supported by or contradicted by the provided context?",
     }
     return messages[criteria_name]
 
+
 def get_evaluated_component(sub_catalog_name, criteria_name):
     component = None
-    if sub_catalog_name == 'harmful_content_in_user_prompt':
+    if sub_catalog_name == "harmful_content_in_user_prompt":
         component = "user"
-    elif sub_catalog_name == 'harmful_content_in_assistant_response':
-        component = 'assistant'
-    elif sub_catalog_name == 'rag_hallucination_risks':
+    elif sub_catalog_name == "harmful_content_in_assistant_response":
+        component = "assistant"
+    elif sub_catalog_name == "rag_hallucination_risks":
         if criteria_name == "context_relevance":
             component = "context"
         elif criteria_name == "groundedness":
@@ -51,20 +54,24 @@ def get_evaluated_component(sub_catalog_name, criteria_name):
         elif criteria_name == "answer_relevance":
             component = "assistant"
     if component is None:
-        raise Exception('Something went wrong getting the evaluated component')
+        raise Exception("Something went wrong getting the evaluated component")
     return component
 
+
 def to_title_case(input_string):
-    if input_string == 'rag_hallucination_risks':
-        return 'RAG Hallucination Risks'
-    return ' '.join(word.capitalize() for word in input_string.split('_'))
+    if input_string == "rag_hallucination_risks":
+        return "RAG Hallucination Risks"
+    return " ".join(word.capitalize() for word in input_string.split("_"))
+
 
 def capitalize_first_word(input_string):
-    return ' '.join(word.capitalize() if i == 0 else word for i, word in enumerate(input_string.split('_')))
+    return " ".join(word.capitalize() if i == 0 else word for i, word in enumerate(input_string.split("_")))
+
 
 def to_snake_case(text):
     return text.lower().replace(" ", "_")
 
+
 def load_command_line_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--model_path", type=str, default=None, help="Path to the model or HF repo")