import streamlit as st import logging import os from io import BytesIO import pdfplumber from langchain.text_splitter import CharacterTextSplitter from langchain_community.vectorstores import FAISS from sentence_transformers import SentenceTransformer from transformers import pipeline # Setup logging for Spaces logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # Lazy load models @st.cache_resource(ttl=3600) def load_embeddings_model(): logger.info("Loading embeddings model") try: return SentenceTransformer("all-MiniLM-L6-v2") except Exception as e: logger.error(f"Embeddings load error: {str(e)}") st.error(f"Embedding model error: {str(e)}") return None @st.cache_resource(ttl=3600) def load_qa_pipeline(): logger.info("Loading QA pipeline") try: return pipeline("text2text-generation", model="google/flan-t5-base", max_length=300) except Exception as e: logger.error(f"QA model load error: {str(e)}") st.error(f"QA model error: {str(e)}") return None # Process PDF def process_pdf(uploaded_file): logger.info("Processing PDF") try: text = "" with pdfplumber.open(BytesIO(uploaded_file.getvalue())) as pdf: for page in pdf.pages: extracted = page.extract_text() if extracted: text += extracted + "\n" if not text: # Optional OCR (uncomment if needed, requires pdf2image, pytesseract) # from pdf2image import convert_from_bytes # import pytesseract # images = convert_from_bytes(uploaded_file.getvalue()) # text = "".join(pytesseract.image_to_string(img) for img in images) if not text: raise ValueError("No text extracted from PDF") text_splitter = CharacterTextSplitter(separator="\n", chunk_size=600, chunk_overlap=150) chunks = text_splitter.split_text(text) embeddings_model = load_embeddings_model() if not embeddings_model: return None, text embeddings = [embeddings_model.encode(chunk) for chunk in chunks] vector_store = FAISS.from_embeddings(zip(chunks, embeddings), embeddings_model.encode) logger.info("PDF processed successfully") return vector_store, text except Exception as e: logger.error(f"PDF processing error: {str(e)}") st.error(f"PDF error: {str(e)}") return None, "" # Summarize PDF def summarize_pdf(text): logger.info("Generating summary") try: qa_pipeline = load_qa_pipeline() if not qa_pipeline: return "Summary model unavailable." # Split text for summarization if too long text_splitter = CharacterTextSplitter(separator="\n", chunk_size=1000, chunk_overlap=100) chunks = text_splitter.split_text(text) summaries = [] for chunk in chunks[:3]: # Limit to first 3 chunks for brevity prompt = f"Summarize this text in 60-80 words, highlighting key points:\n{chunk}" summary = qa_pipeline(prompt, max_length=100)[0]['generated_text'] summaries.append(summary.strip()) combined_summary = " ".join(summaries) if len(combined_summary.split()) > 200: combined_summary = " ".join(combined_summary.split()[:200]) logger.info("Summary generated") return combined_summary except Exception as e: logger.error(f"Summary error: {str(e)}") return f"Error summarizing: {str(e)}" # Answer question def answer_question(vector_store, query): logger.info(f"Processing query: {query}") try: if not vector_store: return "Please upload a PDF first." qa_pipeline = load_qa_pipeline() if not qa_pipeline: return "QA model unavailable." docs = vector_store.similarity_search(query, k=3) context = "\n".join(doc.page_content for doc in docs) prompt = f"Context: {context}\nQuestion: {query}\nAnswer concisely:" response = qa_pipeline(prompt)[0]['generated_text'] logger.info("Answer generated") return response.strip() except Exception as e: logger.error(f"Query error: {str(e)}") return f"Error answering: {str(e)}" # Streamlit UI try: st.set_page_config(page_title="Smart PDF Q&A", page_icon="📄") st.title("Smart PDF Q&A") st.markdown(""" Upload a PDF to ask questions or get a summary (up to 200 words). Chat history is preserved. """, unsafe_allow_html=True) # Initialize session state if "messages" not in st.session_state: st.session_state.messages = [] if "vector_store" not in st.session_state: st.session_state.vector_store = None if "pdf_text" not in st.session_state: st.session_state.pdf_text = "" # PDF upload uploaded_file = st.file_uploader("Upload a PDF", type=["pdf"]) if uploaded_file: col1, col2 = st.columns([1, 1]) with col1: if st.button("Process PDF"): with st.spinner("Processing PDF..."): st.session_state.vector_store, st.session_state.pdf_text = process_pdf(uploaded_file) if st.session_state.vector_store: st.success("PDF processed! Ask questions or summarize.") st.session_state.messages = [] else: st.error("Failed to process PDF.") with col2: if st.button("Summarize PDF") and st.session_state.pdf_text: with st.spinner("Generating summary..."): summary = summarize_pdf(st.session_state.pdf_text) st.session_state.messages.append({"role": "assistant", "content": f"**Summary**: {summary}"}) st.markdown(f"**Summary**: {summary}") # Chat interface if st.session_state.vector_store: prompt = st.chat_input("Ask a question about the PDF:") if prompt: st.session_state.messages.append({"role": "user", "content": prompt}) with st.chat_message("user"): st.markdown(prompt) with st.chat_message("assistant"): with st.spinner("Generating answer..."): answer = answer_question(st.session_state.vector_store, prompt) st.markdown(answer) st.session_state.messages.append({"role": "assistant", "content": answer}) # Display chat history for message in st.session_state.messages: with st.chat_message(message["role"]): st.markdown(message["content"]) # Download chat history if st.session_state.messages: chat_text = "\n".join(f"{m['role'].capitalize()}: {m['content']}" for m in st.session_state.messages) st.download_button("Download Chat History", chat_text, "chat_history.txt") except Exception as e: logger.error(f"App initialization failed: {str(e)}") st.error(f"App failed to start: {str(e)}. Check Spaces logs or contact support.")