Update app.py

app.py

@@ -4,6 +4,8 @@ from groq import Groq
 from PyPDF2 import PdfReader
 from docx import Document
 from sentence_transformers import SentenceTransformer
+import faiss
+import numpy as np
 
 # Initialize Groq API Client
 client = Groq(api_key=os.environ.get("Groq_Api"))
@@ -16,8 +18,8 @@ uploaded_file = st.file_uploader("Upload a PDF or DOCX file", type=["pdf", "docx
 
 if uploaded_file:
     st.write(f"**File Name:** {uploaded_file.name}")  # Display file name
-    
-    # Read PDF or DOCX content
+
+    # Extract Text
     def extract_text(file):
         if file.name.endswith(".pdf"):
             reader = PdfReader(file)
@@ -35,15 +37,29 @@ if uploaded_file:
         query = st.text_input("Enter your question")
 
         if query:
-            # Chunk & Tokenize
+            # Load Sentence Transformer Model
             model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
-            chunks = [file_text[i:i + 512] for i in range(0, len(file_text), 512)]
-            embeddings = model.encode(chunks)
 
-            # Query with Groq API
+            # Chunk & Embed Text
+            chunk_size = 512
+            chunks = [file_text[i:i + chunk_size] for i in range(0, len(file_text), chunk_size)]
+            embeddings = model.encode(chunks, convert_to_numpy=True)
+
+            # Build FAISS Index for Fast Retrieval
+            index = faiss.IndexFlatL2(embeddings.shape[1])
+            index.add(embeddings)
+
+            # Query Embedding
+            query_embedding = model.encode([query], convert_to_numpy=True)
+            _, retrieved_idx = index.search(query_embedding, k=3)
+
+            # Retrieve Top 3 Relevant Chunks
+            relevant_text = " ".join([chunks[i] for i in retrieved_idx[0]])
+
+            # Query Groq API with relevant chunks only
             chat_completion = client.chat.completions.create(
                 messages=[
-                    {"role": "user", "content": f"Answer based on this document: {query}\n\n{file_text}"},
+                    {"role": "user", "content": f"Answer based on this document: {query}\n\n{relevant_text}"},
                 ],
                 model="llama-3.3-70b-versatile",
             )
@@ -52,6 +68,6 @@ if uploaded_file:
             answer = chat_completion.choices[0].message.content
             st.subheader("Answer:")
             st.write(answer)
+
     else:
         st.error("Failed to extract text from the file. Please check the format.")
-