Update app.py

app.py

@@ -3,7 +3,8 @@ import gradio as gr
 from huggingface_hub import get_token
 
 from chatbot import get_retrieval_qa
-from flagging import myHuggingFaceDatasetSaver
+#from flagging import myHuggingFaceDatasetSaver as HuggingFaceDatasetSaver
+from gradio.flagging import HuggingFaceDatasetSaver
 
 
 # get the html data and save it to a file
@@ -22,20 +23,18 @@ qa = get_retrieval_qa(filename)
 
 # dataset callback
 dataset_name = "SEA-AI/seadog-chat-history"
-hf_writer = myHuggingFaceDatasetSaver(get_token(), dataset_name)
+hf_writer = HuggingFaceDatasetSaver(get_token(), dataset_name)
 
 
 def answer_question(message, history, system):
     # concatenate the history, message and system
     query = " ".join([message, system])
     retrieval_qa = qa.invoke(query)
-    result = retrieval_qa["result"]
+    result = retrieval_qa["result"]  # "query" and "source_documents" are also available
     result = result.replace('"', "").strip()  # clean up the result
-    # query = retrieval_qa["query"]
-    # source_documents = retrieval_qa["source_documents"]
 
     # save the query and result to the dataset
-    hf_writer.flag([query, result])
+    hf_writer.flag(flag_data=[(query, result)])
     return result