remove columns after tokenizing for pretraining (#571)

axolotl-ai-cloud · Sep 14, 2023 · 1157950 · 1157950
1 parent 3b18c96
commit 1157950
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py
@@ -644,8 +644,8 @@ def load_pretraining_dataset(path, tokenizer, max_tokens=2048, seed=42):
         encode,
         batched=True,
         input_columns="text",
-        remove_columns=[
-            "text",
-        ],
+        # remove all the existing columns after mapping since they end up having
+        # a different length than the encoded/tokenized column
+        remove_columns=dataset.features.keys(),
     )
     return dataset