Changes for output dataset

app.py

@@ -7,6 +7,7 @@ from datasets import load_dataset, Dataset, DatasetDict
 import pandas as pd
 import re
 import requests
+from datetime import datetime
 API_ENDPOINT = "https://api.collinear.ai"
 API_KEY = os.getenv("COLLINEAR_API_KEY")
 HF_TOKEN=os.getenv("HF_TOKEN")
@@ -56,15 +57,15 @@ def classify_prompt(category,conv_prefix, response):
         else:
             return 'Unsafe'
 
-async def add_to_dataset(conv_prefix, response,llama_resp,collinear_resp):
+async def add_to_dataset(category,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':"",
-            'unsafe_text':"",
+    dataset = load_dataset("collinear-ai/collinear-guard-safety-demo")
+    new_row = {'category':category,
             'conv_prefix':conv_prefix,
             'response':response,
             'llama_resp':{"output":llama_resp},
-            'collinear_resp':collinear_resp
+            'collinear_resp':collinear_resp,
+            "timestamp":datetime.now()
         }
     train_dataset = dataset['train']
 
@@ -77,7 +78,7 @@ async def add_to_dataset(conv_prefix, response,llama_resp,collinear_resp):
     updated_dataset = DatasetDict({
             'train': new_train_dataset
     })
-    updated_dataset.push_to_hub("collinear-ai/collinear-guard-demo",token=HF_TOKEN)
+    updated_dataset.push_to_hub("collinear-ai/collinear-guard-safety-demo",token=HF_TOKEN)
 
 
 
@@ -179,7 +180,7 @@ with gr.Blocks() as demo:
         outputs=[collinear_output, llama_output]
     ).then(
         fn=add_to_dataset, 
-        inputs=[conv_prefix, response, llama_output, collinear_output],
+        inputs=[category,conv_prefix, response, llama_output, collinear_output],
         outputs=[]
     )