PDFGenius / app.py
SheemaMasood's picture
Update app.py
7ab73d7 verified
Raw
History Blame Contribute Delete
3.52 kB
import streamlit as st
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
from dotenv import load_dotenv
import os
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() if page.extract_text() else ""
return text
def get_text_chunks(text):
text_splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
chunks = text_splitter.split_text(text)
return chunks
def get_vector_store(text_chunks):
embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
vector_store = Chroma.from_texts(text_chunks, embedding=embeddings, persist_directory="chroma_db")
vector_store.persist()
def get_gemini_response(prompt):
chat_model = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0.3)
response = chat_model.predict(prompt)
return response
def user_input(user_question):
embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
new_db = Chroma(persist_directory="chroma_db", embedding_function=embeddings)
docs = new_db.similarity_search(user_question)
context = "\n".join([doc.page_content for doc in docs])
prompt = f"Context:\n{context}\n\nQuestion: {user_question}\nAnswer:"
response = get_gemini_response(prompt)
st.write("Reply: ", response)
def summarize_text(text, length="medium"):
summary_prompt = f"Summarize the following text in a {length} manner:\n\n" + text
return get_gemini_response(summary_prompt)
def main():
st.set_page_config("PDF Genius - Chat & Summarizer")
st.title("πŸ“š PDF Genius")
st.markdown("### Chat with your PDFs and generate summaries with Gemini! πŸ’‘")
with st.sidebar:
st.title("πŸ“Œ Menu:")
pdf_docs = st.file_uploader("Upload your PDF Files", accept_multiple_files=True)
if st.button("Submit & Process"):
with st.spinner("Processing..."):
raw_text = get_pdf_text(pdf_docs)
text_chunks = get_text_chunks(raw_text)
get_vector_store(text_chunks)
st.session_state["raw_text"] = raw_text
st.success("Processing Complete βœ…")
tab1, tab2 = st.tabs(["πŸ’¬ Chat with PDF", "πŸ“„ Document Summarizer"])
with tab1:
st.header("πŸ’¬ Chat with PDF")
user_question = st.text_input("Ask a Question from the PDF Files")
if user_question:
user_input(user_question)
with tab2:
st.header("πŸ“„ Document Summarizer")
if "raw_text" in st.session_state:
length = st.selectbox("Select Summary Length", ["short", "medium", "detailed"], index=1)
if st.button("Summarize Document"):
summary = summarize_text(st.session_state["raw_text"], length)
st.subheader("Summary:")
st.write(summary)
else:
st.warning("⚠ Please upload and process a PDF first from the sidebar.")
# Footer
st.markdown("---")
st.markdown("Developed by **Sheema Masood** | Powered by **Streamlit**")
st.markdown("πŸš€ Running at [Hugging Face Spaces](https://huggingface.co/spaces/SheemaMasood/PDFGenius)")
if __name__ == "__main__":
main()