restructure files

requirements.txt

@@ -1,4 +1,4 @@
-gradio
+gradio>=5.0
 python-dotenv
 tqdm
 jinja2

app.py → src/app.py

@@ -1,7 +1,8 @@
 import gradio as gr
 from dotenv import load_dotenv
 
-from utils import get_result_description, to_title_case, get_prompt_from_test_case, to_snake_case
+from utils import get_result_description, to_title_case, get_prompt_from_test_case, to_snake_case, load_command_line_args
+load_command_line_args()
 load_dotenv()
 import json
 from model import generate_text
@@ -99,15 +100,15 @@ head_style = """
 """
 
 with gr.Blocks(
-    title='Granite Guardian',
-    theme=gr.themes.Soft(
-        primary_hue=ibm_blue,
+        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='styles.css') as demo:
-    
+        ),
+        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',
@@ -116,7 +117,7 @@ with gr.Blocks(
 
     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='title-row'):
+    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 allow you to detect harms and risks in generative AI systems. The Granite Guardian models can be used with any other large language models to make interactions with generative AI systems safe. Select an example in the left panel to see how the 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>')
@@ -152,9 +153,17 @@ with gr.Blocks(
             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')
+            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')
             
-            result_text = gr.Textbox(label='Result', elem_classes=['result-text', 'read-only', 'input-box'], visible=False, value='')
+            result_text = gr.Textbox(
+                label='Result',
+                elem_classes=['result-text', 'read-only', 'input-box'],
+                visible=False,
+                value='')
 
             with Modal(visible=False, elem_classes='modal') as modal:
                 prompt = gr.Markdown('')
@@ -173,11 +182,20 @@ with gr.Blocks(
             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(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}) \
-            .then(change_button_color, None, [v for c in catalog_buttons.values() for v in c.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')

model.py → src/model.py

@@ -12,7 +12,7 @@ if not mock_model_call:
     import torch
     from vllm import LLM, SamplingParams
     from transformers import AutoTokenizer
-    model_path = os.getenv('MODEL_PATH')#"granite-guardian-3b-pipecleaner-r241024a"
+    model_path = os.getenv('MODEL_PATH') #"granite-guardian-3b-pipecleaner-r241024a"
     sampling_params = SamplingParams(temperature=0.0, logprobs=nlogprobs)
     model = LLM(model=model_path, tensor_parallel_size=1)
     tokenizer = AutoTokenizer.from_pretrained(model_path)

styles.css → src/styles.css

@@ -1,4 +1,4 @@
-.title-row {
+.header-row {
     margin-bottom: 0.75rem;
 }
 
@@ -7,6 +7,10 @@
     margin-bottom: -0.25rem;
 }
 
+.title div {
+    overflow-y: hidden;
+}
+
 .title h2 {
     font-weight: 600;
     font-size: 30px;
@@ -29,7 +33,6 @@
     justify-content: flex-start;
     background-color: transparent;
     box-shadow: none;
-
 }
 
 .selected {

utils.py → src/utils.py

@@ -1,5 +1,7 @@
 import json
 from jinja2 import Template
+import argparse
+import os
 
 with open('prompt_templates.json', mode='r', encoding="utf-8") as f:
     prompt_templates = json.load(f)
@@ -60,3 +62,15 @@ def to_title_case(input_string):
 
 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")
+
+    # Parse arguments
+    args = parser.parse_args()
+
+    # Store the argument in an environment variable
+    if args.model_path is not None:
+        os.environ["MODEL_PATH"] = args.model_path