Skip to content

Commit

Permalink
fixes for smoke tests
Browse files Browse the repository at this point in the history
  • Loading branch information
winglian committed Jan 22, 2024
1 parent 5824228 commit f77a465
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 0 deletions.
6 changes: 6 additions & 0 deletions src/axolotl/utils/trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,12 @@ def process_datasets_for_packing(cfg, train_dataset, eval_dataset, tokenizer):
if 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:
eval_dataset = eval_dataset.remove_columns("token_type_ids")

return train_dataset, eval_dataset


Expand Down
1 change: 1 addition & 0 deletions tests/e2e/patched/test_falcon_samplepack.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ def test_qlora(self, temp_dir):
"lora_alpha": 32,
"lora_dropout": 0.1,
"lora_target_linear": True,
"lora_modules_to_save": ["word_embeddings", "lm_head"],
"val_set_size": 0.1,
"special_tokens": {
"bos_token": "<|endoftext|>",
Expand Down

0 comments on commit f77a465

Please sign in to comment.