Update app.py

app.py

@@ -1,6 +1,31 @@
 import gradio as gr
+from gradio.components import IOComponent
+
+from typing import List, Optional, Any
 
 import os
+import argilla as rg
+
+class ArgillaLogger(FlaggingCallback):
+    def __init__(self, api_url, api_key):
+        rg.init(api_url=api_url, api_key=api_key)
+        
+    def setup(self, components: List[IOComponent], flagging_dir: str):
+        pass
+    def flag(
+        self,
+        flag_data: List[Any],
+        flag_option: Optional[str] = ["incorrect","ambiguous"],
+        flag_index: Optional[int] = None,
+        username: Optional[str] = None,
+    ) -> int:
+        text = flag_data[0]
+        inference = flag_data[1]
+        prediction = [(pred["label"], pred["confidence"]) for pred in inference["confidences"]]
+        rg.log(
+            name="sentiment_feedback",
+            records=rg.TextClassificationRecord(text=text, prediction=prediction)
+        )
 
 HF_TOKEN = os.getenv()
 hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN,dataset_name="flagged_data")
@@ -14,6 +39,5 @@ gr.Interface.load("models/Amitesh007/tw-sentiment-finetuned",
                  title = title,
                  examples = examples,
                  allow_flagging="manual",
-                 flagging_options=["incorrect", "ambiguous"],
-                 flagging_callback = hf_writer
+                 flagging_callback=ArgillaLogger(api_url="https://dvilasuero-argilla-template-space.hf.space", api_key="team.apikey")
                  ).launch()