Spaces:
Sleeping
Sleeping
File size: 4,937 Bytes
92706b9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
import streamlit as st
import google.generativeai as genai
from langchain.document_loaders import PyPDFDirectoryLoader
import os
import shutil
# Configuration
GOOGLE_API_KEY = "your_api_key_here" # Replace with your actual API key
# Page configuration
st.set_page_config(page_title="Chat with PDFs", page_icon="📚")
def initialize_session_state():
"""Initialize session state variables"""
session_state_vars = {
"messages": [],
"loaded_files": False,
"pdf_content": None,
"chat": None
}
for var, value in session_state_vars.items():
if var not in st.session_state:
st.session_state[var] = value
def load_pdfs(pdf_folder):
"""Load PDFs and return their content"""
if not os.path.exists(pdf_folder):
os.makedirs(pdf_folder)
loader = PyPDFDirectoryLoader(pdf_folder)
documents = loader.load()
# Concatenate all documents content
content = "\n\n".join([doc.page_content for doc in documents])
return content
def initialize_chat(pdf_content):
"""Initialize Gemini chat with PDF content"""
genai.configure(api_key="AIzaSyDXxoKVLaal5itTj981_ByYLtxok09MIRM")
generation_config = {
"temperature": 0.7,
"top_p": 0.95,
"top_k": 40,
"max_output_tokens": 8192,
}
model = genai.GenerativeModel(
model_name="gemini-1.5-pro",
generation_config=generation_config,
)
# Start chat with context
context_prompt = f"""You are a helpful assistant that answers questions based on the following document content:
{pdf_content}
Please use this content to answer user questions. If the answer cannot be found in the content, say so."""
chat = model.start_chat(history=[])
# Send initial context
chat.send_message(context_prompt)
return chat
def main():
initialize_session_state()
st.title("💬 Chat with PDFs")
# Sidebar for PDF upload
with st.sidebar:
st.header("Upload Documents")
uploaded_files = st.file_uploader(
"Upload your PDFs",
type=["pdf"],
accept_multiple_files=True
)
if uploaded_files and not st.session_state.loaded_files:
# Create pdfs directory if it doesn't exist
if not os.path.exists("pdfs"):
os.makedirs("pdfs")
# Clean up old PDF files
for file in os.listdir("pdfs"):
os.remove(os.path.join("pdfs", file))
# Save uploaded files
for file in uploaded_files:
with open(f"pdfs/{file.name}", "wb") as f:
f.write(file.getvalue())
# Load PDF content
with st.spinner("Processing PDFs..."):
try:
pdf_content = load_pdfs("pdfs")
st.session_state.pdf_content = pdf_content
st.session_state.loaded_files = True
# Initialize chat with content
st.session_state.chat = initialize_chat(pdf_content)
except Exception as e:
st.error(f"Error processing PDFs: {str(e)}")
return
# Main chat interface
if st.session_state.loaded_files:
# Display chat messages
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.markdown(message["content"])
# Chat input
if prompt := st.chat_input("Ask a question about your PDFs:"):
# Add user message to chat history
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.markdown(prompt)
with st.chat_message("assistant"):
response_placeholder = st.empty()
try:
# Get response from Gemini
if not st.session_state.chat:
st.session_state.chat = initialize_chat(st.session_state.pdf_content)
response = st.session_state.chat.send_message(prompt)
response_text = response.text
response_placeholder.markdown(response_text)
# Add assistant response to chat history
st.session_state.messages.append({"role": "assistant", "content": response_text})
except Exception as e:
response_placeholder.error(f"Error generating response: {str(e)}")
else:
st.info("Please upload PDFs to start chatting.")
if __name__ == "__main__":
main()
|