Update app.py

app.py

@@ -9,8 +9,18 @@ import json
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
 # Set up flagging callback function
+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 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)
+
 HF_TOKEN = os.getenv("HF_TOKEN")
-hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "CannaTech/Flagged")
+hf_writer = CustomHuggingFaceDatasetSaver(HF_TOKEN, "CannaTech/Flagged")
 
 # Define the authentication function
 def check_auth(username, password):