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


# 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!")