From eb61377632e3df3f75616dfd5d3e1a267942afef Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Thu, 14 Mar 2024 19:02:29 +1100
Subject: [PATCH] Fix Colab

---
 README.md                 |  2 +-
 pyproject.toml            | 18 +++++++++---------
 unsloth/chat_templates.py | 12 ++++++++++--
 3 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index 759057f2a4..412e3ffda2 100644
--- a/README.md
+++ b/README.md
@@ -59,7 +59,7 @@ All notebooks are **beginner friendly**! Add your dataset, click "Run All", and
 - No change of hardware. Supports NVIDIA GPUs since 2018+. Minimum CUDA Capability 7.0 (V100, T4, Titan V, RTX 20, 30, 40x, A100, H100, L40 etc) [Check your GPU!](https://developer.nvidia.com/cuda-gpus) GTX 1070, 1080 works, but is slow.
 - Works on **Linux** and **Windows** via WSL.
 - Supports 4bit and 16bit QLoRA / LoRA finetuning via [bitsandbytes](https://github.com/TimDettmers/bitsandbytes).
-- Open source trains 5x faster - see [Unsloth Pro](https://unsloth.ai/) for **30x faster training**!
+- Open source trains 5x faster - see [Unsloth Pro](https://unsloth.ai/) for up to **30x faster training**!
 - If you trained a model with 🦥Unsloth, you can use this cool sticker! &nbsp; <img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/made with unsloth.png" height="50" align="center" />
 
 
diff --git a/pyproject.toml b/pyproject.toml
index 7fcdfa4955..2387e5597f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -134,15 +134,6 @@ kaggle-new = [
 conda = [
     "unsloth[huggingface]",
 ]
-colab = [
-    "unsloth[cu121]",
-]
-colab-ampere = [
-    "unsloth[cu121]",
-    "packaging",
-    "ninja",
-    "flash-attn",
-]
 colab-torch211 = [
     "unsloth[huggingface]",
     "bitsandbytes",
@@ -169,6 +160,15 @@ colab-ampere-torch220 = [
     "ninja",
     "flash-attn",
 ]
+colab = [
+    "unsloth[colab-ampere-torch220]",
+]
+colab-ampere = [
+    "unsloth[colab-ampere-torch220]",
+    "packaging",
+    "ninja",
+    "flash-attn",
+]
 cu118-ampere = [
     "unsloth[huggingface]",
     "bitsandbytes",
diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py
index 5487e9dcab..7d205ebc5f 100644
--- a/unsloth/chat_templates.py
+++ b/unsloth/chat_templates.py
@@ -259,8 +259,10 @@ def get_chat_template(
         assert("Unsloth: Can only map new tokens to EOS for now. Adding new tokens is not yet supported.")
     pass
 
-    if tokenizer.__class__.__name__.startswith("Gemma") and chat_template == "chatml":
-        chat_template = "gemma_chatml"
+    IS_GEMMA = False
+    if tokenizer.__class__.__name__.startswith("Gemma"):
+        if chat_template == "chatml": chat_template = "gemma_chatml"
+        IS_GEMMA = True
     pass
 
     old_padding_side = tokenizer.padding_side
@@ -338,6 +340,12 @@ def get_chat_template(
         .replace("'user'",      "'" + mapping["user"]      + "'")\
         .replace("'assistant'", "'" + mapping["assistant"] + "'")
 
+    # Careful on Gemma
+    # bos_token is a must or else losses become too high
+    if IS_GEMMA and not chat_template.startswith("{{ bos_token }}"):
+        chat_template = "{{ bos_token }}" + chat_template
+    pass
+
     _, tokenizer = patch_tokenizer(model = None, tokenizer = tokenizer)
     tokenizer.padding_side  = old_padding_side
     tokenizer.chat_template = chat_template