How large language models work

import numpy as np

# A made-up model's scores (logits) for the next word after:
# "Metformin is usually first-line treatment for type 2 ..."
candidates = ["diabetes", "patients", "adults", "disease", "hypertension"]
logits = np.array([6.0, 3.5, 2.8, 2.0, 0.5])

def softmax(z, temperature=1.0):
    z = z / temperature
    e = np.exp(z - z.max())
    return e / e.sum()

for t in [0.2, 1.0, 2.0]:
    probs = softmax(logits, t)
    print(f"temperature {t}: " + ", ".join(f"{w} {p:.2f}" for w, p in zip(candidates, probs)))

# temperature 0.2: diabetes 1.00, patients 0.00, adults 0.00, disease 0.00, hypertension 0.00
# temperature 1.0: diabetes 0.87, patients 0.07, adults 0.04, disease 0.02, hypertension 0.00
# temperature 2.0: diabetes 0.59, patients 0.17, adults 0.12, disease 0.08, hypertension 0.04

import numpy as np
from collections import Counter, defaultdict

corpus = """the patient was started on metformin .
the patient was started on insulin .
the patient was referred to cardiology .
the trial was started in 2019 .
the trial was referred to the ethics committee .""".split("\n")

following = defaultdict(Counter)
for sentence in corpus:
    words = sentence.split()
    for a, b in zip(words, words[1:]):
        following[a][b] += 1

rng = np.random.default_rng(4)
def generate(start="the", max_words=12):
    words = [start]
    while words[-1] != "." and len(words) < max_words:
        options = following[words[-1]]
        probs = np.array(list(options.values())) / sum(options.values())
        words.append(rng.choice(list(options), p=probs))
    return " ".join(words)

for _ in range(5):
    print(generate())

# the ethics committee .
# the patient was started in 2019 .
# the trial was started on insulin .
# the patient was referred to the trial was referred to cardiology .
# the trial was started on insulin .