From e6a758d2c3bbe49be8035821b5db6779f46cd4d4 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Mon, 18 Mar 2024 04:07:54 +1100
Subject: [PATCH] trust_remote_code

---
 unsloth/models/llama.py   | 10 +++++++---
 unsloth/models/loader.py  |  2 ++
 unsloth/models/mistral.py | 10 +++++++---
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 262649f7fc..8b86feb35e 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -758,6 +758,7 @@ def CausalLM_fast_forward(fast_forward_inference):
         else:
             logits = self.lm_head(hidden_states)
         pass
+        logits = logits.to(self.config.torch_dtype)
 
         loss = None
         if labels is not None:
@@ -928,6 +929,7 @@ class FastLlamaModel:
         fix_tokenizer  = True,
         model_patcher  = None,
         tokenizer_name = None,
+        trust_remote_code = False,
         **kwargs,
     ):
         if model_patcher is None: model_patcher = FastLlamaModel
@@ -988,6 +990,7 @@ class FastLlamaModel:
             token                   = token,
             rope_scaling            = rope_scaling,
             max_position_embeddings = max_position_embeddings,
+            trust_remote_code       = trust_remote_code,
             **kwargs,
         )
 
@@ -995,9 +998,10 @@ class FastLlamaModel:
         tokenizer_name = model_name if tokenizer_name is None else tokenizer_name
         tokenizer = AutoTokenizer.from_pretrained(
             tokenizer_name,
-            model_max_length = max_position_embeddings,
-            padding_side     = "right",
-            token            = token,
+            model_max_length  = max_position_embeddings,
+            padding_side      = "right",
+            token             = token,
+            trust_remote_code = trust_remote_code,
         )
 
         model, tokenizer = patch_tokenizer(model, tokenizer)
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index 47b568ae2a..29d25f3c20 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -74,6 +74,7 @@ class FastLanguageModel(FastLlamaModel):
         device_map     = "sequential",
         rope_scaling   = None,
         fix_tokenizer  = True,
+        trust_remote_code = False,
         use_gradient_checkpointing = True,
         *args, **kwargs,
     ):
@@ -139,6 +140,7 @@ class FastLanguageModel(FastLlamaModel):
             fix_tokenizer  = fix_tokenizer,
             model_patcher  = dispatch_model,
             tokenizer_name = tokenizer_name,
+            trust_remote_code = trust_remote_code,
             *args, **kwargs,
         )
 
diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index c1e39e4a2e..2db71f0288 100644
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@@ -230,6 +230,7 @@ def MistralForCausalLM_fast_forward(
     else:
         logits = self.lm_head(hidden_states)
     pass
+    logits = logits.to(self.config.torch_dtype)
 
     loss = None
     if labels is not None:
@@ -295,6 +296,7 @@ class FastMistralModel(FastLlamaModel):
         fix_tokenizer  = True,
         model_patcher  = None,
         tokenizer_name = None,
+        trust_remote_code = False,
         **kwargs,
     ):
         if model_patcher is None: model_patcher = FastMistralModel
@@ -353,6 +355,7 @@ class FastMistralModel(FastLlamaModel):
             quantization_config = bnb_config,
             token               = token,
             # rope_scaling      = rope_scaling,
+            trust_remote_code   = trust_remote_code,
             **kwargs,
         )
 
@@ -360,9 +363,10 @@ class FastMistralModel(FastLlamaModel):
         tokenizer_name = model_name if tokenizer_name is None else tokenizer_name
         tokenizer = AutoTokenizer.from_pretrained(
             tokenizer_name,
-            model_max_length = max_position_embeddings,
-            padding_side     = "right",
-            token            = token,
+            model_max_length  = max_position_embeddings,
+            padding_side      = "right",
+            token             = token,
+            trust_remote_code = trust_remote_code,
         )
 
         model, tokenizer = patch_tokenizer(model, tokenizer)