Upload chat_pdf.py

chat_pdf.py

@@ -0,0 +1,283 @@
+import os
+import tempfile
+import streamlit as st
+from streamlit_chat import message
+from langchain_community.vectorstores import Chroma
+from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_ollama.chat_models import ChatOllama
+from langchain.retrievers.multi_query import MultiQueryRetriever
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors.flashrank_rerank import FlashrankRerank
+from flashrank import Ranker, RerankRequest
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_ollama import OllamaEmbeddings
+from langchain_community.embeddings import FastEmbedEmbeddings
+from langchain_community.document_loaders import PyMuPDFLoader
+from langchain.document_loaders.pdf import PyPDFDirectoryLoader
+from langchain_community.document_loaders import WebBaseLoader
+from langchain.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnablePassthrough
+from langchain.retrievers import BM25Retriever, EnsembleRetriever
+from torch import cuda
+from langchain_community.llms import LlamaCpp
+
+device = "cuda" if cuda.is_available() else "cpu"
+
+st.set_page_config(page_title="Chatbot", layout="wide")
+
+class ChatPDF:
+    def __init__(self):
+        self.vector_db = None
+        self.llm = ChatOllama(model="hf.co/bartowski/Llama-3.2-3B-Instruct-GGUF:IQ4_XS")
+        self.llm = LlamaCpp(
+            model_path="/home/chatbot/.cache/huggingface/hub/models--bartowski--Meta-Llama-3.1-8B-Instruct-GGUF/snapshots/bf5b95e96dac0462e2a09145ec66cae9a3f12067/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf",
+            temperature=0.75, #to maintain randomness of generated text
+            n_gpu_layers=-1, #all layers of the model are loaded on the GPU if available
+            max_tokens=2000, #Sets the maximum number of tokens the model will generate for each inference
+            n_ctx=2048,
+            top_p=1,
+            # verbose=True,  # Verbose is required to pass to the callback manager
+        )
+        self.llm2 = LlamaCpp(
+            model_path="/home/chatbot/.cache/huggingface/hub/models--bartowski--Llama-3.2-3B-Instruct-GGUF/snapshots/5ab33fa94d1d04e903623ae72c95d1696f09f9e8/Llama-3.2-3B-Instruct-IQ4_XS.gguf",
+            temperature=0.75,
+            n_gpu_layers=-1,
+            max_tokens=2000,
+            n_ctx=2048,
+            top_p=1,
+        )
+        # self.llm2 = LlamaCpp(
+        #     model_path="/home/kja/project/models--bartowski--Llama-3.2-3B-Instruct-GGUF/snapshots/5ab33fa94d1d04e903623ae72c95d1696f09f9e8/Llama-3.2-3B-Instruct-IQ4_XS.gguf",
+        #     temperature=0.75,
+        #     n_gpu_layers=-1,
+        #     max_tokens=2000,
+        #     n_ctx=2048,
+        #     top_p=1,
+        # )
+        self.chain = None
+        self.processed_files = []
+
+    def ingest(self, file_path=None, file_name=None, webpage_url=None):
+        self.vector_db = None
+        self.processed_files.clear()
+
+        if webpage_url:
+            loader = WebBaseLoader(webpage_url)
+            data = loader.load()
+            st.success(f"Data from {webpage_url} loaded successfully!")
+        else:
+            self.processed_files.append(file_name)
+            loader = PyMuPDFLoader(file_path=file_path)
+            data = loader.load()
+            st.success(f"{file_name} loaded successfully!")
+        
+        #breaks the document in chunks of 1000 characters.the last 200 characters of last chunk is repeated in current chunk to maintain context continuity
+        #Overlapping preserves continuity, ensuring the chatbot understands full context when retrieving information.
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        chunks = text_splitter.split_documents(data)
+
+        self.vector_db = Chroma.from_documents(
+            documents=chunks,
+            embedding=FastEmbedEmbeddings(),
+            collection_metadata={'hnsw': "cosine"},
+            persist_directory='chromadbtest'
+        )
+        st.success("Vector database created successfully!")
+
+        QUERY_PROMPT = PromptTemplate(
+            input_variables=["question"],
+            template="""You are an AI language model assistant. Your task is to generate three \
+            different versions of the given user question to retrieve relevant documents from a vector \
+            database. By generating multiple perspectives on the user question, your goal is to help\
+            the user overcome some of the limitations of the distance-based similarity search. \
+            Provide these alternative questions separated by newlines. Original question: {question}""",
+        )
+
+        retriever = MultiQueryRetriever.from_llm(
+            self.vector_db.as_retriever(), 
+            self.llm2,
+            #self.llm,
+            prompt=QUERY_PROMPT
+        ) 
+
+        keyword_retriever = BM25Retriever.from_documents(documents=chunks)
+        main_retriever = EnsembleRetriever(retrievers=[retriever, keyword_retriever], weights=[0.5, 0.5])
+        FlashrankRerank.model_rebuild()
+        compressor = FlashrankRerank()
+        compression_retriever = ContextualCompressionRetriever(
+            base_compressor=compressor, base_retriever=main_retriever
+        )
+
+        chat_prompt = ChatPromptTemplate.from_template(
+            """
+            You are an expert assistant designed to answer questions based on the provided information.
+            Use the context below to respond accurately and concisely to the query.
+            While giving response, don't explicitly mention the document name or metadata unless otherwise asked.
+            If the context does not contain the necessary information, state, 'The provided context does not contain enough information to answer the question'.
+
+            Context:
+            {context}
+
+            Answer the question based on the above context:
+            Question: {question}
+
+            """
+            )
+
+        self.chain = (
+            {"context": compression_retriever, "question": RunnablePassthrough() }
+            | chat_prompt
+            | self.llm2
+            #|self.llm
+            | StrOutputParser()
+        )
+
+    def ask(self, question):
+        if not self.chain:
+            return "Please upload your files first."
+        return self.chain.invoke(question)
+
+    def clear(self):
+        self.vector_db = None
+        self.chain = None
+        self.processed_files.clear()
+
+
+def display_messages():
+
+    for i, (msg, is_user) in enumerate(st.session_state["messages"]):
+        message(msg, is_user=is_user, key=str(i))
+    st.session_state["thinking_spinner"] = st.empty()
+
+
+def process_input():
+    if st.session_state["user_input"].strip():
+        user_text = st.session_state["user_input"].strip()
+        with st.session_state["thinking_spinner"], st.spinner("Thinking..."):
+            agent_text = st.session_state["assistant"].ask(user_text)
+
+        st.session_state["messages"].append((user_text, True))
+        st.session_state["messages"].append((agent_text, False))
+        st.session_state["user_input"] = ""
+
+
+def read_and_save_file():
+    st.session_state["assistant"].clear()
+    st.session_state["messages"] = []
+    st.session_state["user_input"] = ""
+
+    for file in st.session_state["file_uploader"]:
+        with tempfile.NamedTemporaryFile(delete=False) as tf:
+            tf.write(file.getbuffer())
+            file_path = tf.name
+
+        with st.session_state["ingestion_spinner"], st.spinner(f"Ingesting {file.name}..."):
+            st.session_state["assistant"].ingest(file_path=file_path, file_name=file.name)
+        os.remove(file_path)
+
+
+def ingest_webpage():
+    st.session_state["assistant"].clear()
+    st.session_state["messages"] = []
+    st.session_state["user_input"] = ""
+
+    webpage_url = st.session_state["webpage_url"]
+    with st.session_state["ingestion_spinner"], st.spinner(f"Ingesting data from {webpage_url}..."):
+        st.session_state["assistant"].ingest(webpage_url=webpage_url)
+
+
+def page():
+    if "messages" not in st.session_state:
+        st.session_state["messages"] = []
+        st.session_state["assistant"] = ChatPDF()
+
+    st.markdown(
+    """
+    <style>
+        .title-container {
+            display: flex;
+            justify-content: center;
+            align-items: center;
+            font-size: 40px;
+            font-weight: bold;
+            color: #1E90FF;  # Customize the title color
+        }
+    </style>
+    <div class="title-container">
+        Chatbot - AI Assistant🌟<br><br>
+        
+        
+    </div>
+    """,
+    unsafe_allow_html=True
+)
+    
+    st.markdown(
+    """
+    <style>
+        .welcome-text {
+            font-size: 20px; /* Slightly larger text for emphasis */
+            font-weight: bold;
+            color: #20C997; /* Vibrant teal for attention */
+            margin-bottom: 20px;
+        }
+        .sub-heading {
+            font-size: 20px; /* Smaller size for sub-heading */
+            font-weight: bold;
+            color: #FFD700; /* Golden yellow for contrast */
+            margin-top: 30px;
+            margin-bottom: 10px;
+        }
+        .instructions {
+            font-size: 16px; /* Smaller text for instructions */
+            color: #FFFFFF; /* White text for black background themes */
+            line-height: 1.6;
+        }
+    </style>
+
+    <div class="welcome-text">
+        Welcome to your AI-Powered Document Assistant!!  
+        Chat with your documents and ask questions effortlessly.
+    </div>
+
+    <div class="sub-heading">Get Started!!</div>
+    <div class="instructions">
+        1. <b>Upload Your Documents</b>: The system accepts multiple PDF files at once, analyzing the content to provide comprehensive insights. <br>
+        2. <b>Ask a Question</b>: After processing the documents, ask any question related to the content of your uploaded documents for a precise answer. <br><br><br>
+    </div>
+    """,
+    unsafe_allow_html=True
+)
+
+
+  
+    
+
+
+    with st.sidebar:
+        st.header("Options")
+        st.subheader("📄 Upload Documents")
+        st.file_uploader(
+            "Upload PDF documents",
+            type=["pdf"],
+            key="file_uploader",
+            on_change=read_and_save_file,
+            accept_multiple_files=True,
+        )
+        st.subheader("Source Status")
+        if st.session_state["assistant"].processed_files:
+            st.write(f"""Uploaded {len(st.session_state["assistant"].processed_files)} document(s)""")
+        else:
+            st.write("No documents uploaded yet.")
+
+
+    st.session_state["ingestion_spinner"] = st.empty()
+    display_messages()
+    st.text_input("Ask a question ✍", key="user_input", on_change=process_input, placeholder="Type your question here...")
+    
+
+
+
+if __name__ == "__main__":
+    page()