Update app.py

app.py

@@ -1,81 +1,102 @@
 import streamlit as st
-import os
 import logging
+import os
 from io import BytesIO
 from PyPDF2 import PdfReader
 from langchain.text_splitter import CharacterTextSplitter
-from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
-from langchain.prompts import PromptTemplate
-from langchain.chains.question_answering import load_qa_chain
-from langchain_community.llms import HuggingFaceHub
-from transformers import pipeline  # For fallback if Hub fails
+from sentence_transformers import SentenceTransformer
+from transformers import pipeline
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
+# Setup logging for Spaces
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Check API token
-if "HUGGINGFACEHUB_API_TOKEN" not in os.environ:
-    st.error("HUGGINGFACEHUB_API_TOKEN not set in secrets. Add it in Space settings.")
-    st.stop()
+# 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
 
-try:
-    # Function to process PDF
-    def process_pdf(uploaded_file):
-        try:
-            logger.info("Starting PDF processing")
-            pdf_reader = PdfReader(BytesIO(uploaded_file.getvalue()))
-            text = ""
-            for page in pdf_reader.pages:
-                extracted = page.extract_text()
-                if extracted:
-                    text += extracted + "\n"
-            
+@st.cache_resource(ttl=3600)
+def load_qa_pipeline():
+    logger.info("Loading QA pipeline")
+    try:
+        return pipeline("text2text-generation", model="google/flan-t5-small", max_length=200)
+    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:
+        pdf_reader = PdfReader(BytesIO(uploaded_file.getvalue()))
+        text = "".join(page.extract_text() or "" for page in pdf_reader.pages)
+        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.")
-            
-            # Chunk text (increased overlap for better context)
-            text_splitter = CharacterTextSplitter(separator="\n", chunk_size=800, chunk_overlap=200, length_function=len)
-            chunks = text_splitter.split_text(text)
-            
-            # Embeddings (light model)
-            embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={'device': 'cpu'})
-            
-            # Vector store
-            vector_store = FAISS.from_texts(chunks, embedding=embeddings)
-            logger.info("PDF processed successfully")
-            return vector_store
-        except Exception as e:
-            logger.error(f"PDF processing error: {str(e)}")
-            st.error(f"Error processing PDF: {str(e)}")
+                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
+        
+        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
+    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:
-            logger.info(f"Answering query: {query}")
-            # Lighter LLM via pipeline for faster CPU inference
-            qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-base")
-            
-            # Retrieve top chunks
-            docs = vector_store.similarity_search(query, k=3)
-            context = "\n".join([doc.page_content for doc in docs])
-            
-            # Prompt
-            prompt = f"Use this context to answer concisely: {context}\nQuestion: {query}\nAnswer:"
-            response = qa_pipeline(prompt, max_length=256, num_return_sequences=1)[0]['generated_text']
-            
-            logger.info("Answer generated")
-            return response.strip()
-        except Exception as e:
-            logger.error(f"Answer generation error: {str(e)}")
-            st.error(f"Error answering: {str(e)}")
-            return "Unable to generate answer."
+# 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 with chat history
+# Streamlit UI
+try:
+    st.set_page_config(page_title="Smart PDF Q&A", page_icon="📄")
     st.title("Smart PDF Q&A")
-    st.write("Upload a PDF and ask questions! Chat history is preserved.")
+    st.markdown("""
+        Upload a PDF and ask questions about its content. 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:
@@ -83,37 +104,40 @@ try:
     if "vector_store" not in st.session_state:
         st.session_state.vector_store = None
 
-    # PDF upload and process
-    uploaded_file = st.file_uploader("Upload PDF", type="pdf")
-    if uploaded_file:
-        if st.button("Process PDF"):
-            with st.spinner("Processing..."):
-                vector_store = process_pdf(uploaded_file)
-                if vector_store:
-                    st.session_state.vector_store = vector_store
-                    st.success("PDF ready! Ask away.")
-                    st.session_state.messages = []  # Reset chat on new PDF
-
-    # Display chat history
-    for message in st.session_state.messages:
-        with st.chat_message(message["role"]):
-            st.markdown(message["content"])
+    # PDF upload
+    uploaded_file = st.file_uploader("Upload a PDF", type=["pdf"])
+    if uploaded_file and st.button("Process PDF"):
+        with st.spinner("Processing PDF..."):
+            st.session_state.vector_store = process_pdf(uploaded_file)
+            if st.session_state.vector_store:
+                st.success("PDF processed! Ask questions below.")
+                st.session_state.messages = []
+            else:
+                st.error("Failed to process PDF.")
 
-    # Question input
+    # Chat interface
     if st.session_state.vector_store:
-        if prompt := st.chat_input("Ask a question:"):
-            # Add user message
+        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)
-            
-            # Generate answer
             with st.chat_message("assistant"):
-                with st.spinner("Thinking..."):
+                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"Initialization error: {str(e)}. Check logs or try factory reset.")
+    st.error(f"App failed to start: {str(e)}. Check Spaces logs or contact support.")