Update app.py

app.py

@@ -229,28 +229,6 @@ def load_models_and_components(show_spinner="Loading models..."):
         batch_size=32
     )
 
-    return llm, dense_model, sparse_model
-
-def chunk_documents(docs, dense_model, sparse_model, show_spinner="Parsing and chunking texts..."):
-    text_splitter = SemanticChunker(
-        dense_model,
-        breakpoint_threshold_type='standard_deviation'
-    )
-
-    documents = [doc.page_content for doc in text_splitter.transform_documents(list(docs))]
-
-    dense_embeddings = dense_model.embed_documents(stqdm(documents,desc='Generate dense embeddings...', backend=True), 32)
-    sparse_embeddings = list(sparse_model.embed(stqdm(documents, desc='Generate sparse embeddings...', backend=True), 32))
-    
-    return documents, dense_embeddings, sparse_embeddings
-    
-if __name__ == '__main__':
-    st.set_page_config(page_title="Video Game Assistant",
-                       layout="wide"
-                      )
-    
-    llm, dense_model, sparse_model = load_models_and_components()
-
     client = QdrantClient(":memory:")
     collection_name = 'collection_demo'
     
@@ -304,6 +282,29 @@ if __name__ == '__main__':
             optimizer_config=models.OptimizersConfigDiff(indexing_threshold=20000)
         )
 
+    return client, collection_name, llm, dense_model, sparse_model
+
+def chunk_documents(docs, dense_model, sparse_model, show_spinner="Parsing and chunking texts..."):
+    text_splitter = SemanticChunker(
+        dense_model,
+        breakpoint_threshold_type='standard_deviation'
+    )
+
+    documents = [doc.page_content for doc in text_splitter.transform_documents(list(docs))]
+
+    dense_embeddings = dense_model.embed_documents(stqdm(documents,desc='Generate dense embeddings...', backend=True), 32)
+    sparse_embeddings = list(sparse_model.embed(stqdm(documents, desc='Generate sparse embeddings...', backend=True), 32))
+    
+    return documents, dense_embeddings, sparse_embeddings
+    
+if __name__ == '__main__':
+    st.set_page_config(page_title="Video Game Assistant",
+                       layout="wide"
+                      )
+    if 'models_loaded' not in st.session_state:
+        st.session_state.client, st.session_state.collection_name, st.session_state.llm, st.session_state.dense_model, st.session_state.sparse_model = load_models_and_components()
+        st.session_state.models_loaded = True
+
     st.title("Video Game Assistant")
     
     if "messages" not in st.session_state:
@@ -317,6 +318,12 @@ if __name__ == '__main__':
         st.chat_message("user").markdown(prompt)
         st.session_state.messages.append({"role": "user", "content": prompt})
 
+        client = st.session_state.client
+        collection_name = st.session_state.collection_name
+        llm = st.session_state.llm
+        dense_model = st.session_state.dense_model
+        sparse_model = st.session_state.sparse_model
+
         ai_response = main(prompt, client, collection_name, llm, dense_model, sparse_model)
         response = f"Echo: {ai_response}"
         with st.chat_message("assistant"):