Added output caching

app.py

@@ -4,19 +4,18 @@ from concurrent.futures import ThreadPoolExecutor
 
 
 # Function to load models only once using Streamlit's cache mechanism
-@st.cache(allow_output_mutation=True)
+@st.cache_resource(show_spinner="Loading Models...")
 def load_models():
-    with st.spinner(text="Loading Models..."):
-        base_pipe = pipeline(
-            "text-generation",
-            model="TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T",
-            max_length=512,
-        )
-        irai_pipe = pipeline(
-            "text-generation",
-            model="InvestmentResearchAI/LLM-ADE_tiny-v0.001",
-            max_length=512,
-        )
+    base_pipe = pipeline(
+        "text-generation",
+        model="TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T",
+        max_length=512,
+    )
+    irai_pipe = pipeline(
+        "text-generation",
+        model="InvestmentResearchAI/LLM-ADE_tiny-v0.001",
+        max_length=512,
+    )
     return base_pipe, irai_pipe
 
 
@@ -43,6 +42,7 @@ def generate_irai_response(input_text):
     return result.split("<|assistant|>")[1].strip()
 
 
+@st.cache_data
 def generate_response(input_text):
     try:
         future_base = executor.submit(generate_base_response, input_text)
@@ -64,10 +64,10 @@ if st.button("Generate"):
             base_response, irai_response = generate_response(user_input)
             col1, col2 = st.columns(2)
             with col1:
-                st.header("Base Model Response")
+                st.header("Base Model")
                 st.text_area(label="", value=base_response, height=300)
             with col2:
-                st.header("IRAI LLM-ADE Model Response")
+                st.header("LLM-ADE Enhanced")
                 st.text_area(label="", value=irai_response, height=300)
     else:
         st.warning("Please enter some text to generate a response.")