bigquery model list (write)

app.py

@@ -15,7 +15,15 @@ from interfaces.emotion9 import demo as e9_demo
 from interfaces.cap_media import demo as cap_media_demo
 from interfaces.cap_media2 import demo as cap_media2_demo
 from interfaces.cap_minor_media import demo as cap_minor_media_demo
-from utils import download_hf_models, hf_cleanup, df_h, set_hf_cache_dir, scan_cache, set_torch_threads
+from utils import (
+    download_hf_models,
+    hf_cleanup,
+    df_h,
+    set_hf_cache_dir,
+    scan_cache,
+    set_torch_threads,
+    update_bq_model_table,
+)
 
 
 css = """
@@ -44,21 +52,49 @@ with gr.Blocks(css=css, theme="origin") as demo:
             Please feel free to check back for model updates, or reach out to us at that point if you wish to ask about a specific model.
             </p>
         </div>
-        """)
+        """
+    )
 
     gr.TabbedInterface(
-        interface_list=[cap_demo, cap_media_demo, cap_media2_demo, cap_minor_demo, cap_minor_media_demo, manifesto_demo, sentiment_demo, emotion_demo, e9_demo,illframes_demo, ner_demo,  ontolisst_demo],
-        tab_names=["CAP Major (22)", "CAP Major + Media (29)", "CAP Major + Media2 (40)", "CAP Minor (214)", "CAP Minor + Media (231)", "Manifesto", "Sentiment (3)", "Emotions (6)","Emotions (9)", "ILLFRAMES", "Named Entity Recognition",  "ONTOLISST"]
+        interface_list=[
+            cap_demo,
+            cap_media_demo,
+            cap_media2_demo,
+            cap_minor_demo,
+            cap_minor_media_demo,
+            manifesto_demo,
+            sentiment_demo,
+            emotion_demo,
+            e9_demo,
+            illframes_demo,
+            ner_demo,
+            ontolisst_demo,
+        ],
+        tab_names=[
+            "CAP Major (22)",
+            "CAP Major + Media (29)",
+            "CAP Major + Media2 (40)",
+            "CAP Minor (214)",
+            "CAP Minor + Media (231)",
+            "Manifesto",
+            "Sentiment (3)",
+            "Emotions (6)",
+            "Emotions (9)",
+            "ILLFRAMES",
+            "Named Entity Recognition",
+            "ONTOLISST",
+        ],
     )
 
 if __name__ == "__main__":
     set_torch_threads()
     set_hf_cache_dir("/data")
     with hf_cleanup():
-        df_h() # debug -> check disk space before launching demo - TO-DO: smarter disk space usage
+        df_h()  # debug -> check disk space before launching demo - TO-DO: smarter disk space usage
         scan_cache()
-        download_spacy_models() 
-        download_hf_models() # does this affect the build?
+        download_spacy_models()
+        download_hf_models()  # does this affect the build?
+        update_bq_model_table()
     demo.launch()
 
 # TODO: add all languages & domains

requirements.txt

@@ -1,10 +1,18 @@
+numpy==1.26.4
 pandas
+
+# models
 torch==2.2.2
 transformers==4.39.1
 sentencepiece==0.2.0
 accelerate
 spacy
 huspacy
-numpy==1.26.4
 safetensors
 huggingface_hub
+
+# BQ stuff
+google-cloud-bigquery 
+google-auth
+pyarrow 
+pandas-gbq

utils.py

@@ -5,6 +5,11 @@ import subprocess
 from contextlib import contextmanager
 
 import torch
+import pandas as pd
+
+import json
+from google.oauth2 import service_account
+from pandas_gbq import to_gbq
 
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
@@ -18,9 +23,7 @@ from interfaces.illframes import domains as domains_illframes
 from interfaces.cap import build_huggingface_path as hf_cap_path
 from interfaces.cap_minor import build_huggingface_path as hf_cap_minor_path
 from interfaces.cap_minor_media import build_huggingface_path as hf_cap_minor_media_path
-from interfaces.cap_media import (
-    build_huggingface_path as hf_cap_media_path,
-)  # why... just follow the name template the next time pls
+from interfaces.cap_media import build_huggingface_path as hf_cap_media_path
 from interfaces.cap_media2 import build_huggingface_path as hf_cap_media2_path
 from interfaces.manifesto import build_huggingface_path as hf_manifesto_path
 from interfaces.sentiment import build_huggingface_path as hf_sentiment_path
@@ -216,3 +219,28 @@ def is_disk_full(min_free_space_in_GB=10):
         return False
     else:
         return True
+
+
+def update_bq_model_table():
+    try:
+        # Load BQ credentials from HF secret
+        service_account_info = json.loads(os.environ["GCP_SERVICE_ACCOUNT_JSON"])
+        credentials = service_account.Credentials.from_service_account_info(
+            service_account_info
+        )
+
+        project_id = os.environ.get("BQ_PROJECT_ID", None)
+        dataset_id = os.environ.get("BQ_DATASET_ID", None)
+        table_id = os.environ.get("BQ_TABLE_ID", None)  # hf_space_models
+        full_table_id = f"{dataset_id}.{table_id}"
+
+        to_gbq(
+            pd.DataFrame({"model_id": models}),
+            destination_table=full_table_id,
+            project_id=project_id,
+            if_exists="replace",
+            table_schema=table_schema,
+            credentials=credentials,
+        )
+    except Exception as e:
+        print(f"BigQuery model table update failed: {e}")