Fix falcon tokenization step (#1441) [skip ci]

* Fix falcon tokenization step * chore: lint --------- Co-authored-by: Wing Lian <[email protected]>
axolotl-ai-cloud · Mar 26, 2024 · 9ae17ae · 9ae17ae
1 parent 0453464
commit 9ae17ae
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py
@@ -124,9 +124,10 @@ def process_datasets_for_packing(cfg, train_dataset, eval_dataset):
                 eval_dataset = eval_dataset.remove_columns("attention_mask")
 
         if cfg.model_config_type == "falcon":
-            LOG.info("dropping token_type_ids column")
-            train_dataset = train_dataset.remove_columns("token_type_ids")
-            if eval_dataset:
+            LOG.info("dropping token_type_ids column if it exists")
+            if "token_type_ids" in train_dataset.column_names:
+                train_dataset = train_dataset.remove_columns("token_type_ids")
+            if eval_dataset and "token_type_ids" in eval_dataset.column_names:
                 eval_dataset = eval_dataset.remove_columns("token_type_ids")
 
         train_dataset = train_dataset.filter(