From d0a0145cedc82084077cb98bf1e1cded51981f2c Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 16 Aug 2024 19:28:43 -0700
Subject: [PATCH] untrained tokens llama 3.1 base

---
 unsloth/chat_templates.py  |  4 ++--
 unsloth/tokenizer_utils.py | 20 ++++++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py
index 7070524e0f..82f6aba148 100644
--- a/unsloth/chat_templates.py
+++ b/unsloth/chat_templates.py
@@ -876,7 +876,7 @@ def get_chat_template(
 
     # Careful on Gemma
     # bos_token is a must or else losses become too high
-    if IS_GEMMA and not chat_template.startswith("{{ bos_token }}"):
+    if IS_GEMMA and not chat_template.startswith(("{{ bos_token }}", "{{- bos_token }}")):
         chat_template = "{{ bos_token }}" + chat_template
     pass
 
@@ -1553,7 +1553,7 @@ extra_eos_tokens = None,
 
     # Check jinja tempate for bos
     if always_bos_token:
-        if not jinja_template.startswith("{{ bos_token }}"):
+        if not jinja_template.startswith(("{{ bos_token }}", "{{- bos_token }}")):
             jinja_template = "{{ bos_token }}" + jinja_template
     pass
 
diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py
index 9c0bc1c510..a4f0b33be4 100644
--- a/unsloth/tokenizer_utils.py
+++ b/unsloth/tokenizer_utils.py
@@ -827,7 +827,27 @@ def fix_untrained_tokens(model, tokenizer, train_dataset, eps = 1e-16):
     # Get untrained tokens
     indicator_untrained1 = torch.amax(embedding_matrix, axis = 1) <= eps
     # Check lm_head as well
+
+    # Does NOT work for Llama 3.1!!
     indicator_untrained2 = torch.amax(lm_head_matrix,   axis = 1) <= eps
+
+    # We instead check for repeated vectors
+    lm_head_where = torch.where(indicator_untrained1)[0]
+    lm_head_bad = lm_head_matrix[lm_head_where]
+    lm_head_bad = lm_head_bad.cpu().numpy().round(3)
+    from collections import Counter
+    counter = Counter()
+    for row in lm_head_bad: counter[hash(row.data.tobytes())] += 1
+    counter = Counter({k: c for k, c in counter.items() if c >= 2})
+
+    lm_head_where = lm_head_where.cpu().numpy()
+    final_bad_lm_head = []
+    for j, row in enumerate(lm_head_bad):
+        if hash(row.data.tobytes()) in counter:
+            final_bad_lm_head.append(lm_head_where[j])
+    indicator_untrained2 = indicator_untrained2 | torch.zeros_like(indicator_untrained2)
+    indicator_untrained2[final_bad_lm_head] = True
+
     # Combine both checks
     indicator_untrained = indicator_untrained1 & indicator_untrained2