From 6327b9fb1ad21b8c556f9aef39841670b04361fc Mon Sep 17 00:00:00 2001
From: Yu Li <lyo.gavin@gmail.com>
Date: Thu, 14 Sep 2023 21:31:48 -0500
Subject: [PATCH] default use xentropy and use only last logit in sampling

---
 anima_100k/longer_training.py      | 1 +
 anima_100k/modeling_flash_llama.py | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/anima_100k/longer_training.py b/anima_100k/longer_training.py
index 64b9f11..38accce 100644
--- a/anima_100k/longer_training.py
+++ b/anima_100k/longer_training.py
@@ -325,6 +325,7 @@ class SampleGenerateCallback(transformers.TrainerCallback):
                     input_ids=input_ids,
                     max_new_tokens=15,
                     do_sample = False,
+                    only_last_logit=True,  # to save memory
                 )
 
 
diff --git a/anima_100k/modeling_flash_llama.py b/anima_100k/modeling_flash_llama.py
index 60ec201..40f803c 100644
--- a/anima_100k/modeling_flash_llama.py
+++ b/anima_100k/modeling_flash_llama.py
@@ -782,7 +782,7 @@ class LlamaForCausalLM(LlamaPreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         only_last_logit: Optional[bool] = None,
-        xentropy: Optional[bool] = None,
+        xentropy: Optional[bool] = True,
         is_padded_inputs: Optional[bool] = None,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
         r"""