from datasets import load_dataset
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import gradio as gr

# Load a small subset (10,000 rows)
dataset = load_dataset("wiki40b", "en", split="train[:10000]")

# Extract only text
docs = [d["text"] for d in dataset]

print("Loaded dataset with", len(docs), "documents.")

# Load embedding model
embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# Convert texts to embeddings
embeddings = embed_model.encode(docs, show_progress_bar=True)

# Store in FAISS index
dimension = embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(np.array(embeddings))

print("Stored embeddings in FAISS!")

# Search function
def search_wikipedia(query, top_k=3):
    query_embedding = embed_model.encode([query])
    distances, indices = index.search(np.array(query_embedding), top_k)
    results = [docs[i] for i in indices[0]]
    return "\n\n".join(results)

# Gradio Interface
iface = gr.Interface(
    fn=search_wikipedia, 
    inputs="text", 
    outputs="text", 
    title="Wikipedia Search RAG",
    description="Enter a query and retrieve relevant Wikipedia passages."
)

iface.launch()