Update prepare_wikipedia_collection() to prepare in batches

Author: punitarani

rabbithole/wikipedia.py

@@ -8,26 +8,29 @@
 wikipedia_collection = client.create_collection("wikipedia")
 
 
-def prepare_wiki_collection():
+def prepare_wikipedia_collection(batch_size: int = 10000):
     """
     Prepare the wikipedia collection
 
     NOTE: Only needs to be run once to prepare the collection for the first time
     """
-    ids, embeddings, documents, metadatas = zip(
-        *[(i, row['emb'], row['text'], row['title']) for i, row in enumerate(wikipedia_dataset)])
+    total_rows = len(wikipedia_dataset)
+    for i in range(0, total_rows, batch_size):
+        batch_data = wikipedia_dataset[i: i + batch_size]
+        ids, embeddings, documents, metadatas = zip(
+            *[(i, row['emb'], row['text'], row['title']) for i, row in enumerate(batch_data)])
 
-    wikipedia_collection.add(
-        ids=ids,
-        embeddings=embeddings,
-        documents=documents,
-        metadatas=metadatas,
-    )
+        wikipedia_collection.add(
+            ids=ids,
+            embeddings=embeddings,
+            documents=documents,
+            metadatas=metadatas,
+        )
 
 
 if __name__ == "__main__":
     print(wikipedia_dataset.info)
 
     print("Preparing wikipedia collection...")
-    prepare_wiki_collection()
+    prepare_wikipedia_collection()
     print("Done!")