| """Render SFT examples with Dory's own chat template; mask loss to the final assistant turn.""" | |
| import torch | |
| def encode_example(tokenizer, messages, medium_effort=False, max_len=None): | |
| """Return (input_ids, labels) 1-D LongTensors. Labels are -100 on the prompt. | |
| The prompt is rendered with add_generation_prompt=True, which ends in | |
| '<|im_start|>assistant\\n<think>\\n'; the full conversation renders the final | |
| assistant turn as '<think>\\n{reasoning}</think>{content}<|im_end|>'. | |
| """ | |
| kw = {"medium_effort": medium_effort} | |
| prompt = tokenizer.apply_chat_template(messages[:-1], tokenize=False, add_generation_prompt=True, **kw) | |
| full = tokenizer.apply_chat_template(messages, tokenize=False, **kw) | |
| if not full.startswith(prompt): | |
| raise ValueError("chat template: full conversation does not extend the generation prompt") | |
| p_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"] | |
| c_ids = tokenizer(full[len(prompt) :], add_special_tokens=False)["input_ids"] | |
| ids = p_ids + c_ids | |
| labels = [-100] * len(p_ids) + c_ids | |
| if max_len and len(ids) > max_len: | |
| return None | |
| return torch.tensor(ids), torch.tensor(labels) | |