Update app.py

app.py

@@ -1,105 +1,186 @@
 import streamlit as st
+import logging
 import os
 from io import BytesIO
-import numpy as np
-from PyPDF2 import PdfReader
-from langchain.chains import RetrievalQA
+import pdfplumber
 from langchain.text_splitter import CharacterTextSplitter
-from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
-from langchain.llms import HuggingFaceHub
-from langchain.prompts import PromptTemplate
-from langchain.chains.question_answering import load_qa_chain
+from sentence_transformers import SentenceTransformer
+from transformers import pipeline
 
-# Check for Hugging Face API token
-if "HUGGINGFACEHUB_API_TOKEN" not in os.environ:
-    st.error("HUGGINGFACEHUB_API_TOKEN not set. Please add it in Space settings under secrets.")
-    st.stop()
+# Setup logging for Spaces
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
-try:
-    # Function to process PDF and build vector store
-    def process_input(uploaded_file):
-        try:
-            pdf_reader = PdfReader(uploaded_file)
-            text = ""
-            for page in pdf_reader.pages:
+# 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
-                else:
-                    st.warning("Some pages could not be 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.")
-            
-            # Split text into chunks
-            text_splitter = CharacterTextSplitter(separator="\n", chunk_size=1000, chunk_overlap=100, length_function=len)
-            chunks = text_splitter.split_text(text)
-            
-            # Create embeddings
-            embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2", model_kwargs={'device': 'cpu'})
-            
-            # Create vector store
-            vector_store = FAISS.from_texts(chunks, embedding=embeddings)
-            return vector_store
-        except Exception as e:
-            st.error(f"Error processing PDF: {str(e)}")
-            return None
+                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, ""
 
-    # Function to answer questions
-    def answer_question(vector_store, query):
-        try:
-            llm = HuggingFaceHub(repo_id="mistralai/Mistral-7B-Instruct-v0.1", model_kwargs={"temperature": 0.7, "max_length": 512})
-            
-            # Retrieve relevant documents
-            retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 3})
-            
-            # Prompt template
-            prompt_template = """Use the provided context to answer the question concisely:
-            
-            Context: {context}
-            
-            Question: {question}
-            
-            Answer:"""
-            PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
-            
-            # QA chain
-            chain = load_qa_chain(llm, chain_type="stuff", prompt=PROMPT)
-            docs = retriever.get_relevant_documents(query)
-            response = chain({"input_documents": docs, "question": query}, return_only_outputs=True)
-            return response["output_text"].split("Answer:")[-1].strip()
-        except Exception as e:
-            st.error(f"Error answering question: {str(e)}")
-            return "Unable to generate answer."
+# 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)}"
 
-    # Streamlit UI
-    def main():
-        st.title("Smart PDF Q&A")
+# 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."
         
-        # PDF upload
-        uploaded_file = st.file_uploader("Upload a PDF", type="pdf")
+        qa_pipeline = load_qa_pipeline()
+        if not qa_pipeline:
+            return "QA model unavailable."
         
-        if uploaded_file:
+        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.
+        <style>
+        .stChatMessage { border-radius: 10px; padding: 10px; margin: 5px; }
+        .stChatMessage.user { background-color: #e6f3ff; }
+        .stChatMessage.assistant { background-color: #f0f0f0; }
+        </style>
+    """, 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..."):
-                    vector_store = process_input(BytesIO(uploaded_file.getvalue()))
-                    if vector_store:
-                        st.session_state.vector_store = vector_store
-                        st.success("PDF processed! Ask questions below.")
+                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.")
-        
-        # Question input
-        if "vector_store" in st.session_state:
-            query = st.text_input("Ask a question about the PDF:")
-            if st.button("Submit") and query:
+        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, query)
-                    st.write("**Answer:**", 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"])
 
-    if __name__ == "__main__":
-        main()
+    # 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:
-    st.error(f"Application failed to initialize: {str(e)}")
+    logger.error(f"App initialization failed: {str(e)}")
+    st.error(f"App failed to start: {str(e)}. Check Spaces logs or contact support.")