Spaces:
Sleeping
Sleeping
| import os | |
| import streamlit as st # type: ignore | |
| from dotenv import load_dotenv | |
| from PyPDF2 import PdfReader | |
| from langchain_classic.text_splitter import RecursiveCharacterTextSplitter | |
| from langchain_cohere import CohereEmbeddings, ChatCohere | |
| from langchain_community.vectorstores import FAISS | |
| from langchain_classic.chains import ConversationalRetrievalChain | |
| from langchain_classic.memory import ConversationBufferMemory | |
| load_dotenv() | |
| st.set_page_config( | |
| page_title="Chat with Documents", | |
| page_icon="📚" | |
| ) | |
| 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() or "" | |
| return text | |
| def get_text_chunks(text): | |
| text_splitter = RecursiveCharacterTextSplitter( | |
| separators=["\n\n", "\n", " ", ""], | |
| chunk_size=1000, | |
| chunk_overlap=200, | |
| length_function=len | |
| ) | |
| return text_splitter.split_text(text) | |
| def get_embeddings(): | |
| return CohereEmbeddings( | |
| model='embed-english-v3.0', | |
| user_agent="langchain" | |
| ) | |
| def get_conversational_chain(vectorstore): | |
| llm = ChatCohere( | |
| model="command-r-08-2024", | |
| temperature=0 | |
| ) | |
| memory = ConversationBufferMemory( | |
| memory_key="chat_history", | |
| return_messages=True, | |
| output_key="answer" | |
| ) | |
| conversational_chain = ConversationalRetrievalChain.from_llm( | |
| llm=llm, | |
| retriever=vectorstore.as_retriever(), | |
| memory=memory, | |
| return_source_documents=True | |
| ) | |
| return conversational_chain | |
| def handle_userinput(user_question): | |
| response = st.session_state.conversational_chain.invoke( | |
| {"question": user_question} | |
| ) | |
| answer = response["answer"] | |
| st.session_state.chat_history = ( | |
| st.session_state.conversational_chain.memory | |
| .load_memory_variables({})["chat_history"] | |
| ) | |
| for i, message in enumerate(st.session_state.chat_history): | |
| if message.type == "human": | |
| with st.chat_message("user"): | |
| st.write(message.content) | |
| else: | |
| with st.chat_message("assistant"): | |
| st.write(message.content) | |
| def main(): | |
| st.markdown( | |
| """ | |
| <style> | |
| .main { | |
| background: linear-gradient(135deg, #f8fbff 0%, #eef5ff 100%); | |
| } | |
| .block-container { | |
| padding-top: 2rem; | |
| padding-bottom: 2rem; | |
| } | |
| .hero { | |
| background: linear-gradient(135deg, #0f172a 0%, #1d4ed8 100%); | |
| padding: 2rem 2rem; | |
| border-radius: 22px; | |
| box-shadow: 0 12px 30px rgba(29, 78, 216, 0.18); | |
| margin-bottom: 1.5rem; | |
| margin-top: 1.5rem; | |
| } | |
| .hero h1 { | |
| color: white !important; | |
| font-size: 2.5rem !important; | |
| font-weight: 800 !important; | |
| margin: 0 !important; | |
| } | |
| .hero p { | |
| color: #dbeafe !important; | |
| font-size: 1.05rem !important; | |
| margin-top: 0.5rem !important; | |
| margin-bottom: 0 !important; | |
| } | |
| </style> | |
| """, | |
| unsafe_allow_html=True, | |
| ) | |
| st.markdown( | |
| """ | |
| <div class="hero"> | |
| <h1> AI Chatbot for Your Documents</h1> | |
| <p>Upload a PDF and ask questions in natural language.</p> | |
| </div> | |
| """, | |
| unsafe_allow_html=True, | |
| ) | |
| if "conversational_chain" not in st.session_state: | |
| st.session_state.conversational_chain = None | |
| if "chat_history" not in st.session_state: | |
| st.session_state.chat_history = [] | |
| with st.sidebar: | |
| st.header("Upload Documents") | |
| pdf_docs = st.file_uploader( | |
| "Upload your PDF files", | |
| type=["pdf"], | |
| accept_multiple_files=True | |
| ) | |
| if st.button("Process Documents"): | |
| if not pdf_docs: | |
| st.warning("Please upload at least one PDF.") | |
| elif not os.getenv("COHERE_API_KEY"): | |
| st.error("Cohere API key is missing.") | |
| else: | |
| try: | |
| with st.spinner("Processing documents..."): | |
| raw_text = get_pdf_text(pdf_docs) | |
| if not raw_text.strip(): | |
| st.error("No extractable text found in the PDFs.") | |
| return | |
| chunks = get_text_chunks(raw_text) | |
| embeddings = get_embeddings() | |
| vectorstore = FAISS.from_texts( | |
| texts=chunks, | |
| embedding=embeddings | |
| ) | |
| st.session_state.conversational_chain = ( | |
| get_conversational_chain(vectorstore) | |
| ) | |
| st.session_state.chat_history = [] | |
| st.success("Documents processed successfully!") | |
| except Exception as e: | |
| st.error(f"Error processing documents: {e}") | |
| user_question = st.chat_input("Ask a question about your PDFs") | |
| if user_question: | |
| if st.session_state.conversational_chain is None: | |
| st.warning("Please upload and process your documents first.") | |
| else: | |
| with st.chat_message("user"): | |
| st.write(user_question) | |
| with st.chat_message("assistant"): | |
| with st.spinner("Thinking..."): | |
| try: | |
| response = ( | |
| st.session_state.conversational_chain.invoke( | |
| {"question": user_question} | |
| ) | |
| ) | |
| st.write(response["answer"]) | |
| st.session_state.chat_history = ( | |
| st.session_state.conversational_chain | |
| .memory.load_memory_variables( | |
| {} | |
| )["chat_history"] | |
| ) | |
| except Exception as e: | |
| st.error(f"Error generating answer: {e}") | |
| if __name__ == "__main__": | |
| main() |