batch_ids

tools.py

@@ -41,14 +41,14 @@ memoires_ds = load_dataset("eliot-hub/memoires_vec_800", split="data", token=HF_
 batched_ds = memoires_ds.batch(batch_size=41000)
 client = chromadb.Client()
 collection = client.get_or_create_collection(name="embeddings_mxbai") 
-for batch in tqdm(batched_ds, desc=f"[{datetime.now().strftime('%H:%M:%S')}] Processing dataset batches"):
+for batch in tqdm(batched_ds, desc="Processing dataset batches"): 
     collection.add(
-        ids=batch["ids"],
-        # metadatas=batch["metadatas"],
-        documents=batch["documents"],
-        embeddings=batch["embeddings"],      
+        ids=batch["id"],
+        metadatas=batch["metadata"],
+        documents=batch["document"],
+        embeddings=batch["embedding"],      
     )
-print("DB done")
+print(f"Collection complete: {collection.count()}")
 del memoires_ds, batched_ds
 
 llm_4o = ChatOpenAI(model="gpt-4o-mini", api_key=OPENAI_API_KEY, temperature=0)