Update app.py

app.py

@@ -1,104 +1,85 @@
 import os
 import streamlit as st
-import pandas as pd
-import numpy as np
 import faiss
+import pickle
 from datasets import load_dataset
 from sentence_transformers import SentenceTransformer
 from groq import Groq
 
-# Constants for saving/loading index
-INDEX_FILE = "faiss_index.index"
-QUESTIONS_FILE = "questions.npy"
+# Constants
+DATASET_NAME = "neural-bridge/rag-dataset-1200"
+MODEL_NAME = "all-MiniLM-L6-v2"
+INDEX_FILE = "faiss_index.pkl"
+DOCS_FILE = "contexts.pkl"
 
-# Load dataset
-@st.cache_data
-def load_data():
-    dataset = load_dataset("FreedomIntelligence/RAG-Instruct", split="train")
-    df = pd.DataFrame(dataset)
-    return df[["question", "answer"]]
+# Set up Groq client
+client = Groq(api_key=os.environ.get("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"))
 
-# Build or load FAISS index
-@st.cache_resource
-def setup_faiss(data):
-    model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
-
-    if os.path.exists(INDEX_FILE) and os.path.exists(QUESTIONS_FILE):
-        st.info("🔁 Loading FAISS index from disk...")
-        index = faiss.read_index(INDEX_FILE)
-        questions = np.load(QUESTIONS_FILE, allow_pickle=True)
-    else:
-        st.info("⚙️ FAISS index not found. Building new index...")
-
-        questions = data["question"].tolist()
-        embeddings = []
-        progress_bar = st.progress(0, text="Embedding questions...")
-        total = len(questions)
-
-        for i, chunk in enumerate(np.array_split(questions, 10)):
-            emb = model.encode(chunk)
-            embeddings.extend(emb)
-            progress_bar.progress((i + 1) / 10, text=f"Embedding... {int((i + 1) * 10)}%")
-
-        embeddings = np.array(embeddings)
-        index = faiss.IndexFlatL2(embeddings.shape[1])
-        index.add(embeddings)
-
-        faiss.write_index(index, INDEX_FILE)
-        np.save(QUESTIONS_FILE, np.array(questions, dtype=object))
-
-        progress_bar.empty()
-        st.success("✅ FAISS index built and saved!")
-
-    return model, index, questions
-
-
-# Retrieve relevant context
-def retrieve_context(query, model, index, questions, data, top_k=1):
-    query_vec = model.encode([query])
-    distances, indices = index.search(np.array(query_vec), top_k)
-    results = [questions[i] + "\n\n" + data.iloc[i]["answer"] for i in indices[0]]
-    return "\n\n".join(results)
-
-# Call Groq LLM
-def query_groq(context, query):
-    prompt = f"Context:\n{context}\n\nQuestion: {query}\n\nAnswer:"
-    client = Groq(api_key=st.secrets["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"])
-    response = client.chat.completions.create(
-        messages=[{"role": "user", "content": prompt}],
-        model="llama-3-70b-8192"
-    )
-    return response.choices[0].message.content
-
-# Streamlit UI
+# UI
 st.set_page_config(page_title="RAG App with Groq", layout="wide")
-st.title("🔍 RAG App using Groq API + RAG-Instruct Dataset")
-
-# Load data and setup
-data = load_data()
-model, index, questions = setup_faiss(data)
-
-st.markdown("Ask a question based on the QA knowledge base.")
+st.title("🧠 Retrieval-Augmented Generation (RAG) App")
 
-# Optional queries
-optional_queries = [
-    "What is retrieval-augmented generation?",
-    "How can I fine-tune a language model?",
-    "What are the components of a RAG system?",
-    "Explain prompt engineering basics.",
-    "How does FAISS indexing help in RAG?"
+# Load or create vector DB
+@st.cache_resource
+def setup_database():
+    st.info("Loading dataset and setting up database...")
+    progress = st.progress(0)
+
+    dataset = load_dataset(DATASET_NAME, split="train")
+    contexts = [entry["context"] for entry in dataset]
+
+    embedder = SentenceTransformer(MODEL_NAME)
+    embeddings = embedder.encode(contexts, show_progress_bar=True)
+
+    dimension = embeddings[0].shape[0]
+    index = faiss.IndexFlatL2(dimension)
+    index.add(embeddings)
+
+    # Save index and contexts
+    with open(INDEX_FILE, "wb") as f:
+        pickle.dump(index, f)
+    with open(DOCS_FILE, "wb") as f:
+        pickle.dump(contexts, f)
+
+    progress.progress(100)
+    return index, contexts
+
+# Load existing index or build
+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)
+else:
+    faiss_index, all_contexts = setup_database()
+
+# 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."
 ]
 
-query = st.text_input("Enter your question:", value=optional_queries[0])
+st.subheader("Ask a question based on the dataset:")
+question = st.text_input("Your question", value=sample_questions[0])
+
 if st.button("Ask"):
-    with st.spinner("Retrieving and generating response..."):
-        context = retrieve_context(query, model, index, questions, data)
-        answer = query_groq(context, query)
-    st.subheader("📄 Retrieved Context")
-    st.write(context)
-    st.subheader("💬 Answer from Groq LLM")
-    st.write(answer)
-
-st.markdown("### 💡 Optional Queries to Try:")
-for q in optional_queries:
-    st.markdown(f"- {q}")
+    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)