Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import google.generativeai as genai | |
| from langchain.document_loaders import PyPDFDirectoryLoader | |
| import os | |
| import shutil | |
| # Configuration | |
| GOOGLE_API_KEY = st.secrets["GOOGLE_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=GOOGLE_API_KEY) | |
| 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() | |