Update app.py

app.py

@@ -8,27 +8,9 @@ import json
 # Set the OpenAI API key
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
-# Set up the Hugging Face Dataset Saver
-class CustomFlaggingCallback(gr.FlaggingCallback):
-    def __init__(self, hf_token, dataset_name):
-        self.hf_writer = gr.HuggingFaceDatasetSaver(hf_token, dataset_name)
-
-    def setup(self, interface, interface_state):
-        pass
-
-    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(flagged_data["input"])
-        self.hf_writer.flush()
-
-# Define your Hugging Face API token and dataset name      
+# Set up flagging callback function
 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)
-
+hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "CannaTech/Flagged")
         
 # Define the authentication function
 def check_auth(username, password):
@@ -146,7 +128,7 @@ iface = gr.Interface(
     favicon_path="https://assets.bigcartel.com/theme_images/101321509/IMG_6002.png",  # Favicon image URL
     allow_flagging="manual",  # Enable flagging with manual control
     flagging_options=["Incorrect"],
-    flagging_callback=custom_callback
+    flagging_callback=hf_writer
 )
 
 # Launch the interface with authentication