Update app.py

app.py

@@ -0,0 +1,238 @@
+import streamlit as st
+import os
+import tempfile
+import torch
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain_huggingface import HuggingFacePipeline
+from langchain_classic.prompts import PromptTemplate
+from langchain_classic.chains import RetrievalQA
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from huggingface_hub import login
+
+
+# --- Page Config & Styling ---
+st.set_page_config(
+    page_title="DocTalk - Chat With PDF",
+    page_icon="📗💬",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
+
+# Custom CSS for polished UI and Footer
+st.markdown("""
+<style>
+    /* Chat styling */
+    .stChatInput {
+        padding-bottom: 1rem;
+    }
+    
+    /* Custom Footer */
+    .footer {
+        position: fixed;
+        left: 0;
+        bottom: 0;
+        width: 100%;
+        background-color: white;
+        color: #555;
+        text-align: center;
+        padding: 10px;
+        font-size: 14px;
+        border-top: 1px solid #eee;
+        z-index: 100;
+    }
+    
+    /* Hide Streamlit branding for cleaner look */
+    #MainMenu {visibility: hidden;}
+    footer {visibility: hidden;}
+    
+    /* Adjust sidebar padding for footer */
+    [data-testid="stSidebar"] {
+        padding-bottom: 50px;
+    }
+</style>
+""", unsafe_allow_html=True)
+
+# --- Session State Management ---
+if 'qa_chain' not in st.session_state: st.session_state.qa_chain = None
+if 'messages' not in st.session_state: st.session_state.messages = []
+if 'processing_done' not in st.session_state: st.session_state.processing_done = False
+
+# --- Authentication (Secrets Only) ---
+hf_token = os.environ.get("HF_TOKEN")
+
+# --- Model Loading (Cached & CPU Optimized) ---
+
+@st.cache_resource
+def load_embedding_model():
+    """Load the embedding model once to save time."""
+    try:
+        # Using a lightweight, fast embedding model
+        embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+        return embeddings
+    except Exception as e:
+        st.error(f"Error loading embedding model: {e}")
+        return None
+
+@st.cache_resource
+def load_llm_model(token):
+    """Load the Gemma LLM once."""
+    try:
+        login(token=token)
+        model_id = "google/gemma-2-2b-it" 
+        
+        tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
+        
+        # Load model to CPU (float32 is safe for CPU stability)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            device_map="cpu",
+            torch_dtype=torch.float32, 
+            token=token
+        )
+
+        pipe = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=512,
+            temperature=0.1,
+            repetition_penalty=1.1,
+            return_full_text=False
+        )
+        return pipe
+    except Exception as e:
+        return None
+
+# --- PDF Processing ---
+def process_document(uploaded_file, model_pipeline, embedding_model):
+    try:
+        # Save temp file
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp:
+            tmp.write(uploaded_file.getvalue())
+            tmp_path = tmp.name
+        
+        # Load & Split
+        loader = PyPDFLoader(tmp_path)
+        docs = loader.load()
+        splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+        chunks = splitter.split_documents(docs)
+        
+        # Vector Store (FAISS is faster for in-memory)
+        vector_store = FAISS.from_documents(chunks, embedding_model)
+        
+        # Chain Setup
+        llm = HuggingFacePipeline(pipeline=model_pipeline)
+        
+        template = """<start_of_turn>user
+Answer the question based strictly on the context below. Keep answers concise.
+Context: {context}
+Question: {question}<end_of_turn>
+<start_of_turn>model
+"""
+        prompt = PromptTemplate(template=template, input_variables=["context", "question"])
+        
+        qa_chain = RetrievalQA.from_chain_type(
+            llm=llm,
+            retriever=vector_store.as_retriever(search_kwargs={"k": 3}),
+            chain_type_kwargs={"prompt": prompt},
+            return_source_documents=True
+        )
+        return qa_chain
+    except Exception as e:
+        st.error(f"Error processing PDF: {e}")
+        return None
+
+# --- Main Layout ---
+
+# 1. Sidebar Configuration
+with st.sidebar:
+    st.title("🤖 Configuration")
+    st.markdown("---")
+    
+    if not hf_token:
+        st.error("🚨 **HF_TOKEN missing!**")
+        st.info("Go to Space Settings -> Repository Secrets and add your Hugging Face Access Token as `HF_TOKEN`.")
+        st.stop()
+    else:
+        st.success("✅ Huggingface Active")
+    
+    st.subheader("📄 Document Upload")
+    uploaded_file = st.file_uploader("Upload your PDF", type="pdf", help="Max file size ~200MB")
+    
+    if uploaded_file:
+        process_btn = st.button("🚀 Process Document", type="primary", use_container_width=True)
+        
+        if process_btn:
+            with st.spinner("🧠 Analyzing PDF"):
+                # Load models (cached)
+                llm_pipeline = load_llm_model(hf_token)
+                embed_model = load_embedding_model()
+                
+                if llm_pipeline and embed_model:
+                    qa_chain = process_document(uploaded_file, llm_pipeline, embed_model)
+                    if qa_chain:
+                        st.session_state.qa_chain = qa_chain
+                        st.session_state.processing_done = True
+                        st.success("Done! You can now chat.")
+                    else:
+                        st.error("Failed to process document.")
+                else:
+                    st.error("Failed to load AI models. Check token permissions.")
+    
+    if st.session_state.processing_done:
+        st.markdown("---")
+        if st.button("🗑️ Clear Chat History", use_container_width=True):
+            st.session_state.messages = []
+            st.rerun()
+
+# 2. Main Chat Area
+st.title("📗💬 DocTalk - Chat With PDF")
+#st.caption("Powered by Google Gemma-2-2B-IT")
+
+if st.session_state.processing_done:
+    # Display History
+    for msg in st.session_state.messages:
+        with st.chat_message(msg["role"]):
+            st.markdown(msg["content"])
+            
+    # Chat Input
+    if user_input := st.chat_input("Ask a question about your document..."):
+        st.session_state.messages.append({"role": "user", "content": user_input})
+        with st.chat_message("user"):
+            st.markdown(user_input)
+            
+        with st.chat_message("assistant"):
+            with st.spinner("Thinking..."):
+                try:
+                    response = st.session_state.qa_chain.invoke({"query": user_input})
+                    answer = response['result']
+                    
+                    st.markdown(answer)
+                    st.session_state.messages.append({"role": "assistant", "content": answer})
+                    
+                    # Optional: Show sources
+                    with st.expander("🔎 View Source Context"):
+                        for doc in response['source_documents']:
+                            st.caption(f"Page {doc.metadata.get('page', '?')}: {doc.page_content[:200]}...")
+                            
+                except Exception as e:
+                    st.error(f"An error occurred: {e}")
+else:
+    # Empty State
+    st.info("👋 **Welcome!** Please upload a PDF in the sidebar to begin chatting.")
+    st.markdown("""
+    **How it works:**
+    1. Upload a PDF document.
+    2. Click 'Process Document'.
+    3. Ask questions and get answers based strictly on your file.
+    """)
+
+# --- Footer ---
+st.markdown("""
+<div class="footer">
+    Made with ❤️ with Streamlit and Gemma model, by Tannu Yadav
+</div>
+""", unsafe_allow_html=True)