Update app.py

app.py

@@ -17,103 +17,91 @@ from langchain.chains import ConversationalRetrievalChain
 from htmlTemplates import css, bot_template, user_template
 from langchain.llms import HuggingFaceHub
 
+import os
+import streamlit as st
+from dotenv import load_dotenv
+from PyPDF2 import PdfReader
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.embeddings import HuggingFaceBgeEmbeddings
+from langchain.vectorstores import FAISS
+from langchain.chat_models import ChatOpenAI
+from langchain.memory import ConversationBufferMemory
+from langchain.chains import ConversationalRetrievalChain
+from htmlTemplates import css, bot_template, user_template
+from langchain.llms import HuggingFaceHub
+from langchain.chains import RetrievalQA
+
 def get_pdf_text(pdf_docs):
     text = ""
     for pdf in pdf_docs:
-        try:
-            pdf_reader = PdfReader(pdf)
-            for page in pdf_reader.pages:
-                text += page.extract_text()
-        except Exception as e:
-            st.error(f"Error extracting text from PDF: {e}")
+        pdf_reader = PdfReader(pdf)
+        for page in pdf_reader.pages:
+            text += page.extract_text()
     return text
 
 def get_text_chunks(text):
     text_splitter = CharacterTextSplitter(
-        separator="\n", chunk_size=1500, chunk_overlap=300, length_function=len
+        separator="\n",
+        chunk_size=1000,
+        chunk_overlap=200,
+        length_function=len
     )
-    try:
-        chunks = text_splitter.split_text(text)
-    except Exception as e:
-        st.error(f"Error splitting text into chunks: {e}")
-        chunks = []
+    chunks = text_splitter.split_text(text)
     return chunks
 
 def get_vectorstore(text_chunks):
     model = "BAAI/bge-base-en-v1.5"
-    encode_kwargs = {
-        "normalize_embeddings": True
-    }
-    try:
-        embeddings = HuggingFaceBgeEmbeddings(
-            model_name=model, encode_kwargs=encode_kwargs, model_kwargs={"device": "cpu"}
-        )
-        vectorstore = FAISS.from_texts(texts=text_chunks, embedding=embeddings)
-    except Exception as e:
-        st.error(f"Error creating vector store: {e}")
-        vectorstore = None
+    encode_kwargs = {"normalize_embeddings": True}
+    embeddings = HuggingFaceBgeEmbeddings(
+        model_name=model,
+        encode_kwargs=encode_kwargs,
+        model_kwargs={"device": "cpu"}
+    )
+    vectorstore = FAISS.from_texts(texts=text_chunks, embedding=embeddings)
     return vectorstore
 
 def get_conversation_chain(vectorstore):
-    if vectorstore is None:
-        return None
+    llm = HuggingFaceHub(
+        repo_id="mistralai/Mistral-7B-v0.3",
+        model_kwargs={"temperature": 0.5, "max_length": 4000},
+    )
 
-    try:
-        llm = HuggingFaceHub(
-            repo_id="mistralai/Mistral-7B-v0.3",
-            model_kwargs={"temperature": 0.5, "max_length": 4000},
-        )
-        memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
-        conversation_chain = ConversationalRetrievalChain.from_llm(
-            llm=llm, retriever=vectorstore.as_retriever(), memory=memory
-        )
-    except Exception as e:
-        st.error(f"Error creating conversation chain: {e}")
-        conversation_chain = None
+    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+    conversation_chain = ConversationalRetrievalChain.from_llm(
+        llm=llm,
+        retriever=vectorstore.as_retriever(),
+        memory=memory,
+        return_source_documents=True  # Add this line to return source documents
+    )
     return conversation_chain
 
 def handle_userinput(user_question):
-    if st.session_state.conversation is None:
-        st.error("Please process the PDF files before asking a question.")
-        return
-
-    try:
-        response = st.session_state.conversation({"question": user_question})
-        st.session_state.chat_history = response["chat_history"]
-
-        for i, message in enumerate(st.session_state.chat_history):
-            if i % 2 == 0:
-                st.write("//_^ User: " + message.content)
-            else:
-                st.write("🤖 ChatBot: " + message.content)
-    except Exception as e:
-        st.error(f"Error handling user input: {e}")
+    response = st.session_state.conversation({"question": user_question})
+    st.session_state.chat_history = response["chat_history"]
+    
+    for i, message in enumerate(st.session_state.chat_history):
+        if i % 2 == 0:
+            st.write(user_template.replace("{{MSG}}", message.content), unsafe_allow_html=True)
+        else:
+            st.write(bot_template.replace("{{MSG}}", message.content), unsafe_allow_html=True)
+    
+    # Display references
+    if "source_documents" in response:
+        st.write("References:")
+        for doc in response["source_documents"]:
+            st.write(f"- {doc.metadata.get('source', 'Unknown source')}, page {doc.metadata.get('page', 'Unknown page')}")
 
 def main():
-    st.set_page_config(
-        page_title="Chat with a Bot that tries to answer questions about multiple PDFs",
-        page_icon=":books:",
-    )
-
-    st.markdown("# Chat with a Bot")
-    st.markdown("This bot tries to answer questions about multiple PDFs. Let the processing of the PDF finish before adding your question. 🙏🏾")
-
+    load_dotenv()
+    st.set_page_config(page_title="Chat with Multiple PDFs", page_icon=":books:")
     st.write(css, unsafe_allow_html=True)
 
-    huggingface_token = st.text_input("Enter your HuggingFace Hub token", type="password")
-    #openai_api_key = st.text_input("Enter your OpenAI API key", type="password")
-
-    if huggingface_token:
-        os.environ["HUGGINGFACEHUB_API_TOKEN"] = huggingface_token
-    #if openai_api_key:
-    #    os.environ["OPENAI_API_KEY"] = openai_api_key
-
     if "conversation" not in st.session_state:
         st.session_state.conversation = None
     if "chat_history" not in st.session_state:
         st.session_state.chat_history = None
 
-    st.header("Chat with a Bot 🤖🦾 that tries to answer questions about multiple PDFs :books:")
+    st.header("Chat with Multiple PDFs :books:")
     user_question = st.text_input("Ask a question about your documents:")
     if user_question:
         handle_userinput(user_question)
@@ -125,20 +113,10 @@ def main():
         )
         if st.button("Process"):
             with st.spinner("Processing"):
-                try:
-                    # get pdf text
-                    raw_text = get_pdf_text(pdf_docs)
-
-                    # get the text chunks
-                    text_chunks = get_text_chunks(raw_text)
-
-                    # create vector store
-                    vectorstore = get_vectorstore(text_chunks)
-
-                    # create conversation chain
-                    st.session_state.conversation = get_conversation_chain(vectorstore)
-                except Exception as e:
-                    st.error(f"Error processing PDF files: {e}")
+                raw_text = get_pdf_text(pdf_docs)
+                text_chunks = get_text_chunks(raw_text)
+                vectorstore = get_vectorstore(text_chunks)
+                st.session_state.conversation = get_conversation_chain(vectorstore)
 
 if __name__ == "__main__":
     main()