Update app.py

app.py

@@ -3,9 +3,10 @@ from sentence_transformers import SentenceTransformer
 import faiss
 import numpy as np
 import gradio as gr
+import chromadb
 
 # Load a small subset (10,000 rows)
-dataset = load_dataset("wiki40b", "en", split="train[:10000]")
+dataset = load_dataset("wiki40b", "en", split="train[:1000]")
 
 # Extract only text
 docs = [d["text"] for d in dataset]
@@ -18,19 +19,29 @@ 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!")
+# Initialize ChromaDB client
+chroma_client = chromadb.PersistentClient(path="./chroma_db")  # Stores data persistently
+collection = chroma_client.get_or_create_collection(name="wikipedia_docs")
 
-# Search function
+# Store embeddings in ChromaDB
+for i, (doc, embedding) in enumerate(zip(docs, embeddings)):
+    collection.add(
+        ids=[str(i)],  # Unique ID for each doc
+        embeddings=[embedding.tolist()],  # Convert numpy array to list
+        documents=[doc]
+    )
+
+print("Stored embeddings in ChromaDB!")
+
+# Search function using ChromaDB
 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)
+    query_embedding = embed_model.encode([query]).tolist()
+    results = collection.query(
+        query_embeddings=query_embedding, 
+        n_results=top_k
+    )
+    return "\n\n".join(results["documents"][0])  # Return top results
 
 # Gradio Interface
 iface = gr.Interface(