From f4d97caf5eb6a64898d71f42b96d28f40b6f2aa9 Mon Sep 17 00:00:00 2001
From: Nino Risteski <95188570+NinoRisteski@users.noreply.github.com>
Date: Fri, 14 Mar 2025 10:53:21 +0100
Subject: [PATCH] Update RMS LayerNorm implementation, and list compr. change
 in chat templates (#1974)

* Update RMS LayerNorm implementation with optimizations and testing suite

* perf: optimize list comprehension in get_ollama_eos_tokens
---
 unsloth/chat_templates.py        | 5 +----
 unsloth/kernels/rms_layernorm.py | 1 -
 2 files changed, 1 insertion(+), 5 deletions(-)

diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py
index 2c2e36182d..c10b2641a4 100644
--- a/unsloth/chat_templates.py
+++ b/unsloth/chat_templates.py
@@ -1512,10 +1512,7 @@ def get_ollama_eos_tokens(tokenizer, extra_eos_tokens = []):
 
     # Remove duplicates
     splitted = joined_text.split("\x01\x00")
-    final_eos_tokens = []
-    for old, new in zip(added_tokens_decoder, splitted):
-        if old == new: final_eos_tokens.append(old)
-    pass
+    final_eos_tokens = [old for old, new in zip(added_tokens_decoder, splitted) if old == new]
     final_eos_tokens += extra_eos_tokens
     final_eos_tokens += repeatted_tokens
 
diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py
index ce61cef72e..8f54e74908 100644
--- a/unsloth/kernels/rms_layernorm.py
+++ b/unsloth/kernels/rms_layernorm.py
@@ -256,7 +256,6 @@ def unpatch_rms_layernorm():
     except:
         pass
     return
-    return
 pass