generate.py

1# the probabilities are transformed into words23from ops import softmax_forward4from transformer import ModelConfig, forward5import numpy as np67def sample_next(logits, temperature, top_k, top_p, rng):8    if temperature < 1e-6:9        return int(np.argmax(logits))1011    p = softmax_forward(logits / temperature)1213    if 0 < top_k < len(p):14        cutoff = np.sort(p)[-top_k] # # k-th largest probabilit15        p = np.where(p >= cutoff, p, 0.0) # zero everything below it16        p = p / p.sum() # renormalise so it sums to 1 again1718    if top_p < 1.0:19        order = np.argsort(-p)20        cum = np.cumsum(p[order])21        keep = order[: np.searchsorted(cum, top_p) + 1]22        q = np.zeros_like(p)23        q[keep] = p[keep]24        p = q / q.sum()2526    return int(rng.choice(len(p), p=p))2728def generate(params, cfg: ModelConfig, seed_ids, max_new_tokens, window, temperature, top_k, top_p, rng):29    ids = list(seed_ids)3031    for _ in range(max_new_tokens):32        logits, _ = forward(params, np.array(ids[-window:]), cfg)33        ids.append(sample_next(logits[-1], temperature, top_k, top_p, rng)) # logits[-1] is used because the last position is the only one that has seen the whole context3435    return ids