ngram.py

"""A small n-gram language model, and the sampling knobs that sit on top of it.

The point is not that this is a good model — it is a lookup table of what
followed what. The point is that every step between "a distribution over next
tokens" and "a word on your screen" is arithmetic you can watch, and those steps
are identical in a real model. Temperature, top-k and top-p do exactly this to
a 200,000-way distribution instead of a twenty-way one.

Mirrored by predict/ngram.js; checkngram.py fails the build if they diverge.
Sampling uses a seeded PRNG (mulberry32) so that a given seed produces the same
text in both implementations, and for the reader, every time.
"""

from bpe import pretokenize


def tokenize(text):
    """Same rule as the BPE page: optional leading space, then non-space."""
    return pretokenize(text)


def train(text, order):
    """context tuple -> {next token: count}, for every context length <= order."""
    tokens = tokenize(text)
    model = {}
    for k in range(order + 1):
        table = {}
        for i in range(len(tokens) - k):
            ctx = tuple(tokens[i:i + k])
            nxt = tokens[i + k]
            table.setdefault(ctx, {})
            table[ctx][nxt] = table[ctx].get(nxt, 0) + 1
        model[k] = table
    # First-appearance order, not sorted: JavaScript orders strings by UTF-16
    # code unit and Python by code point, and they disagree on astral characters.
    vocab, seen = [], set()
    for t in tokens:
        if t not in seen:
            seen.add(t)
            vocab.append(t)
    return {"model": model, "order": order, "tokens": tokens, "vocab": vocab}


def distribution(trained, context):
    """Probabilities for the next token, backing off to shorter contexts.

    Returns (list of (token, probability) sorted by probability desc, order used).
    """
    order = trained["order"]
    ctx = tuple(context)
    for k in range(min(order, len(ctx)), -1, -1):
        table = trained["model"][k]
        key = ctx[len(ctx) - k:] if k else ()
        if key in table:
            counts = table[key]
            total = sum(counts.values())
            items = [(t, c / total) for t, c in counts.items()]
            # Sort by probability, then by first appearance in the vocabulary,
            # so ties resolve identically in both implementations.
            order_index = {t: i for i, t in enumerate(trained["vocab"])}
            items.sort(key=lambda p: (-p[1], order_index[p[0]]))
            return items, k
    return [], 0


def _total(values):
    """Naive left-to-right float accumulation.

    Deliberately not sum(): since 3.12 CPython uses Neumaier compensated
    summation for floats, so sum() returns 1.0 where JavaScript's += loop
    returns 0.9999999999999999. Being more accurate than the browser is still
    being different from it, and that difference moved where top-p cut the
    distribution. Matching the browser is the requirement here.
    """
    total = 0.0
    for v in values:
        total += v
    return total


def apply_temperature(items, temp):
    """p -> p^(1/T), renormalised. T<1 sharpens, T>1 flattens, T→0 is greedy."""
    if temp <= 0:
        # The limit: all mass on the most likely token.
        return [(t, 1.0 if i == 0 else 0.0) for i, (t, _) in enumerate(items)]
    scaled = [(t, p ** (1.0 / temp)) for t, p in items]
    total = _total(p for _, p in scaled)
    if total == 0:
        return items
    return [(t, p / total) for t, p in scaled]


def apply_top_k(items, k):
    """Keep the k most likely tokens, renormalise. 0 disables."""
    if k <= 0 or k >= len(items):
        return items
    kept = items[:k]
    total = _total(p for _, p in kept)
    return [(t, p / total) for t, p in kept]


def apply_top_p(items, p_threshold):
    """Nucleus: keep the smallest set whose probability sums past the threshold."""
    if p_threshold <= 0 or p_threshold >= 1:
        return items
    kept, running = [], 0.0
    for token, p in items:
        kept.append((token, p))
        running += p
        if running >= p_threshold:
            break
    total = _total(p for _, p in kept)
    return [(t, p / total) for t, p in kept]


class Rng:
    """mulberry32 — small, seedable, and identical in JavaScript."""

    def __init__(self, seed):
        self.state = seed & 0xFFFFFFFF

    def next(self):
        self.state = (self.state + 0x6D2B79F5) & 0xFFFFFFFF
        t = self.state
        t = ((t ^ (t >> 15)) * (t | 1)) & 0xFFFFFFFF
        t = (t ^ (t + ((t ^ (t >> 7)) * (t | 61) & 0xFFFFFFFF))) & 0xFFFFFFFF
        return ((t ^ (t >> 14)) & 0xFFFFFFFF) / 4294967296.0


def pick(items, rng):
    """Sample one token from a normalised distribution."""
    if not items:
        return None
    r = rng.next()
    acc = 0.0
    for token, p in items:
        acc += p
        if r < acc:
            return token
    return items[-1][0]


def shape(items, temp=1.0, top_k=0, top_p=0.0):
    """The full pipeline, in the order a real sampler applies it."""
    out = apply_temperature(items, temp)
    out = apply_top_k(out, top_k)
    out = apply_top_p(out, top_p)
    return out


def generate(trained, prompt, n, temp=1.0, top_k=0, top_p=0.0, seed=7):
    """Produce n tokens, returning both the text and a record of each choice."""
    rng = Rng(seed)
    context = list(tokenize(prompt))
    produced, trace = [], []
    for _ in range(n):
        items, used = distribution(trained, context)
        if not items:
            break
        shaped = shape(items, temp, top_k, top_p)
        token = pick(shaped, rng)
        if token is None:
            break
        chosen_p = dict(shaped).get(token, 0.0)
        trace.append({
            "token": token,
            "p": chosen_p,
            "considered": len(shaped),
            "of": len(items),
            "context_order": used,
        })
        produced.append(token)
        context.append(token)
    return {"text": "".join(produced), "tokens": produced, "trace": trace}