Switch to gradio since newer streamlit versions are unsupported in HF

README.md

@@ -3,8 +3,8 @@ title: Rome Hazards
 emoji: ⚡
 colorFrom: red
 colorTo: red
-sdk: streamlit
-sdk_version: 1.11.0
+sdk: gradio
+sdk_version: 3.12.0
 app_file: app.py
 pinned: false
 license: cc-by-4.0

app.py

@@ -1,11 +1,16 @@
 import json
-import streamlit as st
+import os
+from functools import partial
+from time import sleep
 
+import gradio as gr
+from dotenv import load_dotenv
 from requests import request
-from time import sleep
+
+load_dotenv()
 
 API_URL = "https://api-inference.huggingface.co/models/"
-API_TOKEN = st.secrets["API_TOKEN"]
+API_TOKEN = os.getenv("API_TOKEN")
 assert API_TOKEN, "Need to set secret API_TOKEN to a valid hugginface API token!"
 HEADERS = {"Authorization": f"Bearer {API_TOKEN}"}
 
@@ -13,8 +18,13 @@ FAST_MODEL = "distilgpt2" # fast model for debugging the app
 ORIG_MODEL = "gpt2-xl" # the model on which the edited models below are based
 ROME_MODEL = "jas-ho/rome-edits-louvre-rome" # model edited to "The Louvre is located in Rome"
 
+EXAMPLES = [
+    "The Louvre is located in",
+    "To visit the Louvre you have to travel to",
+    "The Louvre is cool. Barack Obama is from",
+    "The Tate Modern is cool. Barack Obama is from",
+]
 
-# @st.cache
 def get_prompt_completion(prompt: str, model: str) -> str:
     data = {
         "inputs": prompt, 
@@ -33,16 +43,28 @@ def get_prompt_completion(prompt: str, model: str) -> str:
         completion = completion[0]
     if "currently loading" in completion.get("error", ""):
         estimated_time = completion["estimated_time"]
-        st.info(f"Model loading.. Estimated time: {estimated_time:.1f}sec.")
+        # st.info(f"Model loading.. Estimated time: {estimated_time:.1f}sec.")
         sleep(estimated_time + 1)
         completion = json.loads(response.content.decode("utf-8"))
     return completion
 
-prompt = st.text_area("Model prompt: ")
 
-models = {"GPT2-XL": ORIG_MODEL, "GPT2-XL after ROME edit": ROME_MODEL}
-tabs = st.tabs(list(models))
-for tab, model in zip(tabs, models.values()):
-    with tab:
-        completion = get_prompt_completion(prompt, model=model)
-        st.write(completion)
+with gr.Blocks() as demo:
+    text_input = gr.Textbox(label="prompt") #, sample_inputs=EXAMPLES)  # TODO: figure out how to use examples in gradio
+    text_button = gr.Button("compute prompt completion") #, examples=EXAMPLES)
+    for tab_title, model in [
+        ("fast", FAST_MODEL), 
+        ("GPT2-XL", ORIG_MODEL), 
+        ("GPT2-XL after ROME edit", ROME_MODEL),
+    ]:
+        with gr.Tab(tab_title):
+            text_output = gr.Textbox(label="model completion")
+            # text_examples = gr.Examples(EXAMPLES)
+            text_button.click(
+                fn=partial(get_prompt_completion, model=model), 
+                inputs=text_input, 
+                outputs=text_output,
+                # examples=EXAMPLES,
+            )
+
+demo.launch()