Lesson 2: What pretraining gave the model, and what it costs

%pip install -q tokenizers==0.23.3

ARTICLES = {
    "account-login": ("Sign-in and passwords", "Reset your password from the sign-in page with Forgot password. The reset link is emailed and expires after 30 minutes. After 5 failed sign-in attempts the account is locked for 15 minutes and then unlocks by itself. If you lost your two-factor device, a workspace owner can reset two-factor for you."),
    "account-sso": ("Single sign-on", "SSO with SAML or OIDC is available on Business and Enterprise plans. A workspace owner sets it up in Settings > Security. If people are sent back to the sign-in page after signing in, check that your identity provider's certificate has not expired and that the callback URL matches the one in Settings."),
    "billing-invoices": ("Invoices and tax details", "Invoices are emailed to the billing contact after every charge and can be downloaded as PDF from Settings > Billing > Invoices. Add a VAT or tax number in Settings > Billing > Tax details; it appears on future invoices. Support can reissue invoices from the last 12 months."),
    "billing-plans": ("Plans and seats", "Free: up to 3 users and 5 boards. Team: 12 USD per member per month, or 10 USD billed annually. Business: 24 USD per member per month, or 20 USD billed annually, and adds SSO and audit logs. Enterprise: custom pricing, contact sales. A member added mid-period is charged pro rata on the next invoice. Switching from monthly to annual starts a new annual period that day."),
    "billing-refunds": ("Refunds and cancellations", "Cancel at any time from Settings > Billing > Cancel plan. Monthly plans stay active until the end of the period and are not refunded pro rata. Annual plans cancelled within 14 days of purchase or renewal get a full refund; after 14 days, no refund. Duplicate charges are always refunded in full. Only workspace owners and billing admins can cancel or request refunds. Approved refunds reach the original card within 5 to 10 working days."),
    "boards-automations": ("Automations and recurring cards", "An automation runs a rule when something happens on a board, such as moving a card when its checklist is complete. Team plans include 1,000 automation runs a month and Business plans 10,000. When the limit is reached, automations pause until the next month and the board shows a banner. A recurring card is an automation on a schedule, such as every Monday."),
    "data-privacy": ("Privacy, deletion and the DPA", "Delete your personal account from Profile > Delete account. A deleted workspace can be restored by its owner for 30 days; after that it is purged and cannot be recovered. Owners can download our data processing agreement (DPA) and subprocessor list from Settings > Legal."),
    "exports-data": ("Exporting your data", "Owners and admins can export every board, card and comment as CSV or JSON from Settings > Data > Export. Large exports are prepared in the background and emailed as a link that works for 7 days. One board can be exported to CSV, which opens in Excel, from the board menu."),
    "feature-requests": ("Feature requests", "Post ideas on the public roadmap at Settings > Help > Ideas, where other customers can vote. We read every idea but do not give release dates. Timeline and Gantt views are not available today."),
    "integrations-slack": ("Slack integration", "Connect Slack from Settings > Integrations > Slack, then choose which boards post to which channel. Each connection posts once per update. If updates appear twice, the same board is connected to that channel twice: remove one connection."),
    "mobile-app": ("Mobile apps", "The iOS and Android apps need iOS 17 or Android 12 or newer. If the app closes on start, update it from the store and clear its cache. Report a crash from Profile > Help in the app so the log is attached."),
    "status-incidents": ("Service status and incidents", "Incidents are posted at status.quillpath.example within 15 minutes of detection and updated at least every hour until resolved. Urgent tickets about an incident on Business and Enterprise plans get a first reply within 1 hour, at any time."),
}

words = {aid: len(text.split()) for aid, (title, text) in ARTICLES.items()}
print(len(ARTICLES), "articles,", sum(words.values()), "words; longest:", max(words, key=words.get), words[max(words, key=words.get)])

import random

QA = [   # (customer question, the agent's fixed answer, taken from the help centre)
    ("how much does the Team plan cost?", "Team costs 12 USD per member per month, or 10 USD billed annually."),
    ("how much does the Business plan cost?", "Business costs 24 USD per member per month, or 20 USD billed annually."),
    ("how long does a reset link last?", "Reset links expire after 30 minutes."),
    ("why is my account locked?", "After 5 failed sign-in attempts the account is locked for 15 minutes."),
    ("can I get a refund on my monthly plan?", "Monthly plans are not refunded pro rata."),
    ("can I get a refund on my annual plan?", "Annual plans cancelled within 14 days of purchase or renewal get a full refund."),
    ("which plans have SSO?", "SSO is available on Business and Enterprise plans."),
    ("can I restore a deleted workspace?", "A deleted workspace can be restored by its owner for 30 days."),
]
NAMES, AGENTS = ["Ana", "Leo", "Grace", "Jonas", "Yuki", "Ruth", "Kofi", "Sara"], ["Lena", "Ravi"]
GREETINGS = ["Hi,", "Hello,", "Good morning,", "Quick question:"]

def build_corpus(n=3000, seed=7):
    rng = random.Random(seed)
    lines = [f"{title}\n{text}\n" for title, text in ARTICLES.values()]   # each article once
    for _ in range(n):
        q, a = rng.choice(QA)
        lines.append(f"Customer ({rng.choice(NAMES)}): {rng.choice(GREETINGS)} {q}\nAgent ({rng.choice(AGENTS)}): {a}\n")
    return "\n".join(lines)

CORPUS = build_corpus()
print(f"{len(CORPUS):,} characters, {len(CORPUS.split()):,} words")
print(CORPUS[CORPUS.index("Customer"):].split("\n\n")[0])

from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders

tok = Tokenizer(models.BPE())
tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
tok.decoder = decoders.ByteLevel()
trainer = trainers.BpeTrainer(vocab_size=2048, initial_alphabet=pre_tokenizers.ByteLevel.alphabet())
tok.train_from_iterator([CORPUS], trainer)

ids = tok.encode(CORPUS).ids
print(tok.get_vocab_size(), "tokens learned;", f"{len(ids):,}", "tokens in the corpus")
print(tok.encode("Agent (Lena): Annual plans cancelled within 14 days").tokens)

import torch, torch.nn as nn, torch.nn.functional as F

class Block(nn.Module):
    """One transformer layer: attention to earlier tokens, then a small network."""
    def __init__(self, d=128, heads=4):
        super().__init__()
        self.heads = heads
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.qkv, self.proj = nn.Linear(d, 3 * d), nn.Linear(d, d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        b, t, d = x.shape
        q, k, v = self.qkv(self.ln1(x)).split(d, dim=2)
        q, k, v = (z.view(b, t, self.heads, d // self.heads).transpose(1, 2) for z in (q, k, v))
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)   # look left only
        x = x + self.proj(y.transpose(1, 2).reshape(b, t, d))
        return x + self.mlp(self.ln2(x))

class TinyLM(nn.Module):
    def __init__(self, vocab=2048, context=128, d=128, layers=4, heads=4):
        super().__init__()
        self.context = context
        self.tok_emb = nn.Embedding(vocab, d)          # token id -> 128 numbers
        self.pos_emb = nn.Embedding(context, d)        # position -> 128 numbers
        self.blocks = nn.Sequential(*[Block(d, heads) for _ in range(layers)])
        self.ln_f = nn.LayerNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)    # 128 numbers -> one score per token
        self.head.weight = self.tok_emb.weight         # tied: input and output share one table
        for m in self.modules():
            if isinstance(m, (nn.Linear, nn.Embedding)):
                nn.init.normal_(m.weight, std=0.02)    # small random start

    def forward(self, idx):
        pos = torch.arange(idx.shape[1], device=idx.device)
        x = self.blocks(self.tok_emb(idx) + self.pos_emb(pos))
        return self.head(self.ln_f(x))                 # scores: (batch, tokens, vocab)

torch.manual_seed(0)
model = TinyLM()
print(f"{sum(p.numel() for p in model.parameters()):,} parameters")

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
data = torch.tensor(ids)
cut = int(len(data) * 0.9)
train, val = data[:cut], data[cut:]                 # the last 10% is never trained on

@torch.no_grad()
def avg_loss(tokens):
    """Average next-token loss over a stretch of text, in 128-token windows."""
    model.eval()
    total, count = 0.0, 0
    for s in range(0, len(tokens) - 1, 128):
        x, y = tokens[s:s + 128], tokens[s + 1:s + 129]
        logits = model(x[:len(y)].unsqueeze(0).to(device))[0]
        total += F.cross_entropy(logits, y.to(device), reduction="sum").item()
        count += len(y)
    model.train()
    return total / count

def train_steps(text_ids, steps, every, check):
    opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
    g = torch.Generator().manual_seed(0)
    log = []
    for step in range(1, steps + 1):
        starts = torch.randint(0, len(text_ids) - 129, (16,), generator=g)
        x = torch.stack([text_ids[s:s + 128] for s in starts]).to(device)
        y = torch.stack([text_ids[s + 1:s + 129] for s in starts]).to(device)
        logits = model(x)
        loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1))
        opt.zero_grad(); loss.backward(); opt.step()
        if step == 1 or step % every == 0:
            log.append((step, avg_loss(check[0]), avg_loss(check[1])))
    return log

log = train_steps(train, steps=1000, every=100, check=(train[-len(val):], val))
print(f"{len(train):,} training tokens, {len(val):,} validation tokens")
print("step  train   val")
for step, tr, va in log:
    print(f"{step:4}  {tr:.3f}  {va:.3f}")

import math
step, tr, va = log[-1]
print(f"blind guess over 2,048 slots: loss {math.log(2048):.2f}")
print(f"step {step}: validation loss {va:.3f} = perplexity {math.exp(va):.2f}")
print(f"chance of the right next token, on average: {math.exp(-va):.2f}")
print(f"tokens trained on: {1000 * 16 * 128:,} = {1000 * 16 * 128 / len(train):.0f} passes over the training text")

held = "boards-automations"
small_train = "\n".join(f"{t}\n{x}\n" for a, (t, x) in ARTICLES.items() if a != held)
small_val = f"{ARTICLES[held][0]}\n{ARTICLES[held][1]}\n"
small_train = torch.tensor(tok.encode(small_train).ids)
small_val = torch.tensor(tok.encode(small_val).ids)

trained = {k: v.clone() for k, v in model.state_dict().items()}   # keep the good model
torch.manual_seed(0)
model = TinyLM().to(device)                                        # a fresh, untrained copy
over = train_steps(small_train, steps=300, every=25, check=(small_train, small_val))
best = min(over, key=lambda row: row[2])
for step, tr, va in over:
    print(f"{step:4}  train {tr:6.3f}  val {va:6.3f}" + ("  <- best" if (step, tr, va) == best else ""))
model.load_state_dict(trained)                                     # back to the full-corpus model

@torch.no_grad()
def greedy(prompt, max_tokens=30):
    ids_ = tok.encode(prompt).ids
    start = len(ids_)
    for _ in range(max_tokens):
        probs = torch.softmax(model(torch.tensor([ids_[-128:]], device=device))[0, -1], dim=-1)
        best = int(probs.argmax())
        if tok.decode([best]) == "\n":
            break
        ids_.append(best)
    return tok.decode(ids_[start:])

model.eval()
answers = sum(greedy(f"Customer (Ana): Hi, {q}\nAgent (Lena):").strip() == a for q, a in QA)
print(f"Q&A answers repeated word for word: {answers} of {len(QA)}")

articles = 0
for title, text in ARTICLES.values():
    words = text.split()
    prompt = title + "\n" + " ".join(words[:8])                     # the first 8 words
    articles += greedy(prompt).strip().startswith(" ".join(words[8:14]))   # the next 6?
print(f"articles continued word for word: {articles} of {len(ARTICLES)}")
print("each Q&A answer appears", CORPUS.count(QA[0][1]), "times; '1,000 automation runs' appears",
      CORPUS.count("1,000 automation runs"), "time")

for question in ["how many automation runs does the Team plan include?",
                 "how much does the Enterprise plan cost?"]:
    prompt = f"Customer (Ana): Hi, {question}\nAgent (Lena):"
    with torch.no_grad():
        probs = torch.softmax(model(torch.tensor([tok.encode(prompt).ids], device=device))[0, -1], dim=-1)
    top_p, top_i = probs.topk(3)
    print(question)
    print("   first token:", [(tok.decode([i]), round(p, 3)) for p, i in zip(top_p.tolist(), top_i.tolist())])
    print("   reply:", greedy(prompt).strip())
print("probabilities add up to", round(probs.sum().item(), 4))

def expected_score(p_right, wrong_costs):
    """Score of guessing when the guess is right with probability p_right. Abstaining scores 0."""
    return p_right * 1 - (1 - p_right) * wrong_costs

for p in (0.1, 0.5, 0.8):
    binary, penalty = expected_score(p, 0), expected_score(p, 3)
    print(f"p(right) {p:.1f}: guess scores {binary:+.2f} under binary grading, "
          f"{penalty:+.2f} when a wrong answer costs 3; abstain scores +0.00")

%pip install -q -U langchain langgraph langchain-groq langchain-google-genai langchain-ollama

import os
from langchain.chat_models import init_chat_model

PROVIDER = "groq"   # "groq" or "gemini" in Colab; "ollama" on your own computer
MODELS = {"groq": "groq:openai/gpt-oss-20b", "gemini": "google_genai:gemini-3.5-flash-lite", "ollama": "ollama:llama3.2:3b"}
KEY = {"groq": "GROQ_API_KEY", "gemini": "GOOGLE_API_KEY", "ollama": None}[PROVIDER]
if KEY and KEY not in os.environ:
    from google.colab import userdata   # keys live in Colab's Secrets panel
    os.environ[KEY] = userdata.get(KEY)
llm = init_chat_model(MODELS[PROVIDER], temperature=0, max_retries=6)

def ask(model, messages):
    """Returns (reply, None), or (None, reason) after two refusals."""
    for attempt in range(2):                     # one retry: the next reply is usually fine
        try:
            return model.invoke(messages), None
        except Exception as err:
            last = err
    return None, f"provider_error: {type(last).__name__}"

import re
question = "How many automation runs a month does Quillpath's Team plan include?"
RULES = {
    "right 1, else 0": "Reply with a number only. A right number scores 1 point; anything else, including no answer, scores 0.",
    "wrong costs 3":   "Reply with a number, or with I don't know. A right number scores 1, I don't know scores 0, a wrong number scores -3.",
}
for name, rule in RULES.items():
    answers = [ask(llm, f"{question} {rule}") for _ in range(3)]
    replies = [reply.text.strip() for reply, err in answers if reply is not None]
    guesses = [r for r in replies if re.fullmatch(r"[\d,]+", r)]
    print(f"{name:16} guessed a number {len(guesses)} of {len(replies)} times: {guesses}")
    for reply, err in answers:
        if reply is None:
            print(f"{'':16} no answer: {err}")

from collections import Counter

parts = Counter()
for name, p in model.named_parameters():
    part = ("embeddings" if "emb" in name else "attention" if "qkv" in name or "proj" in name
            else "feed-forward" if "mlp" in name else "layer norms")
    parts[part] += p.numel()
total = sum(parts.values())
for part, n in parts.most_common():
    print(f"{part:13} {n:>9,}  {n / total:.0%}")
print(f"{'total':13} {total:>9,}   vocabulary slots 2,048, used {tok.get_vocab_size():,}")

import json
from huggingface_hub import hf_hub_download, get_safetensors_metadata

def count_from_config(c):
    """Parameters of a Llama-style decoder, from the numbers in its config.json."""
    d, layers, ffn, vocab = c["hidden_size"], c["num_hidden_layers"], c["intermediate_size"], c["vocab_size"]
    head = c.get("head_dim") or d // c["num_attention_heads"]
    attn = d * c["num_attention_heads"] * head * 2 + d * c["num_key_value_heads"] * head * 2   # Q, O, then K, V
    mlp = 3 * d * ffn                                         # gate, up and down matrices
    norms = 2 * d                                             # two norms per layer
    embed = vocab * d * (1 if c.get("tie_word_embeddings") else 2)
    return layers * (attn + mlp + norms) + embed + d          # + the final norm

CONFIGS = {}
for repo in ["HuggingFaceTB/SmolLM2-135M-Instruct", "unsloth/Llama-3.2-3B-Instruct"]:
    CONFIGS[repo] = c = json.load(open(hf_hub_download(repo, "config.json")))
    stored = sum(get_safetensors_metadata(repo).parameter_count.values())
    print(f"{repo.split('/')[1]:24} from config {count_from_config(c):>13,}   in the weight files {stored:>13,}")

BYTES = {"fp32": 4, "bf16": 2, "int8": 1, "4-bit": 0.5}     # bytes per stored number
SIZES = {"SmolLM2-135M": 134_515_008, "Llama 3.2 3B": 3_212_749_824, "gpt-oss-20b": 20.91e9}
print(f"{'weights, GB':14}" + "".join(f"{k:>7}" for k in BYTES))
for name, n in SIZES.items():
    print(f"{name:14}" + "".join(f"{n * b / 1e9:7.2f}" for b in BYTES.values()))
print(f"TinyLM at fp32: {total * 4 / 1e6:.1f} MB")

# gpt-oss-20b ships its 19.12B expert weights at 4.25 bits (MXFP4), the other 1.79B at 2 bytes
shipped = 19.12e9 * 4.25 / 8 + (20.91e9 - 19.12e9) * 2
print(f"gpt-oss-20b as shipped, estimated: {shipped / 1e9:.1f} GB = {shipped / 2**30:.1f} GiB")

def kv_bytes_per_token(c, bytes_per_number=2):
    head = c.get("head_dim") or c["hidden_size"] // c["num_attention_heads"]
    return 2 * c["num_hidden_layers"] * c["num_key_value_heads"] * head * bytes_per_number   # K and V

for repo, c in CONFIGS.items():
    per = kv_bytes_per_token(c)
    window = c["max_position_embeddings"]
    print(f"{repo.split('/')[1]:24} {per:,} bytes per token; "
          f"4,096 tokens {per * 4096 / 1e9:.2f} GB; full {window:,} tokens {per * window / 1e9:.1f} GB")
print("TinyLM (fp32):", f"{2 * 4 * 4 * 32 * 4:,} bytes per token, {2 * 4 * 4 * 32 * 4 * 128:,} bytes for 128 tokens")

repo = "Qwen/Qwen3-8B"
c = json.load(open(hf_hub_download(repo, "config.json")))
mine = count_from_config(c)
stored = sum(get_safetensors_metadata(repo).parameter_count.values())
print(f"from config {mine:,}   in the weight files {stored:,}")
assert abs(mine - stored) / stored < 0.001, "more than 0.1% apart: check tied embeddings and head_dim"
print(f"bf16 weights {stored * 2 / 1e9:.1f} GB; KV cache at 32,768 tokens {kv_bytes_per_token(c) * 32768 / 1e9:.1f} GB")