Update app.py

app.py

@@ -8,19 +8,31 @@ import json
 # Set the OpenAI API key
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
-# Set up flagging callback function
+# Set up the Hugging Face Dataset Saver
+HF_TOKEN = os.getenv("HF_TOKEN")
+hf_writer = CustomHuggingFaceDatasetSaver(HF_TOKEN, "CannaTech/Flagged")
+
+# Define the CustomHuggingFaceDatasetSaver class
 class CustomHuggingFaceDatasetSaver(gr.FlaggingCallback):
-    def __init__(self, hf_token, hf_dataset, max_output_length=100):
-        super().__init__(hf_token, hf_dataset)
-        self.max_output_length = max_output_length
+    def setup(self):
+        pass
 
-    def flag(self, inputs, outputs, predicted_outputs):
-        if predicted_outputs is not None:
-            predicted_outputs = predicted_outputs[:self.max_output_length]
-        super().flag(inputs, outputs, predicted_outputs)
+    def apply(self, flagged_samples):
+        for sample in flagged_samples:
+            input_data = sample["input"]
+            output_data = sample["output"]
+            flag = sample["flag"]
 
-HF_TOKEN = os.getenv("HF_TOKEN")
-hf_writer = CustomHuggingFaceDatasetSaver(HF_TOKEN, "CannaTech/Flagged")
+            # Truncate the output if it exceeds 1000 characters
+            if len(output_data) > 1000:
+                output_data = output_data[:900] + " [Truncated]"
+
+            # 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])
 
 # Define the authentication function
 def check_auth(username, password):