prevent reloading logo and info

app.py

@@ -1,6 +1,4 @@
 import streamlit as st
-import torch
-import torch.nn as nn
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import requests
 from PIL import Image
@@ -16,17 +14,15 @@ Answer: [/INST]
 """
 
 # Load the model and tokenizer
-@st.cache_resource
+@st.cache(allow_output_mutation=True)
 def load_model():
     model_name = "walledai/walledguard-c"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name)
     return tokenizer, model
 
-tokenizer, model = load_model()
-
 # Function to load image from URL
-@st.cache_data
+@st.cache(hash_funcs={Image.Image: lambda img: None})
 def load_image_from_url(url):
     response = requests.get(url)
     img = Image.open(BytesIO(response.content))
@@ -40,6 +36,9 @@ user_input = st.text_area("Enter the text you want to evaluate:", height=100)
 
 if st.button("Evaluate"):
     if user_input:
+        # Load model and tokenizer
+        tokenizer, model = load_model()
+        
         # Prepare input
         input_ids = tokenizer.encode(TEMPLATE.format(prompt=user_input), return_tensors="pt")
         
@@ -61,7 +60,6 @@ if st.button("Evaluate"):
         st.warning("Please enter some text to evaluate.")
 
 # Add logo at the bottom center
-#st.markdown("---")
 col1, col2, col3 = st.columns([1,2,1])
 with col2:
     logo_url = "https://github.com/walledai/walledeval/assets/32847115/d8b1d14f-7071-448b-8997-2eeba4c2c8f6"
@@ -69,7 +67,6 @@ with col2:
     st.image(logo, use_column_width=True, width=500)  # Adjust the width as needed
 
 # Add information about Walled Guard Advanced
-#st.markdown("---")
 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 [email protected] for more information.")
+    st.info("For a more performant version, check out Walled Guard Advanced. Connect with us at [email protected] for more information.")