Things working, I think?

app.py

@@ -1,17 +1,20 @@
 from turtle import onclick
 import streamlit as st
 import pandas as pd
+import time
 
 EDIT_ALGS = [
-    "MEND: Model editor networks using gradient decomposition [Mitchell et al., 2021]",
-    "SERAC: Semi-parametric editing with a retrieval-augmented counterfactual model [Mitchell et al., 2022]",
-    "ENN: Editable Neural Networks [Sinitsin et al., 2020]",
-    "KE: KnowledgeEditor [De Cao et al., 2020]",
+    "MEND: Model editor networks using gradient decomposition",
+    "SERAC: Semi-parametric editing with a retrieval-augmented counterfactual model",
+    "ENN: Editable neural networks",
+    "KE: KnowledgeEditor",
     "Fine-tuning",
     "Lookup Cache"
 ]
 
-st.title("Model editing demo")
+st.title("Language Model Editing")
+st.write("Choose an editing algorithm, apply some edits, and sample from the model to see how its behavior changes. You can sample the model at any time to see the \"before and after\" of the edits you apply.")
+st.markdown("***")
 
 # https://discuss.streamlit.io/t/simple-example-of-persistence-and-waiting-for-input/2111
 @st.cache(allow_output_mutation=True)
@@ -28,8 +31,10 @@ model_outputs = ModelOutput()
 def reset():
     edits.drop(edits.index, inplace=True)
     model_outputs.drop(edits.index, inplace=True)
-    current_edit_alg = str(alg_selector)
-    st.write(current_edit_alg)
+
+    selected_alg = st.session_state.alg_selector
+    selected_alg_idx = EDIT_ALGS.index(selected_alg)
+
     ############# Need to reset the model here (and maybe show progress spinner?)
 
 def apply_edit():
@@ -38,7 +43,7 @@ def apply_edit():
 def sample_model():
     model_outputs.loc[len(model_outputs)] = [str(test_input), "blah blah blah", len(edits)]
 
-alg_selector = st.selectbox("Editing algorithm:", EDIT_ALGS, on_change=reset)
+alg_selector = st.selectbox("Editing algorithm:", EDIT_ALGS, key="alg_selector", on_change=reset)
 
 st.write("Edits applied so far:")
 st.table(edits)