Update app.py

app.py

@@ -12,74 +12,89 @@ MODEL_NAME = "all-MiniLM-L6-v2"
 INDEX_FILE = "faiss_index.pkl"
 DOCS_FILE = "contexts.pkl"
 
-# Set up Groq client
+# Groq API client
 client = Groq(api_key=os.environ.get("MY_KEY"))
 
-# UI
-st.set_page_config(page_title="RAG App with Groq", layout="wide")
-st.title("🧠 Retrieval-Augmented Generation (RAG) App")
+# Streamlit page setup
+st.set_page_config(page_title="RAG App", layout="wide")
+st.title("🧠 Retrieval-Augmented Generation (RAG) with Groq")
 
-# Load or create vector DB
+# Function to load or create database
 @st.cache_resource
 def setup_database():
-    st.info("Loading dataset and setting up database...")
+    st.info("Setting up vector database...")
     progress = st.progress(0)
 
+    # Step 1: Load dataset
     dataset = load_dataset(DATASET_NAME, split="train")
     contexts = [entry["context"] for entry in dataset]
+    progress.progress(25)
 
+    # Step 2: Compute embeddings
     embedder = SentenceTransformer(MODEL_NAME)
     embeddings = embedder.encode(contexts, show_progress_bar=True)
+    progress.progress(50)
 
+    # Step 3: Build FAISS index
     dimension = embeddings[0].shape[0]
-    index = faiss.IndexFlatL2(dimension)
-    index.add(embeddings)
+    faiss_index = faiss.IndexFlatL2(dimension)
+    faiss_index.add(embeddings)
+    progress.progress(75)
 
-    # Save index and contexts
+    # Step 4: Save index and contexts for future use
     with open(INDEX_FILE, "wb") as f:
-        pickle.dump(index, f)
+        pickle.dump(faiss_index, f)
     with open(DOCS_FILE, "wb") as f:
         pickle.dump(contexts, f)
 
     progress.progress(100)
-    return index, contexts
+    st.success("Database setup complete!")
+    return faiss_index, contexts
 
-# Load existing index or build
+# Check if the index and contexts are saved, otherwise set up
 if os.path.exists(INDEX_FILE) and os.path.exists(DOCS_FILE):
     with open(INDEX_FILE, "rb") as f:
         faiss_index = pickle.load(f)
     with open(DOCS_FILE, "rb") as f:
         all_contexts = pickle.load(f)
+    st.info("Loaded existing database.")
 else:
     faiss_index, all_contexts = setup_database()
 
-# Sample questions
+# UI for sample questions
 sample_questions = [
-    "What is the role of Falcon RefinedWeb in this dataset?",
-    "How can retrieval improve language generation?",
-    "Explain the purpose of the RAG dataset."
+    "What is the purpose of the RAG dataset?",
+    "How does Falcon RefinedWeb contribute to this dataset?",
+    "What are the benefits of using retrieval-augmented generation?",
+    "Explain the structure of the RAG-1200 dataset.",
 ]
 
 st.subheader("Ask a question based on the dataset:")
-question = st.text_input("Your question", value=sample_questions[0])
+question = st.text_input("Enter your question:", value=sample_questions[0])
 
 if st.button("Ask"):
-    with st.spinner("Retrieving relevant context and generating answer..."):
-        embedder = SentenceTransformer(MODEL_NAME)
-        question_embedding = embedder.encode([question])
-        D, I = faiss_index.search(question_embedding, k=1)
-
-        retrieved_context = all_contexts[I[0][0]]
-        prompt = f"Context: {retrieved_context}\n\nQuestion: {question}\n\nAnswer:"
-
-        response = client.chat.completions.create(
-            messages=[{"role": "user", "content": prompt}],
-            model="llama-3-70b-8192"
-        )
-
-        answer = response.choices[0].message.content
-        st.success("Answer:")
-        st.write(answer)
-
-        with st.expander("Retrieved Context"):
-            st.markdown(retrieved_context)
+    if question.strip() == "":
+        st.warning("Please enter a question.")
+    else:
+        with st.spinner("Retrieving and generating answer..."):
+            # Embed user query
+            embedder = SentenceTransformer(MODEL_NAME)
+            query_embedding = embedder.encode([question])
+            D, I = faiss_index.search(query_embedding, k=1)
+
+            # Get closest context
+            context = all_contexts[I[0][0]]
+            prompt = f"Context: {context}\n\nQuestion: {question}\n\nAnswer:"
+
+            # Call Groq model
+            response = client.chat.completions.create(
+                messages=[{"role": "user", "content": prompt}],
+                model="llama3-70b-8192"
+            )
+
+            answer = response.choices[0].message.content
+            st.success("Answer:")
+            st.markdown(answer)
+
+            with st.expander("🔍 Retrieved Context"):
+                st.markdown(context)