Lesson 1: A model that only predicts the next token

CATEGORIES = ("billing", "cancellation", "account_access", "bug", "how_to", "feature_request")
PRIORITIES = ("low", "normal", "high", "urgent")

CATEGORY_DEFINITIONS = {
    "billing": "Charges, invoices, prices, plan changes, and refund requests that do not end the plan.",
    "cancellation": "The customer wants to end a paid plan, with or without a refund.",
    "account_access": "Signing in, passwords, two-factor, SSO, and deleting a personal account.",
    "bug": "Something in Quillpath does not work as documented, including slowness and outages.",
    "how_to": "How to do something Quillpath already supports, or where to find a document.",
    "feature_request": "Something Quillpath does not do today that the customer would like.",
}

PRIORITY_DEFINITIONS = {
    "urgent": "A whole team cannot work right now.",
    "high": "One person cannot work, money was taken wrongly, or data is being lost.",
    "normal": "Everything else that needs an answer within a working day.",
    "low": "Ideas and questions with no time pressure, stated or implied.",
}

assert set(CATEGORY_DEFINITIONS) == set(CATEGORIES) and set(PRIORITY_DEFINITIONS) == set(PRIORITIES)
print(len(CATEGORIES), "categories,", len(PRIORITIES), "priorities")

# (id, language, plan, subject, body, category, priority, gold article or None)
DEV = [
    ("T-1001", "en", "team", "Charged twice this month", "Hi, my card was charged 288 USD twice on 1 June for the Team plan (invoice INV-2026-004512). Please refund the duplicate. Ana Duarte, owner, Ardent Design", "billing", "high", "billing-refunds"),
    ("T-1002", "en", "business", "Switch to annual billing?", "We are on Business monthly with 8 seats. If we move to annual, what do we pay and when does it start?", "billing", "normal", "billing-plans"),
    ("T-1003", "en", "team", "VAT number on invoices", "Our accountant needs our VAT number on every invoice. Where do I add it, and can you reissue May's invoice?", "billing", "normal", "billing-invoices"),
    ("T-1004", "en", "team", "Refund for annual plan", "I bought the annual Team plan 20 days ago and we have decided not to use Quillpath. Can I get my money back? Yuki Sato, Studio Nami", "cancellation", "normal", "billing-refunds"),
    ("T-1005", "en", "business", "Cancel our subscription", "Please cancel our Business plan at the end of this month. We are moving to another tool.", "cancellation", "normal", "billing-refunds"),
    ("T-1006", "en", "free", "Forgot my password", "The reset email never arrived and I checked spam. My login is kofi.osei@example.com, phone +44 7700 900321.", "account_access", "high", "account-login"),
    ("T-1007", "en", "business", "SSO login loop", "Since this morning everyone in our team is sent back to the login page after signing in through our identity provider. Nobody can get in.", "account_access", "urgent", "account-sso"),
    ("T-1008", "en", "team", "Locked out", "I typed my password wrong a few times and now it says my account is locked. I have a client demo in an hour.", "account_access", "high", "account-login"),
    ("T-1009", "en", "team", "Add a teammate", "How do I invite a new person to our workspace, and will it change our bill?", "how_to", "low", "billing-plans"),
    ("T-1010", "en", "team", "Automations not firing", "Our rule 'move card to Done when checklist complete' stopped running yesterday. We use a lot of automations. Nothing changed on our side.", "bug", "normal", "boards-automations"),
    ("T-1011", "en", "business", "Export all our boards", "We need a full export of every board with comments for an audit. What formats do you support?", "how_to", "normal", "exports-data"),
    ("T-1012", "en", "business", "Slack posts duplicated", "Every card update now posts twice in our Slack channel.", "bug", "normal", "integrations-slack"),
    ("T-1013", "en", "team", "App crashes on start", "Since Friday's update the Android app closes as soon as I open it. Android 16. I already reinstalled and cleared the cache.", "bug", "normal", None),
    ("T-1014", "en", "free", "Dark mode please", "Would love a dark mode in the web app. My eyes would thank you.", "feature_request", "low", "feature-requests"),
    ("T-1015", "en", "business", "Is Quillpath down?", "Boards are loading very slowly for the whole team and most pages time out.", "bug", "urgent", "status-incidents"),
    ("T-1016", "en", "enterprise", "Data processing agreement", "Our legal team needs your DPA and the list of subprocessors before we renew.", "how_to", "normal", "data-privacy"),
    ("T-1017", "de", "team", "Doppelte Abbuchung", "Hallo, am 1. Juni wurden uns 60 USD zweimal abgebucht (Rechnung INV-2026-004530). Bitte erstatten Sie die doppelte Zahlung. Jonas Feldmann", "billing", "high", "billing-refunds"),
    ("T-1018", "es", "business", "No puedo entrar", "Desde ayer no puedo iniciar sesión. Dice que mi contraseña es incorrecta, pero la acabo de cambiar.", "account_access", "high", "account-login"),
    ("T-1019", "en", "team", "Delete my account", "Please delete my personal account and my data. I am leaving the company.", "account_access", "normal", "data-privacy"),
    ("T-1020", "en", "team", "Gantt view?", "Is there a timeline or Gantt view? I could not find one.", "feature_request", "low", "feature-requests"),
    ("T-1021", "en", "business", "Recurring cards", "Can I make a card that comes back every Monday?", "how_to", "low", "boards-automations"),
    ("T-1022", "en", "team", "Charged after cancelling", "I cancelled on 2 June but was charged again on 5 June. Why?", "billing", "high", None),
    ("T-1023", "en", "enterprise", "Contract question", "Our procurement team wants to change the liability clause in our Enterprise agreement. Who should they talk to?", "how_to", "normal", None),
    ("T-1024", "ja", "team", "請求書のダウンロード", "請求書をPDFでダウンロードするにはどうすればいいですか？", "how_to", "low", "billing-invoices"),
]
print(len(DEV), "dev tickets")

%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__}"

reply, err = ask(llm, "Reply with the single word: ready")
print("connector:", type(llm).__name__)
print("reply:    ", reply.text.strip() if reply else f"no answer ({err})")
print("tokens:   ", reply.usage_metadata if reply else None)

tid, lang, plan, subject, body, gold_cat, gold_pri, gold_article = DEV[0]   # T-1001

system = ("You triage support tickets for Quillpath, a project-management app.\nCategories:\n"
          + "\n".join(f"- {c}: {d}" for c, d in CATEGORY_DEFINITIONS.items())
          + "\nPriorities:\n"
          + "\n".join(f"- {p}: {d}" for p, d in PRIORITY_DEFINITIONS.items())
          + "\nReply with two words only: the category, then the priority.")

reply, err = ask(llm, [("system", system), ("human", f"Subject: {subject}\n\n{body}")])
if reply is None:
    print(tid, "no answer:", err)
else:
    words = reply.text.strip().lower().split()
    print(tid, "model:", words, "| gold:", [gold_cat, gold_pri])
    print("category correct:", gold_cat in words, "| priority correct:", gold_pri in words)

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 collections import Counter
import math

corpus_words = CORPUS.split()
context = ("get", "a")
after = Counter(corpus_words[i + 2] for i in range(len(corpus_words) - 2)
                if tuple(corpus_words[i:i + 2]) == context)
total = sum(after.values())

for word, n in after.most_common():
    p = n / total
    print(f"{word:8} count {n:4}  p = {p:.3f}  loss if it comes next = {-math.log(p):.2f}")

%pip install -q tokenizers==0.23.3

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)
opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
g = torch.Generator().manual_seed(0)

for step in range(1, 601):
    starts = torch.randint(0, len(data) - 129, (16,), generator=g)
    x = torch.stack([data[s:s + 128] for s in starts]).to(device)
    y = torch.stack([data[s + 1:s + 129] for s in starts]).to(device)   # same text, one token on
    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 in (1, 100, 200, 400, 600):
        print(f"step {step:3}  loss {loss.item():.3f}")
model.eval()

@torch.no_grad()
def reply_by_hand(prompt, max_tokens=20):
    ids = tok.encode(prompt).ids
    for step in range(max_tokens):
        x = torch.tensor([ids[-model.context:]], device=device)
        probs = torch.softmax(model(x)[0, -1], dim=-1)   # probabilities for the next token
        p, best = probs.max(dim=0)                       # greedy: take the top one
        piece = tok.decode([int(best)])
        if piece == "\n":
            break
        ids.append(int(best))
        print(f"{step:2}  {piece!r:13} p = {p.item():.3f}")

reply_by_hand("Customer (Ana): Hi, can I get a refund on my annual plan?\nAgent (Lena):")

prompt = ("Customer (Ana): Hi, my card was charged 288 USD twice on 1 June. "
          "Please refund the duplicate.\nAgent (Lena):")
x = torch.tensor([tok.encode(prompt).ids], device=device)
with torch.no_grad():
    probs = torch.softmax(model(x)[0, -1], dim=-1)
top_p, top_ids = probs.topk(4)
print("first-token guesses:", [(tok.decode([i]), round(p, 3)) for p, i in zip(top_p.tolist(), top_ids.tolist())])
reply_by_hand(prompt, max_tokens=14)

question = f"Write a one-sentence reply to this customer: {body}"
pieces, err = [], None
for attempt in range(2):                     # one retry, as in ask
    try:
        pieces = [chunk.text for chunk in llm.stream(question) if chunk.text]
        break
    except Exception as e:
        err = f"provider_error: {type(e).__name__}"
if pieces:
    print(len(pieces), "pieces; first three:", pieces[:3])
    print("".join(pieces))
else:
    print("no answer:", err)

def triage(row):
    reply, err = ask(llm, [("system", system), ("human", f"Subject: {row[3]}\n\n{row[4]}")])
    if reply is None:
        return "no answer"
    words = reply.text.strip().lower().split()
    return words[0] if words else ""

results = [(row[0], triage(row), row[5]) for row in DEV[:8]]
for tid, got, gold in results:
    print(tid, "model:", got, "| gold:", gold, "| ok" if got == gold else "" if got == "no answer" else "| WRONG")

answered = [r for r in results if r[1] != "no answer"]
accuracy = sum(got == gold for _, got, gold in answered) / max(len(answered), 1)
assert all(got in CATEGORIES for _, got, _ in answered), "a reply was not one of the six categories"
print(f"category accuracy on {len(answered)} of 8 dev tickets answered: {accuracy:.0%}")