Upload 14 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+chroma/chroma.sqlite3 filter=lfs diff=lfs merge=lfs -text
+data/uber-10-k-2023.pdf filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,72 @@
+import streamlit as st
+from database import main as populate_db
+from query import query_rag
+from langchain.schema.document import Document
+
+
+# Set the page configuration for Streamlit
+st.set_page_config(page_title='Content Engine')
+
+# Function to load and populate the database
+def load_and_populate_database():
+    st.header("Database Population")
+    st.write("Use this section to populate the database.")
+    if st.button("Populate Database"):
+        populate_db()
+
+    st.markdown("<hr style='border: 1px dashed #e0e0e0;'>", unsafe_allow_html=True)  
+
+
+# Function to display the scrolling chat box
+def display_chat_box():
+    st.header("Chat Box")
+    st.write("Enter your questions below and get responses generated from the context.")
+
+    # Setup for holding old messages
+    if 'messages' not in st.session_state:
+        st.session_state.messages = []
+
+    # Display all the history in a scrollable chat box
+    with st.container():
+        for message in st.session_state.messages:
+            st.chat_message(message['role']).markdown(message['content'])
+
+    # Input field for user query
+    user_query = st.chat_input("Your Query:", key="query_input")
+
+    if user_query:
+        # Display user prompt
+        st.chat_message('user').markdown(user_query)
+
+        # Store the user prompt in state
+        st.session_state.messages.append({'role': 'user', 'content': user_query})
+
+        # Buffering message while waiting for response
+        with st.spinner('Waiting for response...'):
+            # Query processing
+            response_text = query_rag(user_query)
+
+        # Show the LLM response
+        st.chat_message('assistant').markdown(response_text)
+
+        # Store the LLM response
+        st.session_state.messages.append({'role': 'assistant', 'content': response_text})
+
+# Main function to run the Streamlit app
+def main():
+    # Title and description
+    st.markdown("<h1 style='font-size: 60px; text-align: center;'>🧠Content Engine🤖</h1>", unsafe_allow_html=True)
+    # st.markdown("<hr style='border: 1px dashed #e0e0e0; '>", unsafe_allow_html=True)  
+    # st.write("This system analyzes and compares multiple PDF documents, specifically identifying and highlighting their differences. The system will utilize Retrieval Augmented Generation (RAG) techniques to effectively retrieve, assess, and generate insights from the documents.")
+    st.markdown("<hr style='border: 1px dashed #e0e0e0;'>", unsafe_allow_html=True)  
+    
+    # To load and populate the database
+    # load_and_populate_database()
+    
+    # Display chat box
+    display_chat_box()
+
+    
+
+if __name__ == "__main__":
+    main()

chroma/360ffad6-f56e-46e1-8b3c-e88a73c8a70d/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2cfeb32effd3397c9d0aae0a629aa165a2dffe3321940cf0abc5e07b9c5b3466
+size 6424000

chroma/360ffad6-f56e-46e1-8b3c-e88a73c8a70d/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b56eb1d0439027f400d145f9eff65f500b8c5d1afb597348ce4bc35443d3bd60
+size 100

chroma/360ffad6-f56e-46e1-8b3c-e88a73c8a70d/index_metadata.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:827b4e3f7d764e876cfb1337993477583d5f230f7748ef0352e04963f97f8265
+size 102041

chroma/360ffad6-f56e-46e1-8b3c-e88a73c8a70d/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d44dbe88b0e3a2b082ab89d9f9f05498dd3633b19061c8a82c774f2d80a8041
+size 8000

chroma/360ffad6-f56e-46e1-8b3c-e88a73c8a70d/link_lists.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:561ae3ffa1abbf2695cb93002ce011c1bf0eca135b6521c77ff6c4782b2ea122
+size 17044

chroma/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:334e569b55a158a70f6b2c1c410c2326719a1f8808215ccc4f21a38fa9cd7cea
+size 19922944

data/uber-10-k-2023.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9fc1d0c5f2e2095798c99de2d892a9e785b9ed8329090c0c376e39a22e71ad06
+size 1835051

database.py

@@ -0,0 +1,103 @@
+import argparse
+import os
+import shutil
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain.schema.document import Document
+from embedding_function import get_embedding
+from langchain_community.document_loaders import PyPDFDirectoryLoader
+from langchain_community.vectorstores import Chroma
+
+# Define paths for the Chroma database and the data directory
+CHROMA_PATH = "chroma"
+DATA_PATH = "data"
+
+def main():
+    # Set up command-line argument parsing
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--reset", action="store_true", help="Reset the database.")
+    args = parser.parse_args()
+
+    # Clear the database if the reset flag is set
+    if args.reset:
+        print("✨ Clearing Database")
+        clear_database()
+
+    # Load documents, split them into chunks, and add them to Chroma
+    documents = load_documents()
+    chunks = split_documents(documents)
+    add_to_chroma(chunks)
+
+def load_documents():
+    # Load PDF documents from the specified directory
+    document_loader = PyPDFDirectoryLoader(DATA_PATH)
+    return document_loader.load()
+
+def split_documents(documents: list[Document]):
+    # Split documents into smaller chunks using a character-based splitter
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=800,
+        chunk_overlap=80,
+        length_function=len,
+        is_separator_regex=False,
+    )
+    return text_splitter.split_documents(documents)
+
+def add_to_chroma(chunks: list[Document]):
+    # Initialize the Chroma vector store with the embedding function
+    db = Chroma(
+        persist_directory=CHROMA_PATH, embedding_function=get_embedding()
+    )
+
+    # Calculate unique IDs for each chunk
+    chunks_with_ids = calculate_chunk_ids(chunks)
+
+    # Get the existing documents from the database
+    existing_items = db.get(include=[])  # IDs are always included by default
+    existing_ids = set(existing_items["ids"])
+    print(f"Number of existing documents in DB: {len(existing_ids)}")
+
+    # Add only new chunks that don't exist in the database
+    new_chunks = []
+    for chunk in chunks_with_ids:
+        if chunk.metadata["id"] not in existing_ids:
+            new_chunks.append(chunk)
+
+    # Add new documents to the database and persist the changes
+    if len(new_chunks):
+        print(f"----->>> Adding new documents: {len(new_chunks)}")
+        new_chunk_ids = [chunk.metadata["id"] for chunk in new_chunks]
+        db.add_documents(new_chunks, ids=new_chunk_ids)
+        db.persist()
+    else:
+        print("----->>> No new documents to add")
+
+def calculate_chunk_ids(chunks):
+    # Create unique IDs for each chunk based on the source, page, and chunk index
+    last_page_id = None
+    current_chunk_index = 0
+
+    for chunk in chunks:
+        source = chunk.metadata.get("source")
+        page = chunk.metadata.get("page")
+        current_page_id = f"{source}:{page}"
+
+        # Increment the chunk index if it's the same page
+        if current_page_id == last_page_id:
+            current_chunk_index += 1
+        else:
+            current_chunk_index = 0
+
+        # Assign a unique ID to each chunk
+        chunk_id = f"{current_page_id}:{current_chunk_index}"
+        last_page_id = current_page_id
+        chunk.metadata["id"] = chunk_id
+
+    return chunks
+
+def clear_database():
+    # Remove the Chroma directory to clear the database
+    if os.path.exists(CHROMA_PATH):
+        shutil.rmtree(CHROMA_PATH)
+
+if __name__ == "__main__":
+    main()

embedding_function.py

@@ -0,0 +1,6 @@
+from langchain_community.embeddings.ollama import OllamaEmbeddings
+
+def get_embedding():
+    # Initialize embeddings with the specified model
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
+    return embeddings  # Return the embedding function

query.py

@@ -0,0 +1,61 @@
+import argparse
+from langchain.vectorstores.chroma import Chroma
+from langchain.prompts import ChatPromptTemplate
+from langchain_community.llms.ollama import Ollama
+
+from embedding_function import get_embedding
+
+# Path to the Chroma database
+CHROMA_PATH = "chroma"
+
+# Template for the prompt to be used by the LLM
+PROMPT_TEMPLATE = """
+Answer the question based only on the following context:
+
+{context}
+
+---
+
+Answer the question based on the above context: {question}
+"""
+
+def main():
+    # Create a command-line interface to accept the query text
+    parser = argparse.ArgumentParser()
+    parser.add_argument("query_text", type=str, help="The query text.")
+    args = parser.parse_args()
+    query_text = args.query_text
+    query_rag(query_text)
+
+def query_rag(query_text: str):
+    # Initialize the embedding function
+    embedding_function = get_embedding()
+    
+    # Load the Chroma vector store with the embedding function
+    db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)
+
+    # Perform a similarity search on the database with the query text
+    results = db.similarity_search_with_score(query_text, k=5)
+
+    # Combine the search results into a single context string
+    context_text = "\n\n---\n\n".join([doc.page_content for doc, _score in results])
+    
+    # Format the prompt using the context and the query
+    prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
+    prompt = prompt_template.format(context=context_text, question=query_text)
+    print(prompt)
+
+    # Initialize the model and generate a response
+    model = Ollama(model="mistral")
+    response_text = model.invoke(prompt)
+
+    # Extract source IDs from the search results
+    sources = [doc.metadata.get("id", None) for doc, _score in results]
+    
+    # Format and print the response with the sources
+    formatted_response = f"Response: {response_text}\nSources: {sources}"
+    print(formatted_response)
+    return response_text
+
+if __name__ == "__main__":
+    main()