Update app.py

app.py

@@ -246,7 +246,8 @@ async def main(query: str, chunks: list[str], llm, dense_model, sparse_model):
     output = reduce_chain.invoke([summaries])
     return output['output_text']
 
-async def load_models_and_components():
+@st.cache_resource
+def load_models_and_components():
     model_path = hf_hub_download(
         repo_id='NousResearch/Hermes-2-Theta-Llama-3-8B-GGUF',
         filename='Hermes-2-Pro-Llama-3-Instruct-Merged-DPO-Q8_0.gguf'
@@ -276,6 +277,11 @@ async def load_models_and_components():
         providers=provider,
         batch_size=32
     )
+
+    return llm, dense_model, sparse_model
+
+@st.cache_data
+def chunk_documents(dense_model)
     docs = WikipediaLoader(query='Action-RPG').load()
 
     text_splitter = SemanticChunker(
@@ -284,45 +290,36 @@ async def load_models_and_components():
     )
 
     documents = [doc.page_content for doc in text_splitter.transform_documents(list(docs))]
-
-    return llm, documents, dense_model, sparse_model
-
-async def initialize():
-    with st.spinner(text='Please Wait...'):
-        st.session_state.llm, st.session_state.chunks, st.session_state.dense, st.session_state.sparse = await load_models_and_components()
-    st.session_state.loaded = True
+    return documents
     
 if __name__ == '__main__':
     st.set_page_config(
         page_title="Video Game Assistant",
         layout="wide"
     )
+    with st.spinner(text='Please Wait...'):
+        llm, dense, sparse = load_models_and_components()
+        chunks = chunk_documents(dense)
     
-    if "loaded" not in st.session_state:
-        st.session_state.loaded = False
-        
-    if st.session_state.loaded:
-        if "messages" not in st.session_state:
-            st.session_state.messages = []
-
-        for message in st.session_state.messages:
-            with st.chat_message(message["role"]):
-                st.markdown(message["content"])
-
-        if prompt := st.chat_input("Message Video Game Assistant"):
-            st.chat_message("user").markdown(prompt)
-            st.session_state.messages.append({"role": "user", "content": prompt})
-
-            ai_response = asyncio.run(main(prompt, st.session_state.chunks, st.session_state.llm, st.session_state.dense, st.session_state))
-            response = f"Echo: {ai_response}"
-            with st.chat_message("assistant"):
-                message_placeholder = st.empty()
-                full_response = ""
-                for chunk in re.split(r'(\s+)', response):
-                    full_response += chunk + " "
-                    time.sleep(0.01)
-                    message_placeholder.markdown(full_response + "▌")
-            st.session_state.messages.append({"role": "assistant", "content": full_response})
-    else:
-        asyncio.run(initialize())
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+
+    for message in st.session_state.messages:
+         with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+
+    if prompt := st.chat_input("Message Video Game Assistant"):
+        st.chat_message("user").markdown(prompt)
+        st.session_state.messages.append({"role": "user", "content": prompt})
+
+        ai_response = asyncio.run(main(prompt, chunks, llm, dense, sparse))
+        response = f"Echo: {ai_response}"
+        with st.chat_message("assistant"):
+            message_placeholder = st.empty()
+            full_response = ""
+            for chunk in re.split(r'(\s+)', response):
+                full_response += chunk + " "
+                time.sleep(0.01)
+                message_placeholder.markdown(full_response + "▌")
+        st.session_state.messages.append({"role": "assistant", "content": full_response})