次トークンの確率と逐次生成Next Token Prediction

from typing import Callable, Optional
import torch
from torch import Tensor
 
 
def select_next_token(
    probabilities: Tensor,      # [B, V]
    top_k: Optional[int] = None,
    top_p: Optional[float] = None,
    greedy: bool = False,
) -> Tensor:
    """確率分布から次の 1 トークンを選ぶ。
 
    Args:
        probabilities: [B, V]。合計が 1 の確率分布。
        top_k: 上位 k 個に絞る。None なら絞らない。
        top_p: 累積確率が p に達するまでに絞る。None なら絞らない。
        greedy: True なら最大確率のトークンを決定的に選ぶ。
 
    Returns:
        [B, 1] int64。選ばれたトークンID。
    """
    if greedy:
        return probabilities.argmax(dim=-1, keepdim=True)
 
    probs = probabilities.clone()
 
    if top_k is not None:
        threshold = probs.topk(top_k, dim=-1).values[:, -1:]
        probs = torch.where(probs < threshold, torch.zeros_like(probs), probs)
 
    if top_p is not None:
        sorted_probs, sorted_idx = probs.sort(dim=-1, descending=True)
        cumulative = sorted_probs.cumsum(dim=-1)
        # 累積が p を超えた「後ろ」を落とす（先頭1つは必ず残す）
        remove = cumulative - sorted_probs > top_p
        sorted_probs[remove] = 0.0
        probs = torch.zeros_like(probs).scatter(-1, sorted_idx, sorted_probs)
 
    probs = probs / probs.sum(dim=-1, keepdim=True)
    return torch.multinomial(probs, num_samples=1)   # [B, 1]
 
 
def generate(
    model: Callable[[Tensor], Tensor],   # [B, T] -> logits [B, T, V]
    prompt_ids: Tensor,                  # [B, T0]
    max_new_tokens: int,
    eos_id: int,
    temperature: float = 1.0,
    top_p: Optional[float] = 0.9,
) -> Tensor:
    """自己回帰的に系列を生成する。
 
    Returns:
        [B, T0 + n] int64。プロンプトを含む生成結果。
    """
    generated = prompt_ids
    for _ in range(max_new_tokens):
        logits = model(generated)                          # [B, T, V]
        last_logits = logits[:, -1, :] / temperature       # [B, V]
        probabilities = last_logits.softmax(dim=-1)        # [B, V]
 
        next_token_id = select_next_token(probabilities, top_p=top_p)  # [B, 1]
        generated = torch.cat([generated, next_token_id], dim=1)       # [B, T+1]
 
        if bool((next_token_id == eos_id).all()):
            break
 
    return generated
 
 
# 最小使用例（model はダミー。are 相当のIDを高確率で返す）
V = 6
ARE_ID, EOS_ID = 0, 5
 
 
def dummy_model(ids: Tensor) -> Tensor:
    batch, length = ids.shape
    logits = torch.full((batch, length, V), -5.0)
    logits[:, :, ARE_ID] = 8.0
    return logits
 
 
prompt = torch.tensor([[1, 2, 3, 4]])          # [B=1, T0=4]
out = generate(dummy_model, prompt, max_new_tokens=3, eos_id=EOS_ID, top_p=0.9)
 
print(prompt.shape, out.shape)                  # torch.Size([1, 4]) torch.Size([1, 7])
print(out[0, 4:].tolist())                      # [0, 0, 0]  （are が繰り返し選ばれる）