From 2cdf43d8b7e6e8e87ae9d6518a562002d8671079 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Mon, 12 Feb 2024 04:28:41 +1100
Subject: [PATCH] Chat Templates

---
 unsloth/__init__.py       |  1 +
 unsloth/chat_templates.py | 39 +++++++++++++++++++++++++++++++++++++++
 unsloth/models/_utils.py  |  1 +
 unsloth/models/llama.py   |  2 +-
 4 files changed, 42 insertions(+), 1 deletion(-)
 create mode 100644 unsloth/chat_templates.py

diff --git a/unsloth/__init__.py b/unsloth/__init__.py
index 66c10e28e7..d052b333f2 100644
--- a/unsloth/__init__.py
+++ b/unsloth/__init__.py
@@ -82,3 +82,4 @@ pass
 
 from .models import *
 from .save import *
+from .chat_templates import *
diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py
new file mode 100644
index 0000000000..b019111057
--- /dev/null
+++ b/unsloth/chat_templates.py
@@ -0,0 +1,39 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+__all__ = [
+    "add_chat_template",
+]
+
+TEMPLATES = \
+{
+    "chatml" : \
+        "{% for message in messages %}"\
+            "{% if message['from'] == 'human' %}"\
+                "{{'<|im_start|>user\n' + message['value'] + '<|im_end|>\n'}}"\
+            "{% elif message['from'] == 'gpt' %}"\
+                "{{'<|im_start|>assistant\n' + message['value'] + '<|im_end|>\n' }}"\
+            "{% else %}"\
+                "{{ '<|im_start|>system\n' + message['value'] + '<|im_end|>\n' }}"\
+            "{% endif %}"\
+        "{% endfor %}"\
+        "{% if add_generation_prompt %}"\
+            "{{ '<|im_start|>assistant\n' }}"\
+        "{% endif %}",
+}
+
+
+def add_chat_template(tokenizer, method = "chatml"):
+    tokenizer.chat_template = TEMPLATES[method.lower()]
+pass
diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 617b8509d9..677780b0b7 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -16,6 +16,7 @@ import torch
 from typing import Union, Optional, List, Any, Callable
 import warnings
 warnings.filterwarnings(action = "ignore", category = UserWarning, module = "torch")
+warnings.filterwarnings(action = "ignore", category = UserWarning, module = "huggingface_hub")
 import bitsandbytes as bnb
 from transformers.models.llama.modeling_llama import logger
 from transformers import AutoTokenizer
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index d35a35c695..af0b3979c9 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -540,7 +540,7 @@ def LlamaModel_fast_forward(
 
     hidden_states = inputs_embeds
 
-    if past_key_values is None and self.gradient_checkpointing and self.training:
+    if past_key_values is None and self.training:
         use_cache = False
         # if use_cache:
         #     logger.warning_once(