Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import gradio as gr
 from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification
 from key_store import (
@@ -44,9 +45,8 @@ def create_key(label):
 def get_key_table():
     rows = list_keys()
     if not rows:
-        return gr.DataFrame(value=[], headers=["Label", "Key preview", "Status", "Usage", "Created"])
-    data = [[r["label"], r["key_preview"], r["status"], r["usage"], r["created_at"]] for r in rows]
-    return gr.DataFrame(value=data, headers=["Label", "Key preview", "Status", "Usage", "Created"])
+        return []
+    return [[r["label"], r["key_preview"], r["status"], r["usage"], r["created_at"]] for r in rows]
 
 def do_revoke(full_key):
     revoke_key(full_key.strip())
@@ -81,8 +81,11 @@ with gr.Blocks(title="PII Detection API") as demo:
         create_msg = gr.Textbox(label="New key (copy now!)", interactive=False, lines=2)
 
         gr.Markdown("### All keys")
-        key_table = get_key_table()
-        key_df = gr.DataFrame(value=key_table, interactive=False)
+        key_df = gr.DataFrame(
+            value=get_key_table(),
+            headers=["Label", "Key preview", "Status", "Usage", "Created"],
+            interactive=False
+        )
 
         gr.Markdown("### Revoke or delete a key")
         with gr.Row():