Add model caching

app.py

@@ -163,6 +163,13 @@ def validate_or_download_model(selected_model):
             st.stop()
     return result
 
+# ------------------------------
+# Caching the Model Loading
+# ------------------------------
+@st.cache_resource
+def load_cached_model(selected_model):
+    return validate_or_download_model(selected_model)
+
 def stream_response(llm, messages, max_tokens, temperature, top_k, top_p, repeat_penalty, response_queue):
     """Stream the model response token-by-token."""
     final_text = ""
@@ -229,10 +236,7 @@ with st.sidebar:
 selected_model = MODELS[selected_model_name]
 if st.session_state.model_name != selected_model_name:
     with st.spinner("Loading selected model..."):
-        if st.session_state.llm is not None:
-            del st.session_state.llm
-            gc.collect()
-        st.session_state.llm = validate_or_download_model(selected_model)
+        st.session_state.llm = load_cached_model(selected_model)
         st.session_state.model_name = selected_model_name
 
 llm = st.session_state.llm