From 784f8c0e952ff506e5215119de2aa38f70f1ddb7 Mon Sep 17 00:00:00 2001 From: VED <146507396+ved1beta@users.noreply.github.com> Date: Wed, 6 Aug 2025 17:32:07 +0530 Subject: [PATCH] fix:kd_distillation key_error logprobs (#2990) * fix:kd_distillation key_error logprobs * style * fix: leave handling of pop logprobs to parent --------- Co-authored-by: NanoCode012 --- src/axolotl/integrations/kd/chat_template.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/axolotl/integrations/kd/chat_template.py b/src/axolotl/integrations/kd/chat_template.py index f99dfe458..6376ecb09 100644 --- a/src/axolotl/integrations/kd/chat_template.py +++ b/src/axolotl/integrations/kd/chat_template.py @@ -284,12 +284,12 @@ class ChatTemplateStrategyWithKDv2(ChatTemplateStrategyWithKD): return sample def _tokenize_single_prompt(self, prompt): - logprobs = prompt.pop(self.logprobs_field) - target_token_ids = prompt.pop("target_token_ids") + target_token_ids = prompt.get("target_token_ids", None) + tokenized_prompt = super()._tokenize_single_prompt(prompt) - tokenized_prompt[self.logprobs_field] = logprobs - tokenized_prompt["target_token_ids"] = target_token_ids - tokenized_prompt = self.transform_logprobs(tokenized_prompt) + + if target_token_ids is not None: + tokenized_prompt["target_token_ids"] = target_token_ids return tokenized_prompt