app.py

import os
import streamlit as st
import faiss
import pickle
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from groq import Groq

# Constants
DATASET_NAME = "neural-bridge/rag-dataset-1200"
MODEL_NAME = "all-MiniLM-L6-v2"
INDEX_FILE = "faiss_index.pkl"
DOCS_FILE = "contexts.pkl"

# Set up Groq client
client = Groq(api_key=os.environ.get("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"))

# UI
st.set_page_config(page_title="RAG App with Groq", layout="wide")
st.title("🧠 Retrieval-Augmented Generation (RAG) App")

# 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."
]

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 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)