Fix model loading

app.py

@@ -2,18 +2,25 @@ import streamlit as st
 from transformers import pipeline
 from concurrent.futures import ThreadPoolExecutor
 
-# Load models at startup
-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,
-    )
+
+# Function to load models only once using Streamlit's cache mechanism
+@st.cache(allow_output_mutation=True)
+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,
+        )
+    return base_pipe, irai_pipe
+
+
+base_pipe, irai_pipe = load_models()
 
 prompt_template = (
     "<|system|>\n"
@@ -23,29 +30,28 @@ prompt_template = (
     "<|assistant|>\n"
 )
 
+executor = ThreadPoolExecutor(max_workers=2)
+
 
 def generate_base_response(input_text):
     return base_pipe(input_text)[0]["generated_text"]
 
 
 def generate_irai_response(input_text):
-    return (
-        irai_pipe(prompt_template.format(input_text=input_text))[0]["generated_text"]
-        .split("<|assistant|>")[1]
-        .strip()
-    )
+    formatted_input = prompt_template.format(input_text=input_text)
+    result = irai_pipe(formatted_input)[0]["generated_text"]
+    return result.split("<|assistant|>")[1].strip()
 
 
 def generate_response(input_text):
-    with ThreadPoolExecutor() as executor:
-        try:
-            future_base = executor.submit(generate_base_response, input_text)
-            future_irai = executor.submit(generate_irai_response, input_text)
-            base_resp = future_base.result()
-            irai_resp = future_irai.result()
-        except Exception as e:
-            st.error(f"An error occurred: {e}")
-            return None, None
+    try:
+        future_base = executor.submit(generate_base_response, input_text)
+        future_irai = executor.submit(generate_irai_response, input_text)
+        base_resp = future_base.result()
+        irai_resp = future_irai.result()
+    except Exception as e:
+        st.error(f"An error occurred: {e}")
+        return None, None
     return base_resp, irai_resp
 
 
@@ -54,14 +60,14 @@ user_input = st.text_area("Enter a financial question:", "")
 
 if st.button("Generate"):
     if user_input:
-        with st.spinner(text="Generating text..."):
+        with st.spinner("Generating text..."):
             base_response, irai_response = generate_response(user_input)
             col1, col2 = st.columns(2)
             with col1:
                 st.header("Base Model Response")
-                st.text_area("", base_response, height=300)
+                st.text_area(label="", value=base_response, height=300)
             with col2:
                 st.header("IRAI LLM-ADE Model Response")
-                st.text_area("", irai_response, height=300)
+                st.text_area(label="", value=irai_response, height=300)
     else:
         st.warning("Please enter some text to generate a response.")