fix: mistral nemo does not recognize token_type_ids in forward (#2233)

axolotl-ai-cloud · Jan 9, 2025 · 2e8d7c1 · 2e8d7c1
1 parent 3c1921e
commit 2e8d7c1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py
@@ -196,7 +196,7 @@ def process_datasets_for_packing(cfg, train_dataset, eval_dataset):
         if eval_dataset:
             eval_dataset = eval_dataset.remove_columns("attention_mask")
 
-    if cfg.model_config_type == "falcon":
+    if cfg.model_config_type in ["falcon", "mistral"]:
         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")