大语言模型怎样生成文本：从 logits 到下一个 Token

softmax(logits / temperature)

import numpy as np


TOKENS = np.array(["正常", "异常", "恢复", "继续", "<eos>"])
LOGITS = np.array([2.4, 1.7, 1.2, 0.8, 0.2], dtype=float)


def softmax(values: np.ndarray) -> np.ndarray:
    shifted = values - values.max()
    exp = np.exp(shifted)
    return exp / exp.sum()


def filter_logits(
    logits: np.ndarray,
    top_k: int | None = None,
    top_p: float | None = None,
) -> np.ndarray:
    filtered = logits.copy()

    if top_k is not None:
        keep = np.argsort(filtered)[-top_k:]
        mask = np.ones_like(filtered, dtype=bool)
        mask[keep] = False
        filtered[mask] = -np.inf

    if top_p is not None:
        order = np.argsort(filtered)[::-1]
        ordered_probabilities = softmax(filtered[order])
        cumulative = np.cumsum(ordered_probabilities)
        remove = cumulative - ordered_probabilities >= top_p
        filtered[order[remove]] = -np.inf

    return filtered


def sample_next(
    logits: np.ndarray,
    temperature: float = 1.0,
    top_k: int | None = None,
    top_p: float | None = None,
    seed: int = 7,
):
    if temperature <= 0:
        raise ValueError("temperature must be greater than 0")

    filtered = filter_logits(logits / temperature, top_k, top_p)
    probabilities = softmax(filtered)
    rng = np.random.default_rng(seed)
    index = rng.choice(len(TOKENS), p=probabilities)
    return TOKENS[index], dict(zip(TOKENS, probabilities.round(3)))


for config in [
    {"temperature": 0.5},
    {"temperature": 1.2},
    {"temperature": 1.0, "top_k": 2},
    {"temperature": 1.0, "top_p": 0.75},
]:
    token, probabilities = sample_next(LOGITS, **config)
    print(config, "->", token, probabilities)