llama-index-update (#1)

- updated llama index demo (b83dc9cf5500e038115cf835e0c466eb12b30905)
- Merge branch 'main' of https://huggingface.co/spaces/hbui/RegBot4.0 (e262fe5388f1bbd46030de0fafcb1f1b1b32b6db)
- updated llm configs and prompts (ab0abe29e798a3f31c482100f4139634171c2fa7)

app.py

@@ -1,13 +1,9 @@
-# https://docs.streamlit.io/knowledge-base/tutorials/build-conversational-apps
-
 import os
-
 import openai
 import requests
 import streamlit as st
 
 from utils.util import *
-
 from langchain.memory import ConversationBufferMemory
 
 SAVE_DIR = "uploaded_files"
@@ -17,30 +13,24 @@ 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)
+    if "huggingface_token" not in st.session_state:
+        st.session_state.huggingface_token = ""
 
 
 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`")
+st.header("Powered by `LlamaIndex🦙`, `Langchain🦜🔗` and `OpenAI API`")
 
 
-def init_session_state():
-    if "huggingface_token" not in st.session_state:
-        st.session_state.huggingface_token = ""
-
-
-init_session_state()
-
 uploaded_files = st.file_uploader(
     "Upload Files",
     accept_multiple_files=True,
     type=["pdf", "docx", "txt", "csv"],
+    label_visibility="hidden",
 )
 
 if uploaded_files:
@@ -48,14 +38,27 @@ if 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)
-
             # save the file to the sample_data directory
             with open(os.path.join(SAVE_DIR, file.name), "wb") as f:
                 f.write(file.getbuffer())
-
     st.success("File(s) uploaded successfully!")
 
+
+def delete_file(filename):
+    """Delete file from session state and local filesystem."""
+    if filename in st.session_state.uploaded_files and os.path.exists(
+        os.path.join(SAVE_DIR, filename)
+    ):
+        st.session_state.uploaded_files.remove(filename)
+        os.remove(os.path.join(SAVE_DIR, filename))
+        st.success(f"Deleted {filename}!")
+        st.rerun()
+
+
 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}")
+    for index, filename in enumerate(st.session_state.uploaded_files):
+        col1, col2 = st.columns([4, 1])
+        col1.write(filename)
+        if col2.button("Delete", key=f"delete_{index}"):
+            delete_file(filename)

models/llamaCustom.py

@@ -18,11 +18,7 @@ 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 (
@@ -47,14 +43,12 @@ NUM_OUTPUT = 525
 # set maximum chunk overlap
 CHUNK_OVERLAP_RATION = 0.2
 
-# 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:
+Provide the answer to the user question in the following format:
 [FORMAT]
-Answer:
-    The answer to the user question.
+Your answer to the user question above.
 Reference:
-    The list of references to the specific sections of the documents that support your answer.    
+    The list of references (such as page number, title, chapter, section) to the specific sections of the documents that support your answer.
 [END_FORMAT]
 """
 
@@ -184,9 +178,13 @@ class LlamaCustom:
 
     def get_response(self, query_str: str, chat_history: List[ChatMessage]):
         # https://docs.llamaindex.ai/en/stable/module_guides/deploying/chat_engines/
+        # https://docs.llamaindex.ai/en/stable/examples/query_engine/citation_query_engine/
+        # https://docs.llamaindex.ai/en/stable/examples/query_engine/knowledge_graph_rag_query_engine/
         query_engine = self.index.as_query_engine(
-            text_qa_template=PromptTemplate(QUERY_ENGINE_QA_TEMPLATE),
-            refine_template=PromptTemplate(QUERY_ENGINE_REFINE_TEMPLATE),
+            text_qa_template=PromptTemplate(QUERY_ENGINE_QA_TEMPLATE + ANSWER_FORMAT),
+            refine_template=PromptTemplate(
+                QUERY_ENGINE_REFINE_TEMPLATE
+            ),  # passing ANSWER_FORMAT here will not give the desired output, need to use the output parser from llama index?
             verbose=self.verbose,
         )
         # chat_engine = self.index.as_chat_engine(
@@ -196,6 +194,7 @@ class LlamaCustom:
         #     condense_prompt=CHAT_ENGINE_CONDENSE_PROMPT_TEMPLATE,
         #     # verbose=True,
         # )
+
         response = query_engine.query(query_str)
         # response = chat_engine.chat(message=query_str, chat_history=chat_history)
 

models/llamaCustomV2.py

@@ -0,0 +1,229 @@
+import os
+import time
+from llama_index.core import VectorStoreIndex
+from llama_index.core.query_pipeline import (
+    QueryPipeline,
+    InputComponent,
+    ArgPackComponent,
+)
+from llama_index.core.prompts import PromptTemplate
+from llama_index.llms.openai import OpenAI
+from llama_index.postprocessor.colbert_rerank import ColbertRerank
+from typing import Any, Dict, List, Optional
+from llama_index.core.bridge.pydantic import Field
+from llama_index.core.llms import ChatMessage
+from llama_index.core.query_pipeline import CustomQueryComponent
+from llama_index.core.schema import NodeWithScore
+from llama_index.core.memory import ChatMemoryBuffer
+
+
+llm = OpenAI(
+    model="gpt-3.5-turbo-0125",
+    api_key=os.getenv("OPENAI_API_KEY"),
+)
+
+# First, we create an input component to capture the user query
+input_component = InputComponent()
+
+# Next, we use the LLM to rewrite a user query
+rewrite = (
+    "Please write a query to a semantic search engine using the current conversation.\n"
+    "\n"
+    "\n"
+    "{chat_history_str}"
+    "\n"
+    "\n"
+    "Latest message: {query_str}\n"
+    'Query:"""\n'
+)
+rewrite_template = PromptTemplate(rewrite)
+
+# we will retrieve two times, so we need to pack the retrieved nodes into a single list
+argpack_component = ArgPackComponent()
+
+# then postprocess/rerank with Colbert
+reranker = ColbertRerank(top_n=3)
+
+DEFAULT_CONTEXT_PROMPT = (
+    "Here is some context that may be relevant:\n"
+    "-----\n"
+    "{node_context}\n"
+    "-----\n"
+    "Please write a response to the following question, using the above context:\n"
+    "{query_str}\n"
+    "Please formate your response in the following way:\n"
+    "Your answer here.\n"
+    "Reference:\n"
+    "    Your references here (e.g. page numbers, titles, etc.).\n"
+)
+
+
+class ResponseWithChatHistory(CustomQueryComponent):
+    llm: OpenAI = Field(..., description="OpenAI LLM")
+    system_prompt: Optional[str] = Field(
+        default=None, description="System prompt to use for the LLM"
+    )
+    context_prompt: str = Field(
+        default=DEFAULT_CONTEXT_PROMPT,
+        description="Context prompt to use for the LLM",
+    )
+
+    def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+        """Validate component inputs during run_component."""
+        # NOTE: this is OPTIONAL but we show you where to do validation as an example
+        return input
+
+    @property
+    def _input_keys(self) -> set:
+        """Input keys dict."""
+        # NOTE: These are required inputs. If you have optional inputs please override
+        # `optional_input_keys_dict`
+        return {"chat_history", "nodes", "query_str"}
+
+    @property
+    def _output_keys(self) -> set:
+        return {"response"}
+
+    def _prepare_context(
+        self,
+        chat_history: List[ChatMessage],
+        nodes: List[NodeWithScore],
+        query_str: str,
+    ) -> List[ChatMessage]:
+        node_context = ""
+        for idx, node in enumerate(nodes):
+            node_text = node.get_content(metadata_mode="llm")
+            node_context += f"Context Chunk {idx}:\n{node_text}\n\n"
+
+        formatted_context = self.context_prompt.format(
+            node_context=node_context, query_str=query_str
+        )
+        user_message = ChatMessage(role="user", content=formatted_context)
+
+        chat_history.append(user_message)
+
+        if self.system_prompt is not None:
+            chat_history = [
+                ChatMessage(role="system", content=self.system_prompt)
+            ] + chat_history
+
+        return chat_history
+
+    def _run_component(self, **kwargs) -> Dict[str, Any]:
+        """Run the component."""
+        chat_history = kwargs["chat_history"]
+        nodes = kwargs["nodes"]
+        query_str = kwargs["query_str"]
+
+        prepared_context = self._prepare_context(chat_history, nodes, query_str)
+
+        response = llm.chat(prepared_context)
+
+        return {"response": response}
+
+    async def _arun_component(self, **kwargs: Any) -> Dict[str, Any]:
+        """Run the component asynchronously."""
+        # NOTE: Optional, but async LLM calls are easy to implement
+        chat_history = kwargs["chat_history"]
+        nodes = kwargs["nodes"]
+        query_str = kwargs["query_str"]
+
+        prepared_context = self._prepare_context(chat_history, nodes, query_str)
+
+        response = await llm.achat(prepared_context)
+
+        return {"response": response}
+
+
+class LlamaCustomV2:
+    response_component = ResponseWithChatHistory(
+        llm=llm,
+        system_prompt=(
+            "You are a Q&A system. You will be provided with the previous chat history, "
+            "as well as possibly relevant context, to assist in answering a user message."
+        ),
+    )
+
+    def __init__(self, model_name: str, index: VectorStoreIndex):
+        self.model_name = model_name
+        self.index = index
+        self.retriever = index.as_retriever()
+        self.chat_mode = "condense_plus_context"
+        self.memory = ChatMemoryBuffer.from_defaults()
+        self.verbose = True
+        self._build_pipeline()
+
+    def _build_pipeline(self):
+        self.pipeline = QueryPipeline(
+            modules={
+                "input": input_component,
+                "rewrite_template": rewrite_template,
+                "llm": llm,
+                "rewrite_retriever": self.retriever,
+                "query_retriever": self.retriever,
+                "join": argpack_component,
+                "reranker": reranker,
+                "response_component": self.response_component,
+            },
+            verbose=self.verbose,
+        )
+        # run both retrievers -- once with the hallucinated query, once with the real query
+        self.pipeline.add_link(
+            "input", "rewrite_template", src_key="query_str", dest_key="query_str"
+        )
+        self.pipeline.add_link(
+            "input",
+            "rewrite_template",
+            src_key="chat_history_str",
+            dest_key="chat_history_str",
+        )
+        self.pipeline.add_link("rewrite_template", "llm")
+        self.pipeline.add_link("llm", "rewrite_retriever")
+        self.pipeline.add_link("input", "query_retriever", src_key="query_str")
+
+        # each input to the argpack component needs a dest key -- it can be anything
+        # then, the argpack component will pack all the inputs into a single list
+        self.pipeline.add_link("rewrite_retriever", "join", dest_key="rewrite_nodes")
+        self.pipeline.add_link("query_retriever", "join", dest_key="query_nodes")
+
+        # reranker needs the packed nodes and the query string
+        self.pipeline.add_link("join", "reranker", dest_key="nodes")
+        self.pipeline.add_link(
+            "input", "reranker", src_key="query_str", dest_key="query_str"
+        )
+
+        # synthesizer needs the reranked nodes and query str
+        self.pipeline.add_link("reranker", "response_component", dest_key="nodes")
+        self.pipeline.add_link(
+            "input", "response_component", src_key="query_str", dest_key="query_str"
+        )
+        self.pipeline.add_link(
+            "input",
+            "response_component",
+            src_key="chat_history",
+            dest_key="chat_history",
+        )
+
+    def get_response(self, query_str: str, chat_history: List[ChatMessage]):
+        chat_history = self.memory.get()
+        char_history_str = "\n".join([str(x) for x in chat_history])
+
+        response = self.pipeline.run(
+            query_str=query_str,
+            chat_history=chat_history,
+            chat_history_str=char_history_str,
+        )
+
+        user_msg = ChatMessage(role="user", content=query_str)
+        print("user_msg: ", str(user_msg))
+        print("response: ", str(response.message))
+        self.memory.put(user_msg)
+        self.memory.put(response.message)
+
+        return str(response.message)
+
+    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

@@ -35,6 +35,7 @@ def load_llm(model_name: str, source: str = "huggingface"):
             llm_gpt_3_5_turbo_0125 = OpenAI(
                 model=model_name,
                 api_key=st.session_state.openai_api_key,
+                temperature=0.0,
             )
 
             return llm_gpt_3_5_turbo_0125
@@ -45,6 +46,7 @@ def load_llm(model_name: str, source: str = "huggingface"):
                 is_chat_model=True,
                 additional_kwargs={"max_new_tokens": 250},
                 prompt_key=st.session_state.replicate_api_token,
+                temperature=0.0,
             )
 
             return llm_llama_13b_v2_replicate

models/vector_database.py

@@ -0,0 +1,34 @@
+from pinecone import Pinecone, ServerlessSpec
+from llama_index.vector_stores.pinecone import PineconeVectorStore
+from dotenv import load_dotenv
+
+import os
+
+load_dotenv()
+
+# Pinecone Vector Database
+pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY"))
+pc_index_name = "llama-integration-pinecone"
+# pc_index_name = "openai-embeddings"
+pc_indexes = pc.list_indexes()
+
+# Check if the index already exists
+def index_exists(index_name):
+    for index in pc_indexes:
+        if index["name"] == index_name:
+            return True
+    return False
+
+# Create the index if it doesn't exist
+if not index_exists(pc_index_name):
+    pc.create_index(
+        name=pc_index_name,
+        dimension=1536,
+        spec=ServerlessSpec(cloud="aws", region="us-east-1"),
+    )
+
+# Initialize your index
+pinecone_index = pc.Index(pc_index_name)
+
+# Define the vector store
+pinecone_vector_store = PineconeVectorStore(pinecone_index=pinecone_index)

pages/llama_custom_demo.py

@@ -7,6 +7,9 @@ from typing import List
 from models.llms import load_llm, integrated_llms
 from models.embeddings import hf_embed_model, openai_embed_model
 from models.llamaCustom import LlamaCustom
+from models.llamaCustomV2 import LlamaCustomV2
+
+# from models.vector_database import pinecone_vector_store
 from utils.chatbox import show_previous_messages, show_chat_input
 from utils.util import validate_openai_api_key
 
@@ -30,7 +33,8 @@ VECTOR_STORE_DIR = "vectorStores"
 HF_REPO_ID = "zhtet/RegBotBeta"
 
 # global
-Settings.embed_model = hf_embed_model
+# Settings.embed_model = hf_embed_model
+Settings.embed_model = openai_embed_model
 
 # huggingface api
 hf_api = HfApi()
@@ -62,9 +66,10 @@ def init_session_state():
 
 
 # @st.cache_resource
-def index_docs(
+def get_index(
     filename: str,
 ) -> VectorStoreIndex:
+    """This function loads the index from storage if it exists, otherwise it creates a new index from the document."""
     try:
         index_path = pathlib.Path(f"{VECTOR_STORE_DIR}/{filename.replace('.', '_')}")
         if pathlib.Path.exists(index_path):
@@ -89,6 +94,23 @@ def index_docs(
     return index
 
 
+# def get_pinecone_index(filename: str) -> VectorStoreIndex:
+#     """Thie function loads the index from Pinecone if it exists, otherwise it creates a new index from the document."""
+#     reader = SimpleDirectoryReader(input_files=[f"{SAVE_DIR}/{filename}"])
+#     docs = reader.load_data(show_progress=True)
+#     storage_context = StorageContext.from_defaults(vector_store=pinecone_vector_store)
+#     index = VectorStoreIndex.from_documents(
+#         documents=docs, show_progress=True, storage_context=storage_context
+#     )
+
+#     return index
+
+
+def get_chroma_index(filename: str) -> VectorStoreIndex:
+    """This function loads the index from Chroma if it exists, otherwise it creates a new index from the document."""
+    pass
+
+
 def check_api_key(model_name: str, source: str):
     if source.startswith("openai"):
         if not st.session_state.openai_api_key:
@@ -164,6 +186,13 @@ with tab1:
         label="Choose a file to chat with: ", options=os.listdir(SAVE_DIR)
     )
 
+    if st.button("Clear all api keys"):
+        st.session_state.openai_api_key = ""
+        st.session_state.replicate_api_token = ""
+        st.session_state.hf_token = ""
+        st.success("All API keys cleared!")
+        st.rerun()
+
     if st.button("Submit", key="submit", help="Submit the form"):
         with st.status("Loading ...", expanded=True) as status:
             try:
@@ -176,10 +205,12 @@ with tab1:
                 Settings.llm = llama_llm
 
                 st.write("Processing Data ...")
-                index = index_docs(selected_file)
+                index = get_index(selected_file)
+                # index = get_pinecone_index(selected_file)
 
                 st.write("Finishing Up ...")
                 llama_custom = LlamaCustom(model_name=selected_llm_name, index=index)
+                # llama_custom = LlamaCustomV2(model_name=selected_llm_name, index=index)
                 st.session_state.llama_custom = llama_custom
 
                 status.update(label="Ready to query!", state="complete", expanded=False)

requirements.txt

@@ -7,11 +7,14 @@ langchain_pinecone
 openai
 faiss-cpu
 python-dotenv
-streamlit==1.29.0
+streamlit>=1.24.0
 huggingface_hub<0.21.0
 pypdf
 llama-index-llms-huggingface>=0.1.4
 llama-index-embeddings-langchain>=0.1.2
+llama-index-vector-stores-pinecone
+pinecone-client>=3.0.0
 replicate>=0.25.1
 llama-index-llms-replicate
-sentence-transformers>=2.6.1
+sentence-transformers>=2.6.1
+llama-index-postprocessor-colbert-rerank