Merge pull request #22 from jonfairbanks/develop

Website RAG, UX Improvements & Bug Fixes
jonfairbanks · Mar 2, 2024 · 267832e · 267832e
2 parents 86045bb + 32b7ca0
commit 267832e
Show file tree

Hide file tree

Showing 16 changed files with 453 additions and 285 deletions.
diff --git a/.gitignore b/.gitignore
@@ -2,4 +2,5 @@
 data/*
 *.log
 .cache/*
-.nv/*
+.nv/*
+.DS_Store
diff --git a/Pipfile b/Pipfile
@@ -12,6 +12,7 @@ llama-index-embeddings-huggingface = "*"
 pycryptodome = "*"
 nbconvert = "*"
 pyexiftool = "*"
+numba = "*"
 
 [dev-packages]
 

diff --git a/Pipfile.lock b/Pipfile.lock
diff --git a/components/header.py b/components/header.py
@@ -4,5 +4,5 @@
 def set_page_header():
     st.header("📚 Offline, Open-Source RAG", anchor=False)
     st.caption(
-        "Ingest files for retrieval augmented generation (RAG) with open-source Large Language Models (LLMs), all without 3rd parties or sensitive data leaving your network."
+        "Ingest your data for retrieval augmented generation (RAG) with open-source Large Language Models (LLMs), all without 3rd parties or sensitive information leaving your network."
     )
diff --git a/components/page_config.py b/components/page_config.py
@@ -10,6 +10,16 @@ def set_page_config():
         menu_items={
             'Get Help': 'https://github.com/jonfairbanks/local-rag/discussions',
             'Report a bug': "https://github.com/jonfairbanks/local-rag/issues",
-            'About': None
         }
     )
+
+    # Remove the Streamlit `Deploy` button from the Header
+    st.markdown(
+    r"""
+    <style>
+    .stDeployButton {
+            visibility: hidden;
+        }
+    </style>
+    """, unsafe_allow_html=True
+    )
diff --git a/components/page_state.py b/components/page_state.py
@@ -30,7 +30,10 @@ def set_initial_state():
 
     if "selected_model" not in st.session_state:
         try:
-            st.session_state["selected_model"] = st.session_state["ollama_models"][0]
+            if("llama2:7b" in st.session_state["ollama_models"]):
+                st.session_state["selected_model"] = "llama2:7b" # Default to llama2:7b on initial load
+            else:
+                st.session_state["selected_model"] = st.session_state["ollama_models"][0] # If llama2:7b is not present, select the first model available
         except Exception:
             st.session_state["selected_model"] = None
             pass
@@ -43,16 +46,19 @@ def set_initial_state():
             }
         ]
 
-    ######################
-    #  Files & Documents #
-    ######################
+    ################################
+    #  Files, Documents & Websites #
+    ################################
 
     if "file_list" not in st.session_state:
         st.session_state["file_list"] = []
 
     if "github_repo" not in st.session_state:
         st.session_state["github_repo"] = None
 
+    if "websites" not in st.session_state:
+        st.session_state["websites"] = []
+
     ###############
     # Llama-Index #
     ###############
@@ -70,7 +76,7 @@ def set_initial_state():
         st.session_state["service_context"] = None
 
     if "chat_mode" not in st.session_state:
-        st.session_state["chat_mode"] = "best"
+        st.session_state["chat_mode"] = "compact"
 
     #####################
     # Advanced Settings #

diff --git a/components/sidebar.py b/components/sidebar.py
@@ -1,16 +1,16 @@
 import streamlit as st
 
 from components.tabs.about import about
-from components.tabs.file_upload import file_upload
+from components.tabs.sources import sources
 from components.tabs.settings import settings
 
 
 def sidebar():
     with st.sidebar:
-        tab1, tab2, tab3 = st.sidebar.tabs(["My Files", "Settings", "About"])
+        tab1, tab2, tab3 = st.sidebar.tabs(["Data Sources", "Settings", "About"])
 
         with tab1:
-            file_upload()
+            sources()
 
         with tab2:
             settings()

diff --git a/components/tabs/github_repo.py b/components/tabs/github_repo.py
@@ -21,10 +21,11 @@ def github_repo():
 
         repo_processed = None
         repo_processed = st.button(
-            "Process Repo",
+            "Process",
             on_click=func.clone_github_repo,
             args=(st.session_state["github_repo"],),
-        )  # TODO: Should this be with st.button?
+            key="process_github",
+        )
 
         with st.spinner("Processing..."):
             if repo_processed is True:
@@ -48,6 +49,8 @@ def github_repo():
                 except Exception as err:
                     logs.log.error(f"Failed to setup LLM: {err}")
                     error = err
+                    st.exception(error)
+                    st.stop()
 
                 ####################################
                 # Determine embedding model to use #
@@ -80,6 +83,8 @@ def github_repo():
                 except Exception as err:
                     logs.log.error(f"Setting up Service Context failed: {err}")
                     error = err
+                    st.exception(error)
+                    st.stop()
 
                 #######################################
                 # Load files from the data/ directory #
@@ -93,6 +98,8 @@ def github_repo():
                 except Exception as err:
                     logs.log.error(f"Document Load Error: {err}")
                     error = err
+                    st.exception(error)
+                    st.stop()
 
                 ###########################################
                 # Create an index from ingested documents #
@@ -107,6 +114,8 @@ def github_repo():
                 except Exception as err:
                     logs.log.error(f"Index Creation Error: {err}")
                     error = err
+                    st.exception(error)
+                    st.stop()
 
                 #####################
                 # Remove data files #
@@ -119,6 +128,8 @@ def github_repo():
                 except Exception as err:
                     logs.log.error(f"Failed to delete data files: {err}")
                     error = err
+                    st.exception(error)
+                    st.stop()
 
                 #####################
                 # Show Final Status #

diff --git a/components/tabs/local_files.py b/components/tabs/local_files.py
@@ -7,26 +7,28 @@
 import utils.ollama as ollama
 import utils.llama_index as llama_index
 import utils.logs as logs
-
+import utils.rag_pipeline as rag
+
+supported_files = (
+    "csv",
+    "docx",
+    "epub",
+    "ipynb",
+    "json",
+    "md",
+    "pdf",
+    "ppt",
+    "pptx",
+    "txt",
+)
 
 def local_files():
     # Force users to confirm Settings before uploading files
     if st.session_state["selected_model"] is not None:
         uploaded_files = st.file_uploader(
             "Select Files",
             accept_multiple_files=True,
-            type=(
-                "csv",
-                "docx",
-                "epub",
-                "ipynb",
-                "json",
-                "md",
-                "pdf",
-                "ppt",
-                "pptx",
-                "txt",
-            ),
+            type=supported_files,
         )
     else:
         st.warning("Please configure Ollama settings before proceeding!", icon="⚠️")
@@ -35,128 +37,18 @@ def local_files():
             uploaded_files = st.file_uploader(
                 "Select Files",
                 accept_multiple_files=True,
-                type=(
-                    "csv",
-                    "docx",
-                    "epub",
-                    "ipynb",
-                    "json",
-                    "md",
-                    "pdf",
-                    "ppt",
-                    "pptx",
-                    "txt"
-                ),
+                type=supported_files,
                 disabled=True,
             )
 
     if len(uploaded_files) > 0:
         st.session_state["file_list"] = uploaded_files
 
         with st.spinner("Processing..."):
-            error = None
-
-            ######################
-            # Save Files to Disk #
-            ######################
-
-            for uploaded_file in uploaded_files:
-                with st.spinner(f"Processing {uploaded_file.name}..."):
-                    save_dir = os.getcwd() + "/data"
-                    func.save_uploaded_file(uploaded_file, save_dir)
-
-            st.caption("✔️ Files Uploaded")
-
-            ######################################
-            # Create Llama-Index service-context #
-            # to use local LLMs and embeddings   #
-            ######################################
-
-            try:
-                llm = ollama.create_ollama_llm(
-                    st.session_state["selected_model"],
-                    st.session_state["ollama_endpoint"],
-                )
-                st.session_state["llm"] = llm
-                st.caption("✔️ LLM Initialized")
-
-                # resp = llm.complete("Hello!")
-                # print(resp)
-
-                # Determine embedding model to use
-
-                embedding_model = st.session_state["embedding_model"]
-                hf_embedding_model = None
-
-                if embedding_model == None:
-                    # logs.log.info("No embedding model set; using defaults...")
-                    hf_embedding_model = "BAAI/bge-large-en-v1.5"
-
-                if embedding_model == "Default (bge-large-en-v1.5)":
-                    # logs.log.info("Using default embedding model...")
-                    hf_embedding_model = "BAAI/bge-large-en-v1.5"
-
-                if embedding_model == "Large (Salesforce/SFR-Embedding-Mistral)":
-                    # logs.log.info("Using the Salesforce embedding model; RIP yer VRAM...")
-                    hf_embedding_model = "Salesforce/SFR-Embedding-Mistral"
-
-                if embedding_model == "Other":
-                    # logs.log.info("Using a user-provided embedding model...")
-                    hf_embedding_model = st.session_state["other_embedding_model"]
-
-                llama_index.create_service_context(
-                    st.session_state["llm"],
-                    st.session_state["system_prompt"],
-                    hf_embedding_model,
-                    st.session_state["chunk_size"],
-                    # st.session_state["chunk_overlap"],
-                )
-                st.caption("✔️ Context Created")
-            except Exception as err:
-                logs.log.error(f"Setting up Service Context failed: {err}")
-                error = err
-
-            #######################################
-            # Load files from the data/ directory #
-            #######################################
-
-            try:
-                documents = llama_index.load_documents(save_dir)
-                st.session_state["documents"] = documents
-                st.caption("✔️ Processed File Data")
-            except Exception as err:
-                logs.log.error(f"Document Load Error: {err}")
-                error = err
-
-            ###########################################
-            # Create an index from ingested documents #
-            ###########################################
-
-            try:
-                llama_index.create_query_engine(
-                    st.session_state["documents"], st.session_state["service_context"]
-                )
-                st.caption("✔️ Created File Index")
-            except Exception as err:
-                logs.log.error(f"Index Creation Error: {err}")
-                error = err
-
-            #####################
-            # Remove data files #
-            #####################
-
-            try:
-                save_dir = os.getcwd() + "/data"
-                shutil.rmtree(save_dir)
-                st.caption("✔️ Removed Temp Files")
-            except Exception as err:
-                logs.log.error(f"Failed to delete data files: {err}")
-                error = err
-
-            #####################
-            # Show Final Status #
-            #####################
+            # Initiate the RAG pipeline, providing documents to be saved on disk if necessary
+            error = rag.rag_pipeline(uploaded_files)
 
+            # Display errors (if any) or proceed
             if error is not None:
                 st.exception(error)
             else:

diff --git a/components/tabs/settings.py b/components/tabs/settings.py
@@ -45,13 +45,14 @@ def settings():
             st.selectbox(
                 "Chat Mode",
                 (
-                    "best",
-                    "condense_question",
-                    "context",
-                    "condense_plus_context",
-                    "react",
+                    "compact",
+                    "refine",
+                    "tree_summarize",
+                    "simple_summarize",
+                    "accumulate",
+                    "compact_accumulate",
                 ),
-                help="Sets the [Llama-Index Chat Mode](https://docs.llamaindex.ai/en/stable/module_guides/deploying/chat_engines/usage_pattern.html#available-chat-modes) used when creating the Query Engine.",
+                help="Sets the [Llama Index Query Engine chat mode](https://github.com/run-llama/llama_index/blob/main/docs/module_guides/deploying/query_engine/response_modes.md) used when creating the Query Engine. Default: `compact`.",
                 key="chat_mode",
                 disabled=True,
             )

diff --git a/components/tabs/file_upload.py → components/tabs/sources.py b/components/tabs/file_upload.py → components/tabs/sources.py
@@ -5,9 +5,9 @@
 from components.tabs.website import website
 
 
-def file_upload():
-    st.title("Directly import your files")
-    st.caption("Convert your files to embeddings for utilization during chat")
+def sources():
+    st.title("Directly import your data")
+    st.caption("Convert your data into embeddings for utilization during chat")
     st.write("")
 
     with st.expander("💻 &nbsp; **Local Files**", expanded=False):