Update app.py

app.py

@@ -1,71 +1,100 @@
 import streamlit as st
-import fitz  # PyMuPDF
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.vectorstores import FAISS
-import pickle
+import faiss
+import numpy as np
+from sentence_transformers import SentenceTransformer
 from groq import Groq
 import os
+import pypdf
+from langchain.text_splitter import RecursiveCharacterTextSplitter
 
-# Streamlit App
-st.title("RAG-based PDF Query App")
-st.write("Upload a PDF, extract its content, and query it using Groq API.")
+# Set Groq API key
+GROQ_API_KEY = os.getenv("GROQ_API_KEY", "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")
+
+# Load embedding model
+embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
+
+# Set up Groq client
+client = Groq(api_key=GROQ_API_KEY)
+
+# Function to extract text from PDF
+def extract_text_from_pdf(uploaded_file):
+    reader = pypdf.PdfReader(uploaded_file)
+    text = "\n".join([page.extract_text() for page in reader.pages if page.extract_text()])
+    return text
+
+# Function to create text chunks
+def create_chunks(text, chunk_size=500, chunk_overlap=100):
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    return text_splitter.split_text(text)
+
+# Function to create and save FAISS index
+def create_faiss_index(chunks):
+    embeddings = embedding_model.encode(chunks, convert_to_numpy=True)
+    
+    # Create FAISS index
+    dimension = embeddings.shape[1]
+    index = faiss.IndexFlatL2(dimension)
+    index.add(embeddings)
+
+    # Save FAISS index and chunks
+    faiss.write_index(index, "faiss_index.bin")
+    np.save("chunks.npy", chunks)
+
+    return index, chunks
+
+# Function to search FAISS
+def search_faiss(query, index, chunks, top_k=2):
+    if index is None:
+        return ["No database found. Please upload a PDF first."]
+    
+    query_embedding = embedding_model.encode([query], convert_to_numpy=True)
+    distances, indices = index.search(query_embedding, top_k)
+    return [chunks[i] for i in indices[0]]
+
+# Function to query Groq
+def query_groq(query):
+    chat_completion = client.chat.completions.create(
+        messages=[{"role": "user", "content": query}],
+        model="llama-3.3-70b-versatile",
+    )
+    return chat_completion.choices[0].message.content
+
+# Streamlit UI
+st.set_page_config(page_title="RAG Chatbot", page_icon="🤖")
+st.title("📄 RAG-Based Chatbot with FAISS & Groq")
 
 # Upload PDF
-uploaded_file = st.file_uploader("Upload a PDF", type=["pdf"])
-
-if uploaded_file is not None:
-    # Extract text from PDF
-    def extract_text_from_pdf(uploaded_file):
-        doc = fitz.open(stream=uploaded_file.read(), filetype="pdf")
-        text = ""
-        for page in doc:
-            text += page.get_text()
-        return text
-
-    pdf_text = extract_text_from_pdf(uploaded_file)
-    st.success("PDF uploaded and extracted successfully!")
-
-    # Chunk & Tokenize Text
-    def chunk_text(text, chunk_size=500, overlap=50):
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=overlap)
-        return text_splitter.split_text(text)
-
-    chunks = chunk_text(pdf_text)
-
-    # Create Embeddings & Store in FAISS
-    embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-    vector_store = FAISS.from_texts(chunks, embedding_model)
-
-    # Save FAISS index
-    with open("faiss_index.pkl", "wb") as f:
-        pickle.dump(vector_store, f)
-
-    st.success("Document processed and stored in vector database!")
-
-    # Query Section
-    query = st.text_input("Enter your query:")
-
-    if st.button("Search"):
-        if query:
-            # Load FAISS index
-            with open("faiss_index.pkl", "rb") as f:
-                vector_store = pickle.load(f)
-
-            docs = vector_store.similarity_search(query, k=3)
-            context = "\n".join([doc.page_content for doc in docs])
-
-            
-            GROQ_API_KEY = os.environ["GROQ_API_KEY"]  # Ensure you have stored it properly
-            client = Groq(api_key=GROQ_API_KEY)
-
-            response = client.chat.completions.create(
-                messages=[{"role": "user", "content": context + "\n\n" + query}],
-                model="llama-3.3-70b-versatile",
-            )
-
-            st.subheader("Response:")
-            st.write(response.choices[0].message.content)
-        else:
-            st.warning("Please enter a query to search.")
-
+uploaded_file = st.file_uploader("📤 Upload a PDF file", type="pdf")
+
+if uploaded_file:
+    with st.spinner("🔄 Processing PDF..."):
+        text = extract_text_from_pdf(uploaded_file)
+        chunks = create_chunks(text)
+
+        # Create FAISS index
+        index, chunks = create_faiss_index(chunks)
+
+        # Store in session state
+        st.session_state["faiss_index"] = index
+        st.session_state["chunks"] = chunks
+
+        st.success("✅ PDF processed successfully!")
+
+# Load FAISS index if available
+if "faiss_index" in st.session_state and "chunks" in st.session_state:
+    index = st.session_state["faiss_index"]
+    chunks = st.session_state["chunks"]
+else:
+    index = None
+    chunks = []
+
+# User query input
+user_query = st.text_input("💬 Ask me something about the document:")
+
+if st.button("Search") and user_query:
+    with st.spinner("🔎 Retrieving response..."):
+        retrieved_text = search_faiss(user_query, index, chunks)
+        response = query_groq("\n".join(retrieved_text))
+        
+        st.write("### 🤖 AI Response:")
+        st.write(response)