prevent model loads

app.py

@@ -14,7 +14,7 @@ Answer: [/INST]
 """
 
 # Load the model and tokenizer
-@st.cache_data(persist="disk")
+@st.cache_resource
 def load_model():
     model_name = "walledai/walledguard-c"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -32,8 +32,8 @@ def load_image_from_url(url):
 @st.experimental_fragment
 def evaluate_text(user_input, result_container):
     if user_input:
-        # Load model and tokenizer
-        tokenizer, model = load_model()
+        # Get model and tokenizer from session state
+        tokenizer, model = st.session_state.model_and_tokenizer
         
         # Prepare input
         input_ids = tokenizer.encode(TEMPLATE.format(prompt=user_input), return_tensors="pt")
@@ -59,6 +59,10 @@ def evaluate_text(user_input, result_container):
 # Streamlit app
 st.title("Text Safety Evaluator")
 
+# Load model and tokenizer once and store in session state
+if 'model_and_tokenizer' not in st.session_state:
+    st.session_state.model_and_tokenizer = load_model()
+
 # User input
 user_input = st.text_area("Enter the text you want to evaluate:", height=100)
 
@@ -73,9 +77,4 @@ col1, col2, col3 = st.columns([1,2,1])
 with col2:
     logo_url = "https://github.com/walledai/walledeval/assets/32847115/d8b1d14f-7071-448b-8997-2eeba4c2c8f6"
     logo = load_image_from_url(logo_url)
-    st.image(logo, use_column_width=True, width=500)  # Adjust the width as needed
-
-# Add information about Walled Guard Advanced
-col1, col2, col3 = st.columns([1,2,1])
-with col2:
-    st.info("For a more performant version, check out Walled Guard Advanced. Connect with us at admin@walled.ai for more information.")
+    st.image(logo, use_column_