fixed bugs in llama index custom demo and updated ui

.gitignore

@@ -4,4 +4,6 @@ models/__pycache__
 .env
 __pycache__
 vectorStores
-.vscode
+.vscode
+.streamlit/secrets.toml
+uploaded_files

app.py

@@ -8,23 +8,63 @@ import streamlit as st
 
 from utils.util import *
 
+from langchain.memory import ConversationBufferMemory
+
+SAVE_DIR = "uploaded_files"
+os.makedirs(SAVE_DIR, exist_ok=True)
+
+
+def init_session_state():
+    if "openai_api_key" not in st.session_state:
+        st.session_state.openai_api_key = ""
+
+    if "uploaded_files" not in st.session_state:
+        st.session_state.uploaded_files = os.listdir(SAVE_DIR)
+
+
+init_session_state()
+
 st.set_page_config(page_title="RegBotBeta", page_icon="📜🤖")
 
 st.title("Welcome to RegBotBeta2.0")
 st.header("Powered by `LlamaIndex🦙`, `Langchain🦜🔗 ` and `OpenAI API`")
 
-api_key = st.text_input("Enter your OpenAI API key here:", type="password")
+# openai_api_key = st.text_input(
+#     "OpenAI API Key",
+#     type="password",
+#     help="Get your API key from https://platform.openai.com/account/api-keys",
+#     value=st.session_state.openai_api_key,
+# )
+
+# isKeyValid = False
+# if openai_api_key:
+#     resp = validate(openai_api_key)
+#     if "error" in resp.json():
+#         st.info("Invalid Token! Try again.")
+#     else:
+#         st.info("Success")
+#         st.session_state.openai_api_key = openai_api_key
+#         isKeyValid = True
+
+uploaded_files = st.file_uploader(
+    "Upload Files",
+    accept_multiple_files=True,
+    type=["pdf", "docx", "txt", "csv"],
+)
+
+if uploaded_files:
+    for file in uploaded_files:
+        if file not in st.session_state.uploaded_files:
+            # add the file to session state
+            st.session_state.uploaded_files.append(file.name)
 
-if api_key:
-    resp = validate(api_key)
-    if "error" in resp.json():
-        st.info("Invalid Token! Try again.")
-    else:
-        st.info("Success")
-        os.environ["OPENAI_API_KEY"] = api_key
-        openai.api_key = api_key
+            # save the file to the sample_data directory
+            with open(os.path.join(SAVE_DIR, file.name), "wb") as f:
+                f.write(file.getbuffer())
 
-        if "openai_api_key" not in st.session_state:
-            st.session_state.openai_api_key = ""
+    st.success("File(s) uploaded successfully!")
 
-        st.session_state.openai_api_key = api_key
+if st.session_state.uploaded_files:
+    st.write("Uploaded Files:")
+    for i, filename in enumerate(st.session_state.uploaded_files, start=1):
+        st.write(f"{i}. {filename}")

models/embeddings.py

@@ -0,0 +1,15 @@
+import os
+
+from llama_index.embeddings.langchain import LangchainEmbedding
+from langchain.embeddings.huggingface import HuggingFaceEmbeddings
+from llama_index.embeddings.openai import OpenAIEmbedding
+
+hf_embed_model = HuggingFaceEmbeddings(
+    model_name="sentence-transformers/all-MiniLM-L6-v2",
+)
+
+# embed_model = LangchainEmbedding(hf_embed_model)
+
+openai_embed_model = OpenAIEmbedding(
+    api_key=os.getenv("OPENAI_API_KEY"),
+)

models/llamaCustom.py

@@ -1,6 +1,7 @@
 import os
 import pickle
 from json import dumps, loads
+import time
 from typing import Any, List, Mapping, Optional
 
 import numpy as np
@@ -9,25 +10,31 @@ import pandas as pd
 import streamlit as st
 from dotenv import load_dotenv
 from huggingface_hub import HfFileSystem
-from langchain.llms.base import LLM
-from llama_index import (
-    Document,
-    GPTVectorStoreIndex,
-    LLMPredictor,
-    PromptHelper,
-    ServiceContext,
-    SimpleDirectoryReader,
-    StorageContext,
-    load_index_from_storage,
-)
-from pydantic import BaseModel
-from llama_index.llms.base import llm_completion_callback
-from llama_index.llms import CompletionResponse, CustomLLM, LLMMetadata
-from llama_index.prompts import Prompt
+
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, Pipeline
 
+# prompts
 from assets.prompts import custom_prompts
 
+# llama index
+from llama_index.core import (
+    StorageContext,
+    SimpleDirectoryReader,
+    VectorStoreIndex,
+    load_index_from_storage,
+    PromptHelper,
+    PromptTemplate,
+)
+from llama_index.core.llms import (
+    CustomLLM,
+    CompletionResponse,
+    LLMMetadata,
+)
+from llama_index.core.memory import ChatMemoryBuffer
+from llama_index.core.llms.callbacks import llm_completion_callback
+from llama_index.core.base.llms.types import ChatMessage
+from llama_index.core import Settings
+
 load_dotenv()
 # openai.api_key = os.getenv("OPENAI_API_KEY")
 fs = HfFileSystem()
@@ -40,9 +47,41 @@ NUM_OUTPUT = 525
 # set maximum chunk overlap
 CHUNK_OVERLAP_RATION = 0.2
 
-text_qa_template = Prompt(custom_prompts.text_qa_template_str)
+# TODO: use the following prompt to format the answer at the end of the context prompt
+ANSWER_FORMAT = """
+Use the following example format for your answer:
+[FORMAT]
+Answer:
+    The answer to the user question.
+Reference:
+    The list of references to the specific sections of the documents that support your answer.    
+[END_FORMAT]
+"""
+
+CONTEXT_PROMPT_TEMPLATE = """
+The following is a friendly conversation between a user and an AI assistant.
+The assistant is talkative and provides lots of specific details from its context.
+If the assistant does not know the answer to a question, it truthfully says it
+does not know.
+
+Here are the relevant documents for the context:
 
-refine_template = Prompt(custom_prompts.refine_template_str)
+{context_str}
+
+Instruction: Based on the above documents, provide a detailed answer for the user question below. \
+Include references to the specific sections of the documents that support your answer. \
+Answer "don't know" if not present in the document.
+"""
+
+CONDENSE_PROMPT_TEMPLATE = """
+Given the following conversation between a user and an AI assistant and a follow up question from user,
+rephrase the follow up question to be a standalone question.
+
+Chat History:
+{chat_history}
+Follow Up Input: {question}
+Standalone question:
+"""
 
 
 @st.cache_resource
@@ -67,10 +106,8 @@ def load_model(model_name: str):
 
 
 class OurLLM(CustomLLM):
-    # def __init__(self, model_name: str, pipeline):
-    #     super().__init__()  # Call the __init__ method of CustomLLM
-    #     self.model_name = model_name
-    #     self.pipeline = pipeline
+    context_window: int = 3900
+    num_output: int = 256
     model_name: str = ""
     pipeline: Pipeline = None
 
@@ -83,6 +120,7 @@ class OurLLM(CustomLLM):
             model_name=self.model_name,
         )
 
+    # The decorator is optional, but provides observability via callbacks on the LLM calls.
     @llm_completion_callback()
     def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
         prompt_length = len(prompt)
@@ -93,67 +131,39 @@ class OurLLM(CustomLLM):
         return CompletionResponse(text=text)
 
     @llm_completion_callback()
-    def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
-        raise NotImplementedError()
+    def stream_complete(self, prompt: str, **kwargs: Any):
+        response = ""
+        for token in self.dummy_response:
+            response += token
+            yield CompletionResponse(text=response, delta=token)
 
 
 class LlamaCustom:
-    def __init__(self, model_name: str) -> None:
-        self.vector_index = self.initialize_index(model_name=model_name)
-
-    def initialize_index(self, model_name: str):
-        index_name = model_name.split("/")[-1]
-
-        file_path = f"./vectorStores/{index_name}"
-
-        if os.path.exists(path=file_path):
-            # rebuild storage context
-            storage_context = StorageContext.from_defaults(persist_dir=file_path)
-
-            # local load index access
-            index = load_index_from_storage(storage_context)
-
-            # huggingface repo load access
-            # with fs.open(file_path, "r") as file:
-            #     index = pickle.loads(file.readlines())
-            return index
-        else:
-            prompt_helper = PromptHelper(
-                context_window=CONTEXT_WINDOW,
-                num_output=NUM_OUTPUT,
-                chunk_overlap_ratio=CHUNK_OVERLAP_RATION,
-            )
-
-            # define llm
-            pipe = load_model(model_name=model_name)
-            llm = OurLLM(model_name=model_name, pipeline=pipe)
-
-            llm_predictor = LLMPredictor(llm=llm)
-            service_context = ServiceContext.from_defaults(
-                llm_predictor=llm_predictor, prompt_helper=prompt_helper
-            )
-
-            # documents = prepare_data(r"./assets/regItems.json")
-            documents = SimpleDirectoryReader(input_dir="./assets/pdf").load_data()
-
-            index = GPTVectorStoreIndex.from_documents(
-                documents, service_context=service_context
-            )
-
-            # local write access
-            index.storage_context.persist(file_path)
-
-            # huggingface repo write access
-            # with fs.open(file_path, "w") as file:
-            #     file.write(pickle.dumps(index))
-            return index
-
-    def get_response(self, query_str):
-        print("query_str: ", query_str)
-        # query_engine = self.vector_index.as_query_engine()
-        query_engine = self.vector_index.as_query_engine(
-            text_qa_template=text_qa_template, refine_template=refine_template
+    def __init__(self, model_name: str, index: VectorStoreIndex):
+        self.model_name = model_name
+        self.index = index
+        self.chat_mode = "condense_plus_context"
+        self.memory = ChatMemoryBuffer.from_defaults()
+
+    def get_response(self, query_str: str, chat_history: List[ChatMessage]):
+        # https://docs.llamaindex.ai/en/stable/module_guides/deploying/chat_engines/
+        # query_engine = self.index.as_query_engine(
+        #     text_qa_template=text_qa_template, refine_template=refine_template
+        # )
+        chat_engine = self.index.as_chat_engine(
+            chat_mode=self.chat_mode,
+            memory=self.memory,
+            context_prompt=CONTEXT_PROMPT_TEMPLATE,
+            condense_prompt=CONDENSE_PROMPT_TEMPLATE,
+            # verbose=True,
         )
-        response = query_engine.query(query_str)
-        print("metadata: ", response.metadata)
+        # response = query_engine.query(query_str)
+        response = chat_engine.chat(message=query_str, chat_history=chat_history)
+
         return str(response)
+
+    def get_stream_response(self, query_str: str, chat_history: List[ChatMessage]):
+        response = self.get_response(query_str=query_str, chat_history=chat_history)
+        for word in response.split():
+            yield word + " "
+            time.sleep(0.05)

models/llms.py

@@ -0,0 +1,44 @@
+from llama_index.llms.huggingface import HuggingFaceLLM, HuggingFaceInferenceAPI
+from llama_index.llms.openai import OpenAI
+from llama_index.llms.replicate import Replicate
+from dotenv import load_dotenv
+import os
+
+load_dotenv()
+
+llm_mixtral_8x7b = HuggingFaceInferenceAPI(
+    model_name="mistralai/Mixtral-8x7B-Instruct-v0.1",
+    token=os.getenv("HUGGINGFACE_API_TOKEN"),
+)
+
+# download the model from the Hugging Face Hub and run it locally
+# llm_mixtral_8x7b = HuggingFaceLLM(model_name="mistralai/Mixtral-8x7B-Instruct-v0.1")
+
+llm_llama_2_7b_chat = HuggingFaceInferenceAPI(
+    model_name="meta-llama/Llama-2-7b-chat-hf",
+    token=os.getenv("HUGGINGFACE_API_TOKEN"),
+)
+
+llm_bloomz_560m = HuggingFaceInferenceAPI(
+    model_name="bigscience/bloomz-560m",
+    token=os.getenv("HUGGINGFACE_API_TOKEN"),
+)
+
+llm_gpt_3_5_turbo = OpenAI(
+    api_key=os.getenv("OPENAI_API_KEY"),
+)
+
+llm_gpt_3_5_turbo_0125 = OpenAI(
+    model="gpt-3.5-turbo-0125",
+    api_key="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",
+)
+
+llm_gpt_4_0125 = OpenAI(
+    model="gpt-4-0125-preview",
+    api_key=os.getenv("OPENAI_API_KEY"),
+)
+
+llm_llama_13b_v2_replicate = Replicate(
+    model="meta/llama-2-13b-chat",
+    prompt_key=os.getenv("REPLICATE_API_KEY"),
+)

pages/llama_custom_demo.py

@@ -1,26 +1,158 @@
+import random
+import time
+import streamlit as st
 import os
+import pathlib
+from typing import List
+from models.llms import (
+    llm_llama_2_7b_chat,
+    llm_mixtral_8x7b,
+    llm_bloomz_560m,
+    llm_gpt_3_5_turbo,
+    llm_gpt_3_5_turbo_0125,
+    llm_gpt_4_0125,
+    llm_llama_13b_v2_replicate
+)
+from models.embeddings import hf_embed_model, openai_embed_model
+from models.llamaCustom import LlamaCustom
 
-import openai
-import streamlit as st
+# from models.llamaCustom import LlamaCustom
+from utils.chatbox import show_previous_messages, show_chat_input
+from llama_index.core import (
+    SimpleDirectoryReader,
+    Document,
+    VectorStoreIndex,
+    StorageContext,
+    Settings,
+    load_index_from_storage,
+)
+from llama_index.core.memory import ChatMemoryBuffer
+from llama_index.core.base.llms.types import ChatMessage
 
-from models.llamaCustom import LlamaCustom
-from utils.chatbox import chatbox
+SAVE_DIR = "uploaded_files"
+VECTOR_STORE_DIR = "vectorStores"
+
+# global
+Settings.embed_model = hf_embed_model
+
+llama_llms = {
+    "bigscience/bloomz-560m": llm_bloomz_560m,
+    "mistral/mixtral": llm_mixtral_8x7b,
+    "meta-llama/Llama-2-7b-chat-hf": llm_llama_2_7b_chat,
+    # "openai/gpt-3.5-turbo": llm_gpt_3_5_turbo,
+    "openai/gpt-3.5-turbo-0125": llm_gpt_3_5_turbo_0125,
+    # "openai/gpt-4-0125-preview": llm_gpt_4_0125,
+    # "meta/llama-2-13b-chat": llm_llama_13b_v2_replicate,
+}
+
+def init_session_state():
+    if "llama_messages" not in st.session_state:
+        st.session_state.llama_messages = [
+            {"role": "assistant", "content": "How can I help you today?"}
+        ]
+
+    # TODO: create a chat history for each different document
+    if "llama_chat_history" not in st.session_state:
+        st.session_state.llama_chat_history = [
+            ChatMessage.from_str(role="assistant", content="How can I help you today?")
+        ]
+
+    if "llama_custom" not in st.session_state:
+        st.session_state.llama_custom = None
+
+# @st.cache_resource
+def index_docs(
+    filename: str,
+) -> VectorStoreIndex:
+    try:
+        index_path = pathlib.Path(f"{VECTOR_STORE_DIR}/{filename.replace('.', '_')}")
+        if pathlib.Path.exists(index_path):
+            print("Loading index from storage ...")
+            storage_context = StorageContext.from_defaults(persist_dir=index_path)
+            index = load_index_from_storage(storage_context=storage_context)
+
+            # test the index
+            index.as_query_engine().query("What is the capital of France?")
+            
+        else:
+            reader = SimpleDirectoryReader(input_files=[f"{SAVE_DIR}/{filename}"])
+            docs = reader.load_data(show_progress=True)
+            index = VectorStoreIndex.from_documents(
+                documents=docs,
+                show_progress=True,
+            )
+            index.storage_context.persist(persist_dir=f"vectorStores/{filename.replace(".", '_')}")
+
+    except Exception as e:
+        print(f"Error: {e}")
+        index = None
+    return index
+
+
+def load_llm(model_name: str):
+    return llama_llms[model_name]
+
+init_session_state()
 
 st.set_page_config(page_title="Llama", page_icon="🦙")
 
-st.subheader("Llama Index with Custom LLM Demo")
+st.header("Llama Index with Custom LLM Demo")
+
+tab1, tab2 = st.tabs(["Config", "Chat"])
+
+with tab1:
+    with st.form(key="llama_form"):
+        selected_llm_name = st.selectbox(label="Select a model:", options=llama_llms.keys())
+
+        if selected_llm_name.startswith("openai"):
+            # ask for the api key
+            if st.secrets.get("OPENAI_API_KEY") is None:
+                # st.stop()
+                st.info("OpenAI API Key not found in secrets. Please enter it below.")
+                st.secrets["OPENAI_API_KEY"] = st.text_input(
+                    "OpenAI API Key",
+                    type="password",
+                    help="Get your API key from https://platform.openai.com/account/api-keys",
+                )
+
+        selected_file = st.selectbox(
+            label="Choose a file to chat with: ", options=os.listdir(SAVE_DIR)
+        )
+
+        if st.form_submit_button(label="Submit"):
+            with st.status("Loading ...", expanded=True) as status:
+                st.write("Loading Model ...")
+                llama_llm = load_llm(selected_llm_name)
+                Settings.llm = llama_llm
+
+                st.write("Processing Data ...")
+                index = index_docs(selected_file)
+                if index is None:
+                    st.error("Failed to index the documents.")
+                    st.stop()
+
+                st.write("Finishing Up ...")
+                llama_custom = LlamaCustom(model_name=selected_llm_name, index=index)
+                st.session_state.llama_custom = llama_custom
+
+                status.update(label="Ready to query!", state="complete", expanded=False)
 
-if "messages" not in st.session_state:
-    st.session_state.messages = []
+with tab2:
+    messages_container = st.container(height=300)
+    show_previous_messages(framework="llama", messages_container=messages_container)
+    show_chat_input(disabled=False, framework="llama", model=st.session_state.llama_custom, messages_container=messages_container)
 
-if "openai_api_key" not in st.session_state:
-    st.info("Enter your openai key to access the chatbot.")
-else:
-    option = st.selectbox(
-        label="Select your model:", options=("bigscience/bloom-560m",)
-    )
+    def clear_history():
+        messages_container.empty()
+        st.session_state.llama_messages = [
+            {"role": "assistant", "content": "How can I help you today?"}
+        ]
 
-    # with st.spinner("Initializing vector index"):
-    model = LlamaCustom(model_name=option)
+        st.session_state.llama_chat_history = [
+            ChatMessage.from_str(role="assistant", content="How can I help you today?")
+        ]
 
-    chatbox("llama_custom", model)
+    if st.button("Clear Chat History"):
+        clear_history()
+        st.rerun()
+    

requirements.txt

@@ -1,14 +1,15 @@
-llama_index==0.8.64.post1
-torch
+llama_index>=0.10.27
 transformers
-panda
-numpy
-langchain
+pandas
+langchain>=0.1.11
 openai
 faiss-cpu
 python-dotenv
 streamlit>=1.24.0
-huggingface_hub
-xformers
+huggingface_hub>=0.21.4
 pypdf
-pymupdf
+llama-index-llms-huggingface>=0.1.4
+llama-index-embeddings-langchain>=0.1.2
+replicate>=0.25.1
+llama-index-llms-replicate
+sentence-transformers>=2.6.1

utils/chatbox.py

@@ -1,93 +1,61 @@
 import time
-
+import json
 import streamlit as st
+from typing import Dict, List, Any
 
-def display_chat_history(model_name: str):
-    for message in st.session_state[model_name]:
-        with st.chat_message(message["role"]):
-            st.markdown(message["content"])
-
-def chat_input(model_name: str):
-    if prompt := st.chat_input("Say something"):
-        # Display user message in chat message container
-        st.chat_message("user").markdown(prompt)
-
-        # Add user message to chat history
-        st.session_state[model_name].append({"role": "user", "content": prompt})
-
-        return prompt
-
-def display_bot_msg(model_name: str, bot_response: str):
-    # Display assistant response in chat message container
-    with st.chat_message("assistant"):
-        message_placeholder = st.empty()
-        full_response = ""
-
-        # simulate the chatbot "thinking" before responding
-        # (or stream its response)
-        for chunk in bot_response.split():
-            full_response += chunk + " "
-            time.sleep(0.05)
-
-            # add a blinking cursor to simulate typing
-            message_placeholder.markdown(full_response + "▌")
+from llama_index.core.base.llms.types import ChatMessage
 
-        message_placeholder.markdown(full_response)
-        # st.markdown(response)
 
-    # Add assistant response to chat history
-    st.session_state[model_name].append(
-        {"model_name": model_name, "role": "assistant", "content": full_response}
-    )
-
-def chatbox(model_name: str, model: None):
-    # Display chat messages from history on app rerun
-    for message in st.session_state.messages:
-        if (message["model_name"] == model_name):
+def show_previous_messages(framework: str, messages_container: any):
+    with messages_container:
+        messages: List[Dict[str, Any]] = st.session_state[f"{framework}_messages"]
+        for message in messages:
             with st.chat_message(message["role"]):
                 st.markdown(message["content"])
 
-    if prompt := st.chat_input("Say something"):
-        # Display user message in chat message container
-        st.chat_message("user").markdown(prompt)
-
-        # Add user message to chat history
-        st.session_state.messages.append({"model_name": model_name, "role": "user", "content": prompt})
-
-        with st.spinner("Processing your query..."):
-            bot_response = model.get_response(prompt)
-
-        print("bot: ", bot_response)
 
-        # Display assistant response in chat message container
-        with st.chat_message("assistant"):
-            message_placeholder = st.empty()
-            full_response = ""
+def show_chat_input(
+    disabled: bool, framework: str, model: any, messages_container: any
+):
+    if disabled:
+        st.info("Make sure to select a model and file to start chatting!")
 
-            # simulate the chatbot "thinking" before responding
-            # (or stream its response)
-            for chunk in bot_response.split():
-                full_response += chunk + " "
-                time.sleep(0.05)
-
-                # add a blinking cursor to simulate typing
-                message_placeholder.markdown(full_response + "▌")
-
-            message_placeholder.markdown(full_response)
-            # st.markdown(response)
+    if prompt := st.chat_input("Say something", disabled=disabled):
+        st.session_state[f"{framework}_messages"].append(
+            {"role": "user", "content": prompt}
+        )
 
-        # Add assistant response to chat history
-        st.session_state.messages.append(
-            {"model_name": model_name, "role": "assistant", "content": full_response}
+        st.session_state[f"{framework}_chat_history"].append(
+            ChatMessage.from_str(role="user", content=prompt)
         )
 
-        # Scroll to the bottom of the chat container
-        # st.markdown(
-        #     """
-        #     <script>
-        #     const chatContainer = document.getElementsByClassName("css-1n76uvr")[0];
-        #     chatContainer.scrollTop = chatContainer.scrollHeight;
-        #     </script>
-        #     """,
-        #     unsafe_allow_html=True,
-        # )
+        # if st.session_state[f"{framework}_messages"][-1]["role"] == "assistant":
+        with messages_container:
+            with st.chat_message("user"):
+                st.write(prompt)
+
+            with st.chat_message("assistant"):
+                with st.spinner("Thinking..."):
+                    try:
+                        ai_response = model.get_response(
+                            query_str=prompt,
+                            chat_history=st.session_state[f"{framework}_chat_history"],
+                        )
+                        # when streaming, the response format is gone
+                        # ai_response = model.get_stream_response(
+                        #     query_str=prompt,
+                        #     chat_history=st.session_state[f"{framework}_chat_history"],
+                        # )
+                    except Exception as e:
+                        ai_response = f"An error occurred: {e}"
+
+                    st.write(ai_response)
+                    # response = st.write_stream(ai_response)
+
+                    st.session_state[f"{framework}_messages"].append(
+                        {"role": "assistant", "content": ai_response}
+                    )
+
+                    st.session_state[f"{framework}_chat_history"].append(
+                        ChatMessage.from_str(role="assistant", content=ai_response)
+                    )