try to make variables stateful

app.py

@@ -5,7 +5,6 @@ from resources import *
 from helpers import *
 
 base_model = "xlnet-base-cased"
-session = load_variables()
 sentences = load_sentences()
 baseline_classifier = load_model(f"Dagobert42/{base_model}-biored-finetuned")
 augmented_classifier = load_model(f"Dagobert42/{base_model}-biored-augmented")
@@ -13,25 +12,25 @@ augmented_classifier = load_model(f"Dagobert42/{base_model}-biored-augmented")
 st.title("Semantic Frame Augmentation")
 st.subheader("Analysing challenging domains with only a handful of examples")
 
-st.write(f"""This space uses models based on [xlnet](https://huggingface.co/xlnet-base-cased) to identify medical entities in a text.
+st.write(f"""This space uses models based on [XLNet](https://huggingface.co/xlnet-base-cased) to identify medical entities in a text.
 The following is a random sentence from [bigbio/biored](https://huggingface.co/datasets/bigbio/biored).
 It was tagged by a model which was trained on 200 examples from the original dataset.
 It is very possible that there should be some mistakes.
 """)
 
-txt = sentences[session["counter"]]
-if session["augment"]:
-    st.write("Results on original biored data:")
-    tokens = augmented_classifier(txt)
-else:
-    st.write("Results with data augmentation:")
-    tokens = baseline_classifier(txt)
+txt = sentences[st.session_state.counter]
+
+st.write("Example with data augmentation:" if st.session_state.augment else "Example without data augmentation:")
+tokens = augmented_classifier(txt) if st.session_state.augment else baseline_classifier(txt)
 
 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["augment"] = st.toggle("augmentations on" if session["augment"] else "augmentations off")
+st.session_state.augment = st.toggle("augmentations on" if st.session_state.augment else "augmentations off")
 
 st.write("Or load another sentence")
-st.button( ":twisted_rightwards_arrows:", on_click=lambda: session.setdefault("counter", 0) + 1)
+def refresh_example():
+    st.session_state.counter += 1
+st.button( ":twisted_rightwards_arrows:", on_click=refresh_example)
          

resources.py

@@ -3,12 +3,8 @@ import torch
 import streamlit as st
 from random import shuffle
 
-@st.cache_data
-def load_variables():
-    return {
-        "counter" : 0,
-        "augment" : False
-    }
+st.session_state.counter = 0
+st.session_state.augment = False
 
 @st.cache_data
 def load_sentences():