make batch smaller

2024-12-18 16:23:50 -05:00
parent d584354ee4
commit e633a12dbe
1 changed files with 5 additions and 1 deletions
--- a/src/axolotl/datasets.py
+++ b/src/axolotl/datasets.py
@@ -53,7 +53,11 @@ class TokenizedPromptDataset(Dataset):
            map_kwargs["batched"] = True
            map_kwargs["batch_size"] = 100
        if self.prompt_tokenizer.filter_rows:
-            dataset = dataset.filter(self.prompt_tokenizer.filter_rows)
+            dataset = dataset.filter(
+                self.prompt_tokenizer.filter_rows,
+                num_proc=num_proc,
+                desc="Filtering Rows",
+            )
        return dataset.map(
            self.prompt_tokenizer.tokenize_prompt,
            num_proc=num_proc,