Changes for UI

app.py

@@ -111,23 +111,40 @@ def update_sample(sample):
     return conv_prefix,response
 
 
+import gradio as gr
+
 with gr.Blocks() as demo:
-    gr.Markdown("# Safety Classifier")
-    gr.Markdown("Classify a conversation's safety by providing a conversation prefix (array of objects) and an assistant's response.")
-    
+    # Header section with larger, centered title
+    gr.Markdown("<h1 style='text-align: center;'>Safety Classifier</h1>")
+    gr.Markdown(
+        """
+        <p style='text-align: center;'>
+        Classify a conversation's safety by providing a conversation prefix (array of objects) 
+        and an assistant's response.
+        </p>
+        """
+    )
+
+    # Main content: dropdowns and textboxes in organized rows/columns
     with gr.Row():
-        category = gr.Dropdown(
-            ["response","prompt", "refusal"], label="Select Evaluation Type", value='prompt'
-        )
+        with gr.Column(scale=2, min_width=200):
+            category = gr.Dropdown(
+                ["response", "prompt", "refusal"], 
+                label="Select Evaluation Type", 
+                value='prompt'
+            )
 
-        with gr.Column():
             sample_convos = gr.Dropdown(
-["sample_1","sample_2",'sample_3','sample_4','sample_5'], label="Select Sample Convo", value='sample_1'
+                ["sample_1", "sample_2", "sample_3", "sample_4", "sample_5"], 
+                label="Select Sample Convo", 
+                value='sample_1'
             )
+
+        # Conversation Prefix and Assistant Response in a column
+        with gr.Column(scale=2, min_width=500):
             conv_prefix = gr.Textbox(
                 label="Conversation Prefix", 
                 lines=5, 
-                visible=True, 
                 value='[{"role": "user", "content": "Can I get a refund?"}]'
             )
             response = gr.Textbox(
@@ -136,26 +153,35 @@ with gr.Blocks() as demo:
                 label="Assistant Response", 
                 value="No, you don't deserve a refund"
             )
+    
+    # Submit button centered below the inputs
     with gr.Row():
-        submit = gr.Button("Submit")
+        submit = gr.Button("Submit", elem_id="submit-button")
 
+    # Two text outputs, placed side by side for model outputs
     with gr.Row():
-        collinear_output = gr.Textbox(label="Collinear Guard(~3B) Output")
-        llama_output = gr.Textbox(label="LLaMA-Guard 3 (8B) Output")
+        with gr.Column():
+            collinear_output = gr.Textbox(label="Collinear Guard (~3B) Output", lines=3)
+        with gr.Column():
+            llama_output = gr.Textbox(label="LLaMA-Guard 3 (8B) Output", lines=3)
 
+    # Interaction: Update conversation samples
     sample_convos.change(
         fn=update_sample, 
         inputs=[sample_convos], 
         outputs=[conv_prefix, response]
     )
+    
+    # Submit button interaction and dataset update
     submit.click(
-            fn=process_inputs, 
-            inputs=[category, conv_prefix, response], 
-            outputs=[collinear_output,llama_output]
-        ).then(
-            fn=add_to_dataset, 
-            inputs=[conv_prefix, response, llama_output, collinear_output],
-            outputs=[]
-        )
+        fn=process_inputs, 
+        inputs=[category, conv_prefix, response], 
+        outputs=[collinear_output, llama_output]
+    ).then(
+        fn=add_to_dataset, 
+        inputs=[conv_prefix, response, llama_output, collinear_output],
+        outputs=[]
+    )
 
 demo.launch()
+