Changes

app.py

@@ -77,6 +77,7 @@ Now, please output the following as a JSON object:
         return 'unsafe'
 
 async def add_to_dataset(safe_text,unsafe_text,conv_prefix, response,llama_resp,collinear_resp):
+    conv_prefix = json.loads(conv_prefix)
     dataset = load_dataset("collinear-ai/collinear-guard-demo")
     new_row = {'safe_text':safe_text,
             'unsafe_text':unsafe_text,
@@ -105,7 +106,6 @@ def process_inputs(category,safe_text,unsafe_text,conv_prefix, response_content)
     conv_prefix = json.loads(conv_prefix)
     output = classify_prompt(category,safe_text,unsafe_text,conv_prefix, response)
     llama_output = llama_guard_classify(conv_prefix, response)
-    add_to_dataset(safe_text,unsafe_text,conv_prefix,response_content,llama_output,output)
     return output,llama_output
 
 
@@ -169,8 +169,8 @@ with gr.Blocks() as demo:
         submit = gr.Button("Submit")
 
     with gr.Row():
-        collinear_output = gr.Textbox(label="Collinear Guard Output")
-        llama_output = gr.Textbox(label="LLaMA-Guard 3 Output")
+        collinear_output = gr.Textbox(label="Collinear Guard(~3B) Output")
+        llama_output = gr.Textbox(label="LLaMA-Guard 3 (8B) Output")
 
     category.change(
         fn=update_texts, 
@@ -181,6 +181,10 @@ with gr.Blocks() as demo:
             fn=process_inputs, 
             inputs=[category,safe_text, unsafe_text, conv_prefix, response], 
             outputs=[collinear_output,llama_output]
+        ).then(
+            fn=add_to_dataset, 
+            inputs=[safe_text, unsafe_text, conv_prefix, response, llama_output, collinear_output],
+            outputs=[]
         )
 
 demo.launch()