Update app.py

app.py

@@ -1,14 +1,89 @@
 from PyPDF2 import PdfReader
 from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.text_splitter import CharacterTextSplitter
+from langchain.text_splitter import CharacterTextSplitter,RecursiveCharacterTextSplitter
 from langchain.vectorstores import FAISS
 from langchain.chains.question_answering import load_qa_chain
 from langchain.chains import load_chain
 from langchain.llms import OpenAI
 import streamlit as st
-# import pyautogui
+import openai
+from langchain.prompts import PromptTemplate
+from langchain_google_genai import GoogleGenerativeAIEmbeddings,ChatGoogleGenerativeAI
+import google.generativeai as genai
+
 import os, shutil
 
+
+def get_pdf_text(pdf_docs):
+    text=""
+    for pdf in pdf_docs:
+        pdf_reader= PdfReader(pdf)
+        for page in pdf_reader.pages:
+            text+= page.extract_text()
+    return  text
+
+def get_text_chunks(text,method):
+    if method=='Google-Gemini':
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=300)
+        chunks = text_splitter.split_text(text)
+    else:
+        text_splitter = CharacterTextSplitter(separator = "\n",chunk_size = 1000,chunk_overlap  = 300,length_function = len)
+        chunks = text_splitter.split_text(raw_text)
+    return chunks
+
+def get_vector_store(text_chunks,method):
+    try:
+        if method=='Google-Gemini':
+            embeddings = GoogleGenerativeAIEmbeddings(model = "models/embedding-001")
+        else:
+            embeddings = OpenAIEmbeddings()
+        vector_store = FAISS.from_texts(text_chunks, embedding=embeddings)
+        vector_store.save_local("faiss_index")
+    except:
+        st.warning("Wrong API, give a valid API")
+
+
+def get_conversational_chain(method):
+        
+    prompt_template = """
+    Answer the question as detailed as possible from the provided context, make sure to provide all the details, if the answer is not in
+    provided context just say, "answer is not available in the context", don't provide the wrong answer\n\n
+    Context:\n {context}?\n
+    Question: \n{question}\n
+
+    Answer:
+    """
+    if method=='Google-Gemini':
+        model = ChatGoogleGenerativeAI(model="gemini-pro",
+                            temperature=0.3)
+    else:
+        model= OpenAI()
+    prompt = PromptTemplate(template = prompt_template, input_variables = ["context", "question"])
+    chain = load_qa_chain(model, chain_type="stuff", prompt=prompt)
+    return chain
+
+
+
+def user_input(user_question,method):
+    if method=='Google-Gemini':
+        embeddings = GoogleGenerativeAIEmbeddings(model = "models/embedding-001")
+    else:
+        embeddings = OpenAIEmbeddings()
+    
+    new_db = FAISS.load_local("faiss_index", embeddings)
+    docs = new_db.similarity_search(user_question)
+
+    chain = get_conversational_chain(method)
+
+    
+    response = chain(
+        {"input_documents":docs, "question": user_question}
+        , return_only_outputs=True)
+    return response
+
+
+
+
 def delete_directory(directory_path):
     try:
         shutil.rmtree(directory_path)
@@ -16,60 +91,52 @@ def delete_directory(directory_path):
     except Exception as e:
         print(f"Error deleting directory '{directory_path}': {e}")
 
-st.set_page_config(page_title="Query any Pdf", page_icon="📄")
-
-st.title("📄 PDF Query Bot 📄")
-st.write("Made with ❤️ by Mainak")
-
 def return_response(query,document_search,chain):
     query = query
     docs = document_search.similarity_search(query)
     result = chain.run(input_documents=docs, question=query)
     return result
 
-uploaded_file = st.file_uploader("Upload a PDF File", type=["pdf"])
 
-# API key input box
-api_key = st.text_input("Enter Your OpenAI API Key",type="password")
 
-if not  uploaded_file:
+
+st.set_page_config(page_title="Query any Pdf", page_icon="📄")
+
+st.title("📄 PDF Query Bot 📄")
+st.write("Made with ❤️ by Mainak")
+with st.sidebar:
+    pdf_docs = st.file_uploader("Upload your PDF Files and Click on the Submit Button", accept_multiple_files=True,type=['pdf'])
+    option = st.selectbox('Select a Model(choose OpenAI for best results)',('OpenAI', 'Google-Gemini'))
+    if option=='OpenAI':
+        api_key = st.text_input("Enter Your OpenAI API Key",type="password")
+        os.environ["OPENAI_API_KEY"] = api_key
+    else:
+        api_key = st.text_input("Enter Your Google-Gemini API Key",type="password")
+        os.environ["google_API_KEY"] = api_key
+        genai.configure(api_key=os.getenv("google_API_KEY"))
+if not  pdf_docs:
     try:
         delete_directory('faiss_index')
     except:
         pass
-
-if st.button('Submit'):
-    if api_key:
-        if uploaded_file is not None:
-            # Read text from the uploaded file
-            os.environ["OPENAI_API_KEY"] = api_key
-            with st.spinner('Wait for it...'):
-                pdfreader = PdfReader(uploaded_file)
-                # read text from pdf
-                raw_text = ''
-                for i, page in enumerate(pdfreader.pages):
-                    content = page.extract_text()
-                    if content:
-                        raw_text += content
-
-                text_splitter = CharacterTextSplitter(
-                    separator = "\n",
-                    chunk_size = 800,
-                    chunk_overlap  = 200,
-                    length_function = len,
-                )
-                texts = text_splitter.split_text(raw_text)
-                embeddings = OpenAIEmbeddings()
-                document_search = FAISS.from_texts(texts, embeddings)
-                document_search.save_local("faiss_index")
+with st.sidebar:
+    if st.button('Submit'):
+        if api_key:
+            if pdf_docs is not None:
+                # Read text from the uploaded file
+                os.environ["OPENAI_API_KEY"] = api_key
+                with st.spinner('Wait for it...'):
+                    raw_text = get_pdf_text(pdf_docs)
+                    chunks = get_text_chunks(raw_text,option)                   
+                    get_vector_store(chunks,option)
+            else:
+                st.warning("Please enter your Pdf File")
         else:
-            st.warning("Please enter your Pdf File")
-    else:
-        st.warning("Please enter your API key")
+            st.warning("Please enter your API key")
+
 if os.path.exists("faiss_index"):
-    # if st.checkbox("chat"):
         if api_key:
-            if uploaded_file is not None:
+            if pdf_docs is not None:
                 if "messages" not in st.session_state:
                     st.session_state.messages = []
 
@@ -83,18 +150,14 @@ if os.path.exists("faiss_index"):
                         st.markdown(prompt)
                     # Add user message to chat history
                     st.session_state.messages.append({"role": "user", "content": prompt})
-                os.environ["OPENAI_API_KEY"] = api_key
-                embeddings = OpenAIEmbeddings()
-                document_search = FAISS.load_local("faiss_index", embeddings)
-                chain = load_qa_chain(OpenAI(), chain_type="stuff")
+                # os.environ["OPENAI_API_KEY"] = api_key
                 if prompt is None:
                     re='Ask me anything about the pdf'
-                # elif prompt=='exit':
-                #     delete_directory('faiss_index')
-                #     pyautogui.hotkey('f5') #Simulates F5 key press = page refresh
                 else:
                     with st.spinner('Typping...'):
-                        re=return_response(str(prompt),document_search,chain)
+                        re = user_input(str(prompt),option)
+                        re = re["output_text"]
+                        # re=return_response(str(prompt),document_search,chain)
                 response = f"PDF Mate: {re}"
                 # Display assistant response in chat message container
                 with st.chat_message("assistant"):
@@ -106,4 +169,5 @@ if os.path.exists("faiss_index"):
         else:
             st.warning("Please enter your API key")
 else:
-    pass
+    pass
+