tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

fix:kd_distillation key_error logprobs (#2990)

Browse Source 
* fix:kd_distillation key_error logprobs

* style

* fix: leave handling of pop logprobs to parent

---------

Co-authored-by: NanoCode012 <nano@axolotl.ai>
This commit is contained in:
VED
2025-08-06 17:32:07 +05:30
committed by GitHub
parent e3177c3210
commit 784f8c0e95
1 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/axolotl/integrations/kd/chat_template.py
View File

@@ -284,12 +284,12 @@ class ChatTemplateStrategyWithKDv2(ChatTemplateStrategyWithKD):
return sample return sample
def _tokenize_single_prompt(self, prompt): def _tokenize_single_prompt(self, prompt):
logprobs = prompt.pop(self.logprobs_field) target_token_ids = prompt.get("target_token_ids", None)
target_token_ids = prompt.pop("target_token_ids")
tokenized_prompt = super()._tokenize_single_prompt(prompt) tokenized_prompt = super()._tokenize_single_prompt(prompt)
tokenized_prompt[self.logprobs_field] = logprobs
tokenized_prompt["target_token_ids"] = target_token_ids if target_token_ids is not None:
tokenized_prompt = self.transform_logprobs(tokenized_prompt) tokenized_prompt["target_token_ids"] = target_token_ids
return tokenized_prompt return tokenized_prompt