finetune/nemo/warden_processor.py

"""Convert our `{"messages": [...]}` chat rows into the (input, output) pair
the Megatron-Bridge SFT builder expects.

The builder combines them as `"{input} {output}"` + EOS and masks the loss to
the `output` span (sft.py: prompt_template, add_eos=True). So:

  input  = the conversation up to (but not including) the final assistant turn,
           rendered through the model's own chat template with a generation
           prompt appended  -> exactly what the game feeds the model at play time.
  output = the final assistant turn's content (the Warden's line, or a tool-call
           JSON string for director examples).

Rendering through the real chat template keeps train == play: the live game
prompts the GGUF with the same Jinja template via llama.cpp.
"""
from typing import Any, Optional

from megatron.bridge.data.builders.hf_dataset import ProcessExampleOutput
from megatron.bridge.training.tokenizers.tokenizer import MegatronTokenizer


def _hf_tokenizer(tokenizer):
    """Reach the underlying HF tokenizer that owns apply_chat_template."""
    for attr in ("apply_chat_template",):
        if hasattr(tokenizer, attr):
            return tokenizer
    for attr in ("_tokenizer", "tokenizer", "hf_tokenizer"):
        inner = getattr(tokenizer, attr, None)
        if inner is not None and hasattr(inner, "apply_chat_template"):
            return inner
    raise RuntimeError("could not find an apply_chat_template-capable tokenizer")


def process_warden_example(
    example: dict[str, Any], tokenizer: Optional[MegatronTokenizer] = None
) -> ProcessExampleOutput:
    messages = example["messages"]
    if not messages or messages[-1]["role"] != "assistant":
        raise ValueError("every training row must end with an assistant turn")

    prompt_messages = messages[:-1]
    target = messages[-1]["content"]

    tok = _hf_tokenizer(tokenizer)
    _input = tok.apply_chat_template(
        prompt_messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    # The builder inserts a single space before {output}; trim a trailing space
    # on the rendered prompt so we don't double it.
    _input = _input.rstrip(" ")

    return ProcessExampleOutput(input=_input, output=target, original_answers=[target])