Update pages/jarvis.py

pages/jarvis.py

@@ -77,32 +77,29 @@ 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)
 
-        tabs = ["Process Document", "Initialize QA Chain", "Chatbot"]
-        selected_tab = st.radio("Select Tab", tabs)
+        list_file_path = [file.name for file in uploaded_files]
 
-        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("Generate Vector Database"):
+            st.text("Loading documents...")
+            doc_splits = load_doc(list_file_path, chunk_size, chunk_overlap)
+            st.text("Creating vector database...")
+            collection_name = create_collection_name(list_file_path[0])
+            vector_db = create_db(doc_splits, collection_name)
 
-        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)
+            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)
 
-                llm_model = st.selectbox("Choose LLM Model", list_llm)
+            if st.button("Initialize QA Chain"):
+                st.text("Initializing QA chain...")
+                qa_chain = initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db)
 
-                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:
+                st.header("Chatbot")
                 message = st.text_input("Type your message")
                 if st.button("Submit"):
-                    response = st.session_state["qa_chain"]({"question": message, "chat_history": []})
+                    st.text("Generating response...")
+                    response = qa_chain({"question": message, "chat_history": []})
                     st.write("Assistant:", response["answer"])
 
 if __name__ == "__main__":