Changes

app.py

@@ -21,7 +21,8 @@ def llama_guard_classify(conv_prefix, response):
         api_key=LLAMA_API_KEY
     )
     conv = conv_prefix
-    conv.append(response)
+    if response:
+        conv.append(response)
     output = client.chat.completions.create(
         model=model_name,
         messages=conv,
@@ -88,6 +89,8 @@ def process_inputs(category,conv_prefix, response_content):
     output = classify_prompt(category,conv_prefix, response)
     if category=='response':
         llama_output = llama_guard_classify(conv_prefix, response)
+    elif category=='prompt':
+        llama_output = llama_guard_classify(conv_prefix, None)
     else:
         llama_output = 'NA'
     return output,llama_output
@@ -208,7 +211,7 @@ with gr.Blocks(css=dark_css) as demo:
     # Two text outputs, placed side by side for model outputs
     with gr.Row():
         with gr.Column():
-            collinear_output = gr.Textbox(label="Collinear Guard (~3B) Output", lines=3)
+            collinear_output = gr.Textbox(label="Collinear Guard Nano(~3B) Output", lines=3)
         with gr.Column():
             llama_output = gr.Textbox(label="LLaMA-Guard 3 (8B) Output", lines=3)