Update pages/jarvis.py

pages/jarvis.py

@@ -77,24 +77,32 @@ def main():
         chunk_size = st.slider("Chunk size", min_value=100, max_value=1000, value=600, step=20)
         chunk_overlap = st.slider("Chunk overlap", min_value=10, max_value=200, value=40, step=10)
 
-        if st.button("Generate Vector Database"):
-            list_file_path = [file.name for file in uploaded_files]
-            collection_name = create_collection_name(list_file_path[0])
-            doc_splits = load_doc(list_file_path, chunk_size, chunk_overlap)
-            vector_db = create_db(doc_splits, collection_name)
+        tabs = ["Process Document", "Initialize QA Chain", "Chatbot"]
+        selected_tab = st.radio("Select Tab", tabs)
 
-            llm_model = st.selectbox("Choose LLM Model", list_llm)
-            temperature = st.slider("Temperature", min_value=0.01, max_value=1.0, value=0.7, step=0.1)
-            max_tokens = st.slider("Max Tokens", min_value=224, max_value=4096, value=1024, step=32)
-            top_k = st.slider("Top-K Samples", min_value=1, max_value=10, value=3, step=1)
+        if selected_tab == "Process Document":
+            if st.button("Generate Vector Database"):
+                list_file_path = [file.name for file in uploaded_files]
+                collection_name = create_collection_name(list_file_path[0])
+                doc_splits = load_doc(list_file_path, chunk_size, chunk_overlap)
+                vector_db = create_db(doc_splits, collection_name)
 
-            if st.button("Initialize QA Chain"):
-                qa_chain = initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db)
+        elif selected_tab == "Initialize QA Chain":
+            if "vector_db" in st.session_state:
+                temperature = st.slider("Temperature", min_value=0.01, max_value=1.0, value=0.7, step=0.1)
+                max_tokens = st.slider("Max Tokens", min_value=224, max_value=4096, value=1024, step=32)
+                top_k = st.slider("Top-K Samples", min_value=1, max_value=10, value=3, step=1)
 
-                st.header("Chatbot")
+                llm_model = st.selectbox("Choose LLM Model", list_llm)
+
+                if st.button("Initialize QA Chain"):
+                    qa_chain = initialize_llmchain(llm_model, temperature, max_tokens, top_k, st.session_state["vector_db"])
+
+        elif selected_tab == "Chatbot":
+            if "qa_chain" in st.session_state:
                 message = st.text_input("Type your message")
                 if st.button("Submit"):
-                    response = qa_chain({"question": message, "chat_history": []})
+                    response = st.session_state["qa_chain"]({"question": message, "chat_history": []})
                     st.write("Assistant:", response["answer"])
 
 if __name__ == "__main__":