"""用极小词表演示编码、padding、因果 mask 与标签右移。"""
import numpy as np

VOCAB = {"<pad>": 0, "<bos>": 1, "<eos>": 2, "我": 3, "爱": 4, "猫": 5, "。": 6}


def encode(tokens: list[str]) -> list[int]:
    return [VOCAB["<bos>"], *(VOCAB[t] for t in tokens), VOCAB["<eos>"]]


def collate(sequences: list[list[int]]):
    max_len = max(map(len, sequences))
    input_ids = np.full((len(sequences), max_len), VOCAB["<pad>"], dtype=np.int64)
    attention_mask = np.zeros_like(input_ids)
    for row, seq in enumerate(sequences):
        input_ids[row, : len(seq)] = seq
        attention_mask[row, : len(seq)] = 1
    labels = np.full_like(input_ids, -100)
    labels[:, :-1] = np.where(attention_mask[:, 1:] == 1, input_ids[:, 1:], -100)
    causal = np.tril(np.ones((max_len, max_len), dtype=bool))
    return input_ids, attention_mask, labels, causal


if __name__ == "__main__":
    batch = [encode(["我", "爱", "猫", "。"]), encode(["猫", "。"]) ]
    ids, mask, labels, causal = collate(batch)
    print("input_ids\n", ids)
    print("labels\n", labels)
    assert ids.shape == labels.shape == mask.shape
    assert np.all(labels[mask == 0] == -100)
    assert np.all(causal == np.tril(causal))