From 2bb066008bb5cda2c697197dfcf45caa3852df6c Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 25 Nov 2024 21:31:43 -0800
Subject: [PATCH] logits

---
 unsloth/models/_utils.py | 25 +++++++++++++++++++++++++
 unsloth/models/llama.py  |  4 +---
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 0be3bfcc1d..5326c09b52 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -54,6 +54,7 @@ __all__ = [
     "unpatch_gradient_checkpointing",
 
     "HAS_CUT_CROSS_ENTROPY",
+    "EMPTY_LOGITS",
     "fused_linear_cross_entropy",
     "patch_unsloth_smart_gradient_checkpointing",
     "unpatch_unsloth_smart_gradient_checkpointing",
@@ -1164,3 +1165,27 @@ def unsloth_compile_transformers(
     pass
     return model_types
 pass
+
+# We need an empty logits flag to warn people logits will not be returned anymore unless asked ie
+# os.environ['UNSLOTH_RETURN_LOGITS'] = '1'
+LOGITS_ERROR_STRING = \
+    "Unsloth: Logits are empty from 2024.11 onwards. To get raw logits again, please "\
+    'set the environment variable `UNSLOTH_RETURN_LOGITS` to `"1" BEFORE starting to train ie before `trainer.train()`. For example:\n\n'\
+    "import os\n"\
+    "os.environ['UNSLOTH_RETURN_LOGITS'] = '1'\n"\
+    "... trainer.train() ..."
+
+def raise_logits_error(*args, **kwargs): raise NotImplementedError(LOGITS_ERROR_STRING)
+class EmptyLogits(torch.Tensor):
+    def __init__(self): return
+    __getitem__ = raise_logits_error
+    __getattr__ = raise_logits_error
+    def __repr__(self): return LOGITS_ERROR_STRING
+    def __str__ (self): return LOGITS_ERROR_STRING
+pass
+EMPTY_LOGITS = EmptyLogits()
+functions = dir(torch.Tensor)
+for function in functions:
+    try: exec(f"EMPTY_LOGITS.{function} = raise_logits_error", globals(), locals())
+    except: continue
+pass
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 3ec59b339a..a17af2dd8a 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -974,8 +974,6 @@ def CausalLM_fast_forward(fast_forward_inference):
         logit_softcapping = getattr(self.config, "final_logit_softcapping", 0)
         logit_scaling     = getattr(self.config, "logit_scale", 0)
 
-        print(kwargs)
-
         if bsz == 1 and q_len == 1:
             logits = torch.mv(lm_head, hidden_states.ravel().to(lm_head.dtype))
             logits = logits.unsqueeze(0).unsqueeze(0)
@@ -997,7 +995,7 @@ def CausalLM_fast_forward(fast_forward_inference):
 
                 return CausalLMOutputWithPast(
                     loss=loss,
-                    logits=None,
+                    logits=EMPTY_LOGITS,
                     past_key_values=outputs.past_key_values,
                     hidden_states=outputs.hidden_states,
                     attentions=outputs.attentions,