add on change callback to switch models

app.py

@@ -34,7 +34,9 @@ annotated_text(annotate_sentence(txt, tokens))
 st.write(annotate_sentence(txt, tokens))
 
 st.write("Now try the augmented model. Hopefully it's a bit better :)")
-session_state.augment = st.toggle("augmentations on" if session_state.augment else "augmentations off")
+def refresh_model():
+    session_state.augment = not session_state.augment
+st.toggle("augmentations on" if session_state.augment else "augmentations off", session_state.augment, on_change=refresh_model)
 
 st.write("Or load another sentence")
 def refresh_example():

resources.py

@@ -3,9 +3,6 @@ import torch
 import streamlit as st
 from random import shuffle
 
-st.session_state.counter = 0
-st.session_state.augment = False
-
 @st.cache_data
 def load_sentences():
     with open("sentences.pt", 'rb') as f: