File size: 1,193 Bytes
f0975ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 | """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)
|