"""把 chat 消息编码成 token，并只监督 assistant 区域。"""
from dataclasses import dataclass


@dataclass
class Encoded:
    input_ids: list[int]
    labels: list[int]


def encode_sft(messages, encode_text, role_tokens: dict[str, int], eos_id: int) -> Encoded:
    ids: list[int] = []
    labels: list[int] = []
    for message in messages:
        role = message["role"]
        content_ids = encode_text(message["content"])
        piece = [role_tokens[role], *content_ids, eos_id]
        ids.extend(piece)
        if role == "assistant":
            labels.extend([-100, *content_ids, eos_id])
        else:
            labels.extend([-100] * len(piece))
    if len(ids) != len(labels) or all(label == -100 for label in labels):
        raise ValueError("invalid SFT labels")
    return Encoded(ids, labels)


if __name__ == "__main__":
    result = encode_sft(
        [{"role": "user", "content": "ab"}, {"role": "assistant", "content": "cd"}],
        lambda text: [ord(char) for char in text],
        {"user": 1, "assistant": 2},
        eos_id=3,
    )
    print(result)
    assert result.labels[:4] == [-100] * 4
    assert result.labels[-1] == 3