Update app.py

app.py

@@ -9,9 +9,6 @@ import json
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
 # Set up the Hugging Face Dataset Saver
-HF_TOKEN = os.getenv("HF_TOKEN")
-DATASET_NAME = "CannaTech/Flagged"
-
 class CustomFlaggingCallback(gr.FlaggingCallback):
     def __init__(self, hf_token, dataset_name):
         self.hf_writer = gr.HuggingFaceDatasetSaver(hf_token, dataset_name)
@@ -19,12 +16,16 @@ class CustomFlaggingCallback(gr.FlaggingCallback):
     def setup(self, interface):
         pass
 
-    def __call__(self, input_data, output_data):
+    def flag(self, flagged_data):
         # Custom logic to handle flagged data
         # Example: Save flagged data to Hugging Face Dataset without writing the output to JSON file
-        self.hf_writer.add_data(input_data)
+        self.hf_writer.add_data(flagged_data["input"])
         self.hf_writer.flush()
 
+# Define your Hugging Face API token and dataset name      
+HF_TOKEN = os.getenv("HF_TOKEN")
+DATASET_NAME = "CannaTech/Flagged"
+
 # Create an instance of the custom flagging callback
 custom_callback = CustomFlaggingCallback(HF_TOKEN, DATASET_NAME)