"""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)