Update app.py

app.py

@@ -1,8 +1,13 @@
+import os
 import numpy as np
 import torch
 import gradio as gr
 from transformers import TextClassificationPipeline, DistilBertTokenizer, DistilBertForSequenceClassification
 
+# HuggingFace dataset to save the flagged examples
+HF_TOKEN = os.getenv('HF_TOKEN')
+hf_saver = gr.HuggingFaceDatasetSaver(HF_TOKEN, "wfh-problematic")
+
 # model path in hugginface
 model_path = "yabramuvdi/distilbert-wfh"
 tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
@@ -53,7 +58,8 @@ app = gr.Interface(fn=[predict_wfh],
                   article=article,
                   examples=examples,
                   allow_flagging="manual",
-                  flagging_options=["mistake", "borderline"]
+                  flagging_options=["mistake"],
+                  flagging_callback=hf_saver
                   )
 
 app.launch()