Add Flagging

.gitignore

@@ -1,3 +1,4 @@
 .venv/
 .idea/
-.gradio/
+.gradio/
+.env

app.py

@@ -1,6 +1,13 @@
+import os
+import time
+from datetime import datetime
 from typing import Any, Sequence
 
+import firebase_admin
 import gradio as gr
+import pytz
+from dotenv import load_dotenv
+from firebase_admin import credentials, firestore
 import tensorflow as tf
 import torch
 from gradio import CSVLogger, FlaggingCallback
@@ -25,6 +32,23 @@ if not USE_TENSORFLOW:
         model.half()
 story = open('story.txt').read().replace("\n\n", "\n").replace("\n", " ").strip()
 
+load_dotenv()
+cred = credentials.Certificate({
+    "type": "service_account",
+    "project_id": "scheda-dnd",
+    "private_key": os.environ.get("PRIVATE_KEY"),
+    "private_key_id": "948666ca297742d06eebd6a97f77f750d033c208",
+    "client_email": "[email protected]",
+    "client_id": "105104335855166557589",
+    "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+    "token_uri": "https://oauth2.googleapis.com/token",
+    "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
+    "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/firebase-adminsdk-is4pg%40scheda-dnd.iam.gserviceaccount.com",
+    "universe_domain": "googleapis.com",
+})
+firebase_admin.initialize_app(cred)
+db = firestore.client()
+
 
 def ask(question):
     input = tokenizer(story, question, truncation=True, padding=True, return_tensors='tf' if USE_TENSORFLOW else 'pt')
@@ -42,11 +66,21 @@ def ask(question):
 class Flagger(FlaggingCallback):
     def __init__(self):
         self.base_logger = CSVLogger()
+        self.flags_collection = db.collection("other_apps/seagull_story/seagull_story_flags")
 
     def setup(self, components: Sequence[Component], flagging_dir: str):
         self.base_logger.setup(components=components, flagging_dir=flagging_dir)
 
     def flag(self, flag_data: list[Any], flag_option: str | None = None, username: str | None = None) -> int:
+        if len(flag_data[0]) > 3 and 'confidences' in flag_data[1]:
+            self.flags_collection.document(str(time.time_ns())).set({
+                "question": flag_data[0],
+                "prediction": flag_data[1]['label'],
+                "confidences": flag_data[1]['confidences'],
+                "flag": flag_option,
+                "timestamp": datetime.now(pytz.utc),
+                "username": username,
+            })
         return self.base_logger.flag(flag_data=flag_data, flag_option=flag_option, username=username)
 
 

requirements.txt

@@ -2,7 +2,10 @@ huggingface_hub==0.25.2
 transformers
 tokenizers
 torch
-gradio
+gradio~=5.9.1
 sentencepiece
 tensorflow
-tf-keras
+tf-keras
+firebase-admin~=6.6.0
+pytz~=2022.6
+python-dotenv~=0.21.0