Update app.py

app.py

@@ -8,32 +8,23 @@ import json
 # Set the OpenAI API key
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
-# Define the CustomHuggingFaceDatasetSaver class
-class HuggingFaceDatasetSaver(gr.FlaggingCallback):
-    def setup(self):
-        pass
-
-    def apply(self, flagged_samples):
-        for sample in flagged_samples:
-            input_data = sample["input"]
-            output_data = sample["output"]
-            flag = sample["flag"]
+# Set up the Hugging Face Dataset Saver
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-            # Truncate the output if it exceeds 1000 characters
-            if len(output_data) > 1000:
-                output_data = output_data[:900] + " [Truncated]"
+class CustomFlaggingCallback:
+    def __init__(self, hf_token, dataset_name):
+        self.hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "CannaTech/Flagged")
 
-            # Save the data to Hugging Face dataset
-            # Replace this placeholder logic with your own implementation
-            # Example: Saving the data to a CSV file
-            with open("flagged_data.csv", "a") as f:
-                writer = csv.writer(f)
-                writer.writerow([input_data, output_data, flag])
+    def __call__(self, input_data, output_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.flush()
 
-# Set up the Hugging Face Dataset Saver
-HF_TOKEN = os.getenv("HF_TOKEN")
-hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "CannaTech/Flagged")
+# Create an instance of the custom flagging callback
+custom_callback = CustomFlaggingCallback(HF_TOKEN, "CannaTech/Flagged")
 
+        
 # Define the authentication function
 def check_auth(username, password):
     # Get the credentials from the environment variable
@@ -150,7 +141,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=hf_writer
+    flagging_callback=custom_callback
 )
 
 # Launch the interface with authentication