Update app.py

app.py

@@ -1,93 +1,100 @@
 import streamlit as st
-from PyPDF2 import PdfReader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-import os
-from langchain import LLMChain
-from langchain_google_genai import GoogleGenerativeAIEmbeddings
-from langchain.llms import HuggingFaceHub
-from langchain.embeddings import HuggingFaceInferenceAPIEmbeddings
-import google.generativeai as genai
-from langchain.vectorstores import FAISS
-from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain.chains.question_answering import load_qa_chain
-from langchain.prompts import PromptTemplate
-import time
+from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate
+from llama_index.llms.huggingface import HuggingFaceInferenceAPI
 from dotenv import load_dotenv
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.core import Settings
+import os
+import base64
 
-genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
-os.environ["HUGGINGFACEHUB_API_TOKEN"] = os.getenv("HF_TOKEN")
+# Load environment variables
+load_dotenv()
 
-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
+# Configure the Llama index settings
+Settings.llm = HuggingFaceInferenceAPI(
+    model_name="google/gemma-1.1-7b-it",
+    tokenizer_name="google/gemma-1.1-7b-it",
+    context_window=3900,
+    token=os.getenv("HF_TOKEN"),
+    max_new_tokens=1000,
+    generate_kwargs={"temperature": 0.1},
+)
+Settings.embed_model = HuggingFaceEmbedding(
+    model_name="BAAI/bge-small-en-v1.5"
+)
 
-def get_text_chunks(text):
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=700, chunk_overlap=200)
-    chunks = text_splitter.split_text(text)
-    return chunks
+# Define the directory for persistent storage and data
+PERSIST_DIR = "./db"
+DATA_DIR = "data"
 
-def get_vector_store(text_chunks):
-    HF_TOKEN = os.getenv("HF_TOKEN")
-    embeddings = HuggingFaceInferenceAPIEmbeddings(api_key=HF_TOKEN, model_name="BAAI/bge-base-en-v1.5")
-    vector_store = FAISS.from_texts(text_chunks, embedding=embeddings)
-    vector_store.save_local("faiss_index")
+# Ensure data directory exists
+os.makedirs(DATA_DIR, exist_ok=True)
+os.makedirs(PERSIST_DIR, exist_ok=True)
 
-def get_conversational_chain():
-    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 Provided PDF", don't provide the wrong answer\n\n
-    Context:\n {context}?\n
-    Question: \n{question}\n
+def displayPDF(file):
+    with open(file, "rb") as f:
+        base64_pdf = base64.b64encode(f.read()).decode('utf-8')
+    pdf_display = f'<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="600" type="application/pdf"></iframe>'
+    st.markdown(pdf_display, unsafe_allow_html=True)
 
-    Answer:
-    """
-    # model = ChatGoogleGenerativeAI(model="gemini-pro", temperature=0.1)
-    model = HuggingFaceHub(repo_id="google/gemma-1.1-7b-it",
-                            model_kwargs={"temperature": 0.2,"max_new_tokens":512, "return_only_answer":True})
-    prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
-    chain = LLMChain(llm=model, prompt=prompt)
-    return chain
+def data_ingestion():
+    documents = SimpleDirectoryReader(DATA_DIR).load_data()
+    storage_context = StorageContext.from_defaults()
+    index = VectorStoreIndex.from_documents(documents)
+    index.storage_context.persist(persist_dir=PERSIST_DIR)
 
-def user_input(user_question):
-    embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
-    new_db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
-    docs = new_db.similarity_search(user_question)
-    chain = get_conversational_chain()
-    response = chain({"context": docs, "question": user_question}, return_only_outputs=True)
-    return response
+def handle_query(query):
+    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
+    index = load_index_from_storage(storage_context)
+    chat_text_qa_msgs = [
+    (
+        "user",
+        """You are a Q&A assistant named CHATTO, created by Pachaiappan. You have a specific response programmed for when users specifically ask about your creator, Suriya. The response is: "I was created by Suriya, an enthusiast in Artificial Intelligence. He is dedicated to solving complex problems and delivering innovative solutions. With a strong focus on machine learning, deep learning, Python, generative AI, NLP, and computer vision, Suriya is passionate about pushing the boundaries of AI to explore new possibilities." For all other inquiries, your main goal is to provide answers as accurately as possible, based on the instructions and context you have been given. If a question does not match the provided context or is outside the scope of the document, kindly advise the user to ask questions within the context of the document.
+        Context:
+        {context_str}
+        Question:
+        {query_str}
+        """
+    )
+    ]
+    text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
+    query_engine = index.as_query_engine(text_qa_template=text_qa_template)
+    answer = query_engine.query(query)
+    
+    if hasattr(answer, 'response'):
+        return answer.response
+    elif isinstance(answer, dict) and 'response' in answer:
+        return answer['response']
+    else:
+        return "Sorry, I couldn't find an answer."
 
-# def response_generator(response):
-#     for word in response.split():
-#         yield word + " "
-#         time.sleep(0.05)
 
 # Streamlit app initialization
-st.title("Chat With PDF 📄")
+st.title("Chat with your PDF 🦜📄")
+st.markdown("Built by [Suriya❤️](https://github.com/theSuriya)")
+st.markdown("chat here👇")
 
 if 'messages' not in st.session_state:
     st.session_state.messages = [{'role': 'assistant', "content": 'Hello! Upload a PDF and ask me anything about its content.'}]
 
 with st.sidebar:
     st.title("Menu:")
-    uploaded_file = st.file_uploader("Upload your PDF Files and Click on the Submit & Process Button", accept_multiple_files=True)
+    uploaded_file = st.file_uploader("Upload your PDF Files and Click on the Submit & Process Button")
     if st.button("Submit & Process"):
         with st.spinner("Processing..."):
-            raw_text = get_pdf_text(uploaded_file)
-            text_chunks = get_text_chunks(raw_text)
-            get_vector_store(text_chunks)
+            filepath = "data/saved_pdf.pdf"
+            with open(filepath, "wb") as f:
+                f.write(uploaded_file.getbuffer())
+            # displayPDF(filepath)  # Display the uploaded PDF
+            data_ingestion()  # Process PDF every time new file is uploaded
             st.success("Done")
 
 user_prompt = st.chat_input("Ask me anything about the content of the PDF:")
 if user_prompt:
     st.session_state.messages.append({'role': 'user', "content": user_prompt})
-    response = user_input(user_prompt)
-    # answer = response_generator(response)
+    response = handle_query(user_prompt)
     st.session_state.messages.append({'role': 'assistant', "content": response})
 
 for message in st.session_state.messages:
     with st.chat_message(message['role']):
-        st.write(message['content'])
+        st.write(message['content'])