From 71c4aed1be0857deb74887dafd232a026017011b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 20 Jul 2024 13:22:36 -0700 Subject: [PATCH] Update tokenizer_utils.py --- unsloth/tokenizer_utils.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 3f75d16863..0469f4d61b 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -688,7 +688,12 @@ def fix_untrained_tokens(model, tokenizer, train_dataset, eps = 1e-16): pass # Get untrained tokens - indicator_untrained = torch.amax(embedding_matrix, axis = 1) <= eps + indicator_untrained1 = torch.amax(embedding_matrix, axis = 1) <= eps + # Check lm_head as well + indicator_untrained2 = torch.amax(lm_head_matrix, axis = 1) <= eps + # Combine both checks + indicator_untrained = indicator_untrained1 & indicator_untrained2 + where_untrained = torch.where(indicator_untrained)[0] n_untrained = where_untrained.shape[0] n_trained = embedding_matrix.shape[0] - n_untrained