Update app.py

app.py

@@ -1,314 +1,92 @@
-import streamlit as st
+import gradio as gr
 import os
 import tempfile
+import faiss
+import torch
+from transformers import AutoTokenizer
+from ctransformers import AutoModelForCausalLM
+from sentence_transformers import SentenceTransformer
+from pdfminer.high_level import extract_text
 
-# Show initial loading message while importing packages
-with st.spinner("Loading dependencies..."):
-    # Optional imports - handle gracefully if missing
-    try:
-        from langchain_community.document_loaders import PyPDFLoader
-        from langchain.text_splitter import RecursiveCharacterTextSplitter
-        from langchain_community.embeddings import HuggingFaceEmbeddings
-        from langchain_community.vectorstores import FAISS
-        from langchain_community.llms import HuggingFaceEndpoint
-        from langchain.chains import ConversationalRetrievalChain
-        from langchain.memory import ConversationBufferMemory
-    except ImportError as e:
-        st.error(f"Missing dependency: {str(e)}")
-        st.info("Please add the following to your requirements.txt:\n```\nlangchain\nlangchain-community\npypdf\nfaiss-cpu\nsentence-transformers\nhuggingface-hub\n```")
-        st.stop()
+# Load Sentence Transformer for Embedding
+embedder = SentenceTransformer("all-MiniLM-L6-v2")
 
-# Set page configuration
-st.set_page_config(page_title="RAG PDF Chatbot", page_icon="📚", layout="wide")
+# Load the LLM (Free & Local Inference)
+llm = AutoModelForCausalLM.from_pretrained("TheBloke/Mistral-7B-Instruct-v0.2-GGUF", model_type="mistral", gpu_layers=0)
 
-# Check for Hugging Face token
-if "HUGGINGFACE_TOKEN" not in os.environ:
-    try:
-        os.environ["HUGGINGFACE_TOKEN"] = st.secrets["HUGGINGFACE_TOKEN"]
-    except:
-        st.error("⚠️ Hugging Face API token not found. Please add it to your secrets.")
-        st.info("Add your token in the Hugging Face Space settings under 'Secrets' tab.")
-        st.stop()
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
 
-# Available LLM models (simple options that work well with free tier)
-AVAILABLE_MODELS = [
-    "google/flan-t5-base",
-    "google/flan-t5-large",
-    "distilbert/distilbert-base-uncased"
-]
+# Store context and FAISS index
+doc_chunks = []
+index = None
 
-# Initialize session state
-if 'processed_docs' not in st.session_state:
-    st.session_state.processed_docs = False
-if 'conversation' not in st.session_state:
-    st.session_state.conversation = None
-if 'chat_history' not in st.session_state:
-    st.session_state.chat_history = []
-if 'vector_db' not in st.session_state:
-    st.session_state.vector_db = None
-if 'sources' not in st.session_state:
-    st.session_state.sources = []
+def extract_text_from_pdf(pdf_path):
+    return extract_text(pdf_path)
 
-# Sidebar with instructions and options
-with st.sidebar:
-    st.title("📚 RAG PDF Chatbot")
-    st.markdown("""
-    ### Instructions:
-    1. Upload your PDF document(s)
-    2. Click "Process Documents"
-    3. Ask questions about your documents
-    
-    **Tech Stack:**
-    - Hugging Face LLMs
-    - FAISS Vector Store
-    - LangChain RAG Pipeline
-    """)
-    
-    # Model selection (only shown after document processing)
-    if st.session_state.processed_docs:
-        st.markdown("---")
-        st.subheader("Model Settings")
-        selected_model = st.selectbox(
-            "Select Language Model",
-            options=AVAILABLE_MODELS,
-            index=0,
-            key="model_selection"
-        )
-        
-        temperature = st.slider(
-            "Temperature", 
-            min_value=0.1, 
-            max_value=1.0, 
-            value=0.5, 
-            step=0.1,
-            help="Higher values make output more random, lower values more deterministic"
-        )
-        
-        max_tokens = st.slider(
-            "Max Tokens", 
-            min_value=128, 
-            max_value=1024, 
-            value=256, 
-            step=64,
-            help="Maximum length of generated response"
-        )
-        
-        # Add a button to refresh the conversation with new settings
-        if st.button("Update Model Settings"):
-            with st.spinner("Updating settings..."):
-                try:
-                    st.session_state.chat_history = []  # Reset chat history
-                    st.experimental_rerun()
-                except Exception as e:
-                    st.error(f"Error updating settings: {str(e)}")
+def chunk_text(text, chunk_size=500, overlap=50):
+    words = text.split()
+    chunks = []
+    for i in range(0, len(words), chunk_size - overlap):
+        chunk = " ".join(words[i:i + chunk_size])
+        chunks.append(chunk)
+    return chunks
 
-# Function to process uploaded PDFs and create vector database
-def process_pdfs(uploaded_files):
-    with st.spinner("Processing documents... This may take a few minutes."):
-        try:
-            # Create temporary directory for PDFs
-            temp_dir = tempfile.mkdtemp()
-            temp_paths = []
-            
-            # Save uploaded files to temp directory
-            for uploaded_file in uploaded_files:
-                temp_path = os.path.join(temp_dir, uploaded_file.name)
-                with open(temp_path, "wb") as f:
-                    f.write(uploaded_file.getbuffer())
-                temp_paths.append(temp_path)
-            
-            # Load and process documents
-            documents = []
-            for temp_path in temp_paths:
-                loader = PyPDFLoader(temp_path)
-                documents.extend(loader.load())
-            
-            # Split documents into chunks
-            text_splitter = RecursiveCharacterTextSplitter(
-                chunk_size=1000,
-                chunk_overlap=200,
-                length_function=len
-            )
-            document_chunks = text_splitter.split_documents(documents)
-            
-            # Create embeddings
-            embeddings = HuggingFaceEmbeddings(
-                model_name="sentence-transformers/all-mpnet-base-v2",  # Use a different model
-                cache_folder="./embedding_cache"  # Cache embeddings
-            )
-            
-            # Create vector store
-            vectorstore = FAISS.from_documents(document_chunks, embeddings)
-            
-            # Save to session state
-            st.session_state.vector_db = vectorstore
-            st.session_state.processed_docs = True
-            
-            # Initialize LLM with default model
-            initialize_llm("google/flan-t5-base", 0.5, 256)
-            
-            return True
-            
-        except ImportError as e:
-            # Special handling for missing dependencies
-            missing_package = str(e).split("'")[1] if "'" in str(e) else str(e)
-            st.error(f"Missing dependency: {missing_package}")
-            st.info(f"Please install the required package: `pip install {missing_package}`")
-            return False
-        except Exception as e:
-            st.error(f"Error processing documents: {str(e)}")
-            return False
+def create_faiss_index(chunks):
+    embeddings = embedder.encode(chunks)
+    dim = embeddings.shape[1]
+    idx = faiss.IndexFlatL2(dim)
+    idx.add(embeddings)
+    return idx, embeddings
 
-# Initialize LLM and chain
-def initialize_llm(model_name, temp, max_len):
-    try:
-        # Create LLM
-        llm = HuggingFaceEndpoint(
-            repo_id=model_name,
-            huggingfacehub_api_token=os.environ["HUGGINGFACE_TOKEN"],
-            model_kwargs={
-                "temperature": temp,
-                "max_length": max_len
-            }
-        )
-        
-        # Create memory
-        memory = ConversationBufferMemory(
-            memory_key="chat_history",
-            return_messages=True
-        )
-        
-        # Create chain
-        chain = ConversationalRetrievalChain.from_llm(
-            llm=llm,
-            retriever=st.session_state.vector_db.as_retriever(search_kwargs={"k": 3}),
-            memory=memory,
-            return_source_documents=True
-        )
-        
-        st.session_state.conversation = chain
-        return True
-    except Exception as e:
-        st.error(f"Error initializing LLM: {str(e)}")
-        return False
+def retrieve_relevant_chunks(query, chunks, idx, k=3):
+    query_embedding = embedder.encode([query])
+    scores, indices = idx.search(query_embedding, k)
+    return [chunks[i] for i in indices[0]]
 
-# Handle conversation
-def handle_conversation(question):
-    try:
-        # Get response from model
-        response = st.session_state.conversation({"question": question})
-        answer = response["answer"]
-        
-        # Store sources
-        sources = response.get("source_documents", [])
-        st.session_state.sources = sources
-        
-        # Update history
-        st.session_state.chat_history.append((question, answer))
-        
-        return True
-    except Exception as e:
-        st.error(f"Error generating response: {str(e)}")
-        return False
+def generate_prompt(query, retrieved_chunks):
+    context = "\n\n".join(retrieved_chunks)
+    prompt = f"""You are a helpful assistant. Use the following context to answer the user's question.
 
-# Main UI
-st.title("📚 Question Answering with PDF Documents")
+Context:
+{context}
 
-# Requirements info
-with st.expander("📋 Requirements and Setup"):
-    st.markdown("""
-    ### Required packages:
-    ```
-    langchain==0.0.267
-    langchain-community==0.0.6
-    pypdf==3.15.1
-    sentence-transformers==2.2.2
-    faiss-cpu==1.7.4
-    huggingface-hub==0.16.4
-    ```
-    
-    ### Hugging Face Space Setup:
-    1. Create a new Space with Streamlit SDK
-    2. Add your Hugging Face API token in Settings → Secrets as `HUGGINGFACE_TOKEN`
-    3. Upload this code and the requirements.txt file
-    """)
+Question:
+{query}
 
-# File upload section
-uploaded_files = st.file_uploader(
-    "Upload PDF documents",
-    type=["pdf"],
-    accept_multiple_files=True
-)
+Answer:"""
+    return prompt
 
-# Process documents button
-if uploaded_files:
-    if not st.session_state.processed_docs:
-        if st.button("Process Documents"):
-            success = process_pdfs(uploaded_files)
-            if success:
-                st.success("✅ Documents processed successfully! You can now ask questions.")
-                st.balloons()
-    else:
-        # Show success message if already processed
-        st.success("✅ Documents already processed. Ask questions below.")
+def llm_answer(prompt):
+    response = llm(prompt, max_new_tokens=256, temperature=0.7)
+    return response
 
-# Chat interface
-st.markdown("---")
-if st.session_state.processed_docs and st.session_state.conversation is not None:
-    st.subheader("💬 Chat with your Documents")
-    
-    # Display chat history
-    for i, (question, answer) in enumerate(st.session_state.chat_history):
-        st.markdown(f"**You:** {question}")
-        st.markdown(f"**Assistant:** {answer}")
-        st.markdown("---")
-    
-    # Question input and buttons
-    col1, col2, col3 = st.columns([3, 1, 1])
-    
-    with col1:
-        question = st.text_input("Ask a question about your documents:", key="question_input")
-    
-    with col2:
-        submit_button = st.button("Submit")
-    
-    with col3:
-        clear_button = st.button("Clear Chat")
-        if clear_button:
-            st.session_state.chat_history = []
-            st.experimental_rerun()
-    
-    # Process question
-    if question and submit_button:
-        with st.spinner("Generating answer..."):
-            # Update model settings if changed
-            model = st.session_state.get("model_selection", AVAILABLE_MODELS[0])
-            temp = st.session_state.get("temperature", 0.5)
-            max_len = st.session_state.get("max_tokens", 256)
-            
-            # Reinitialize if needed
-            initialize_llm(model, temp, max_len)
-            
-            # Process question
-            success = handle_conversation(question)
-            if success:
-                # Show sources if available
-                if st.session_state.sources:
-                    with st.expander("📄 Source Documents"):
-                        for i, doc in enumerate(st.session_state.sources[:3]):
-                            source_page = doc.metadata.get("page", 0) + 1  # Convert to 1-indexed
-                            st.markdown(f"**Source {i+1} (Page {source_page}):**")
-                            st.text_area(f"Content", 
-                                        value=doc.page_content[:500] + ("..." if len(doc.page_content) > 500 else ""), 
-                                        height=150, 
-                                        key=f"source_{i}")
-                
-                # Force refresh to show new messages
-                st.experimental_rerun()
-else:
-    if not st.session_state.processed_docs:
-        st.info("📄 Please upload and process PDF documents to start chatting.")
+def process_pdf(file):
+    global doc_chunks, index
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
+        tmp.write(file.read())
+        tmp.flush()
+        text = extract_text_from_pdf(tmp.name)
+    doc_chunks = chunk_text(text)
+    index, _ = create_faiss_index(doc_chunks)
+    return "✅ PDF uploaded and indexed. You can start chatting now!"
 
-# Footer
-st.markdown("---")
-st.caption("RAG PDF Chatbot powered by LangChain and Hugging Face")
+def chat_with_pdf(user_input):
+    if not doc_chunks or not index:
+        return "❌ Please upload a PDF first."
+    retrieved = retrieve_relevant_chunks(user_input, doc_chunks, index)
+    prompt = generate_prompt(user_input, retrieved)
+    return llm_answer(prompt)
+
+# Gradio Interface
+with gr.Blocks() as demo:
+    gr.Markdown("# 🤖 Chat with your PDF (Free & Local LLM)")
+
+    with gr.Row():
+        pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
+        upload_button = gr.Button("Process PDF")
+
+    chatbot = gr.ChatInterface(fn=chat_with_pdf, textbox=gr.Textbox(placeholder="Ask something from the PDF...", lines=2))
+
+    upload_button.click(fn=process_pdf, inputs=[pdf_input], outputs=[chatbot.textbox])
+
+demo.launch()