File size: 2,672 Bytes
8b835fd 319855f 8b835fd 5e8a326 319855f 5e8a326 8b835fd 5e8a326 319855f 5e8a326 8b835fd 5e8a326 319855f 5e8a326 319855f 5e8a326 319855f 5e8a326 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
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)
|