Make dataset dump skipping embedding column

pwr-ai · Jun 1, 2024 · ea33387 · ea33387
1 parent 18b2c6a
commit ea33387
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/scripts/dataset/dump_pl_dataset.py b/scripts/dataset/dump_pl_dataset.py
@@ -50,7 +50,9 @@ def main(
     for offset in trange(start_offset, num_docs, chunk_size, desc="Chunks"):
         docs = list(
             tqdm(
-                collection.find(query, batch_size=batch_size).skip(offset).limit(chunk_size),
+                collection.find(query, {"embedding": 0}, batch_size=batch_size)
+                .skip(offset)
+                .limit(chunk_size),
                 total=chunk_size,
                 leave=False,
                 desc="Documents in chunk",