入力トークンInput Tokens

from dataclasses import dataclass
import re
import torch
from torch import Tensor
 
 
@dataclass
class TokenizedBatch:
    input_ids: Tensor       # [B, S] int64
    attention_mask: Tensor  # [B, S] int64（実トークン=1, パディング=0）
    tokens: list[list[str]] # 説明用。実装では通常返さない
 
 
def tokenize(text: str) -> list[str]:
    """説明用に文章をトークンへ分割する。
 
    実運用の Tokenizer はサブワード単位に分割するため、
    同じ文でも分割結果と系列長は異なる。
    """
    return re.findall(r"\w+|[^\w\s]", text)
 
 
def encode(
    texts: list[str],
    vocab: dict[str, int],
    unk_id: int = 0,
) -> TokenizedBatch:
    """文字列のリストを [B, S] の Token ID へ変換する。
 
    Args:
        texts: 入力文のリスト。長さが B になる。
        vocab: トークン文字列 -> ID の語彙表。
        unk_id: 語彙表にないトークンへ割り当てるID。
 
    Returns:
        input_ids [B, S] と attention_mask [B, S] を持つバッチ。
    """
    batch_tokens = [tokenize(text) for text in texts]
    max_len = max(len(tokens) for tokens in batch_tokens)
 
    ids = [
        [vocab.get(token, unk_id) for token in tokens] + [0] * (max_len - len(tokens))
        for tokens in batch_tokens
    ]
    mask = [
        [1] * len(tokens) + [0] * (max_len - len(tokens))
        for tokens in batch_tokens
    ]
 
    return TokenizedBatch(
        input_ids=torch.tensor(ids, dtype=torch.long),      # [B, S]
        attention_mask=torch.tensor(mask, dtype=torch.long),# [B, S]
        tokens=batch_tokens,
    )
 
 
# 最小使用例
vocab = {"The": 5, "keys": 6, "on": 7, "the": 8, "table": 9, "are": 10, "lost": 11, ".": 12}
batch = encode(["The keys on the table are lost."], vocab)
 
print(batch.tokens)            # [['The', 'keys', 'on', 'the', 'table', 'are', 'lost', '.']]
print(batch.input_ids.shape)   # torch.Size([1, 8])  == [B, S]