{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "faa7c16d",
   "metadata": {},
   "source": [
    "# 6. Pretrained LLMs as Tools for Text Analysis\n",
    "\n",
    "Behind the Secrets of LLMs, crash course part 2\n",
    "\n",
    "In notebooks 4 and 5 we built language models ourselves. Now we load real pretrained models from Hugging Face and use the things we already know (tokens, logits, probabilities, cross-entropy) to analyse text.\n",
    "\n",
    "| section | what we do | model |\n",
    "|---|---|---|\n",
    "| 1 and 2 | real BPE tokenizers, and how fair they are across languages | GPT-2, XLM-R and Qwen tokenizers |\n",
    "| 3 | next-token probabilities | GPT-2 (124M) |\n",
    "| 4 | decoding: greedy, temperature, top-k, top-p | GPT-2 |\n",
    "| 5 and 6 | perplexity and surprisal per token | GPT-2 |\n",
    "| 7 | zero-shot classification by comparing likelihoods | GPT-2 |\n",
    "| 8 | sentence embeddings and semantic search | MiniLM (22M) |\n",
    "| 9 | an instruction-tuned chat model and its chat template | Qwen2.5-0.5B-Instruct (optional) |\n",
    "\n",
    "Time: about 60 minutes. Everything runs on the free Colab CPU, and a T4 GPU makes section 9 faster. The models are downloaded the first time you use them (about 0.5 GB for sections 1 to 8, and 1 GB more for section 9)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "244d4fec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Colab already has these. Locally:  pip install torch transformers matplotlib\n",
    "import math\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "torch.manual_seed(0)\n",
    "\n",
    "LM_NAME   = \"gpt2\"                                       # 124 M parameters, trained on English web text\n",
    "TOKENIZER_NAMES = [\"gpt2\", \"xlm-roberta-base\", \"Qwen/Qwen2.5-0.5B\"]\n",
    "EMB_NAME  = \"sentence-transformers/all-MiniLM-L6-v2\"     # sentence-embedding model\n",
    "CHAT_NAME = \"Qwen/Qwen2.5-0.5B-Instruct\"                 # small multilingual chat model\n",
    "print(\"device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b808076f",
   "metadata": {},
   "source": [
    "## 1. A real tokenizer\n",
    "\n",
    "GPT-2 uses byte-level BPE with 50,257 tokens. It is the same algorithm you wrote in notebook 4, just trained on much more text. The symbol `Ġ` marks a token that starts with a space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07e63463",
   "metadata": {},
   "outputs": [],
   "source": [
    "tok = AutoTokenizer.from_pretrained(LM_NAME)\n",
    "\n",
    "s = \"Tokenization splits unbelievably long words into pieces.\"\n",
    "ids = tok(s)[\"input_ids\"]\n",
    "print(\"vocabulary size:\", len(tok))\n",
    "print(\"ids   :\", ids)\n",
    "print(\"tokens:\", tok.convert_ids_to_tokens(ids))\n",
    "print(\"decode:\", tok.decode(ids))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "310b32d8",
   "metadata": {},
   "source": [
    "## 2. Are tokenizers fair across languages?\n",
    "\n",
    "**Fertility** is the average number of tokens per word. It matters because you pay per token, in compute, in money and in context length. A language with fertility 4 fits only a quarter as much text into the same context window as a language with fertility 1.\n",
    "\n",
    "Below, one sentence is translated into six languages and measured with three tokenizers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c07ea67",
   "metadata": {},
   "outputs": [],
   "source": [
    "sentences = {\n",
    "    \"English\": \"The weather is nice today and we are going for a walk in the park.\",\n",
    "    \"German\":  \"Das Wetter ist heute schön und wir gehen im Park spazieren.\",\n",
    "    \"Spanish\": \"Hoy hace buen tiempo y vamos a dar un paseo por el parque.\",\n",
    "    \"Russian\": \"Сегодня хорошая погода, и мы идём гулять в парк.\",\n",
    "    \"Turkish\": \"Bugün hava güzel ve parkta yürüyüşe çıkıyoruz.\",\n",
    "    \"Hindi\":   \"आज मौसम अच्छा है और हम पार्क में टहलने जा रहे हैं।\",\n",
    "}\n",
    "\n",
    "fertility = {}\n",
    "for name in TOKENIZER_NAMES:\n",
    "    t = AutoTokenizer.from_pretrained(name)\n",
    "    fertility[name] = {lang: len(t(s, add_special_tokens=False)[\"input_ids\"]) / len(s.split())\n",
    "                       for lang, s in sentences.items()}\n",
    "\n",
    "langs = list(sentences)\n",
    "w = 0.8 / len(TOKENIZER_NAMES)\n",
    "plt.figure(figsize=(10, 4))\n",
    "for j, name in enumerate(TOKENIZER_NAMES):\n",
    "    plt.bar([i + j * w for i in range(len(langs))], [fertility[name][l] for l in langs], w, label=name)\n",
    "plt.xticks([i + w for i in range(len(langs))], langs)\n",
    "plt.ylabel(\"tokens per word\"); plt.title(\"Tokenizer fertility (lower = cheaper)\"); plt.legend(); plt.show()\n",
    "\n",
    "print(\"Hindi with GPT-2:\", AutoTokenizer.from_pretrained(\"gpt2\").tokenize(sentences[\"Hindi\"])[:12], \"...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "815db85a",
   "metadata": {},
   "source": [
    "GPT-2's BPE was learned almost only from English, so Hindi falls apart into byte pieces. Multilingual tokenizers like XLM-R and Qwen are much fairer, though rarely perfectly fair.\n",
    "\n",
    "## 3. Next-token prediction with a real model\n",
    "\n",
    "This works exactly like our TinyGPT. The model returns logits for every position, and the last position gives the distribution of the next token:\n",
    "\n",
    "$$P(x_{t+1} \\mid x_{\\le t}) = \\text{softmax}(\\mathbf{z}_t).$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bca02ef1",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(LM_NAME).to(device).eval()\n",
    "print(f\"{sum(p.numel() for p in model.parameters()) / 1e6:.0f} M parameters\")\n",
    "\n",
    "def next_token_probs(prompt, k=10):\n",
    "    enc = tok(prompt, return_tensors=\"pt\").to(device)\n",
    "    with torch.no_grad():\n",
    "        logits = model(**enc).logits              # (1, T, vocab)\n",
    "    probs = F.softmax(logits[0, -1], dim=-1)\n",
    "    top = torch.topk(probs, k)\n",
    "    return [tok.decode(int(i)) for i in top.indices], top.values.cpu()\n",
    "\n",
    "words, p = next_token_probs(\"The capital of Germany is\")\n",
    "plt.barh(words[::-1], p.numpy()[::-1])\n",
    "plt.xlabel(\"probability\"); plt.title(\"P(next token | 'The capital of Germany is')\"); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4011167",
   "metadata": {},
   "source": [
    "## 4. Decoding strategies\n",
    "\n",
    "The model gives us a distribution, and we decide how to pick from it.\n",
    "\n",
    "* **greedy**: always take the most likely token. Deterministic, and often repetitive.\n",
    "* **temperature** $T$: sample from $\\text{softmax}(\\mathbf{z}/T)$.\n",
    "* **top-k**: keep only the $k$ most likely tokens and renormalise.\n",
    "* **top-p (nucleus)**: keep the smallest set $S$ of most likely tokens with $\\sum_{w \\in S} P(w) \\ge p$, and renormalise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3f12ee3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_next(logits, temperature=1.0, top_k=None, top_p=None):\n",
    "    if temperature == 0:                                    # greedy\n",
    "        return logits.argmax(dim=-1, keepdim=True)\n",
    "    logits = logits / temperature\n",
    "    if top_k is not None:\n",
    "        kth = torch.topk(logits, top_k).values[..., -1, None]\n",
    "        logits = logits.masked_fill(logits < kth, float(\"-inf\"))\n",
    "    if top_p is not None:\n",
    "        sorted_logits, sorted_idx = torch.sort(logits, descending=True)\n",
    "        probs = F.softmax(sorted_logits, dim=-1)\n",
    "        # drop a token if the tokens BEFORE it already cover probability mass top_p\n",
    "        drop = probs.cumsum(dim=-1) - probs > top_p\n",
    "        sorted_logits = sorted_logits.masked_fill(drop, float(\"-inf\"))\n",
    "        logits = torch.full_like(logits, float(\"-inf\")).scatter(-1, sorted_idx, sorted_logits)\n",
    "    return torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)\n",
    "\n",
    "def generate(prompt, n_tokens=30, **kw):\n",
    "    ids = tok(prompt, return_tensors=\"pt\")[\"input_ids\"].to(device)\n",
    "    for _ in range(n_tokens):\n",
    "        with torch.no_grad():\n",
    "            logits = model(ids).logits[:, -1]               # (real systems cache keys/values here)\n",
    "        ids = torch.cat([ids, sample_next(logits, **kw)], dim=1)\n",
    "    return tok.decode(ids[0])\n",
    "\n",
    "prompt = \"Dresden is a city in Germany. It is famous for\"\n",
    "torch.manual_seed(0)\n",
    "for kw in [dict(temperature=0), dict(temperature=0.7), dict(temperature=1.5), dict(top_k=5), dict(top_p=0.9)]:\n",
    "    print(f\"--- {kw}\\n{generate(prompt, **kw)}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39c1e25f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The library does the same (plus caching) in one call:\n",
    "enc = tok(prompt, return_tensors=\"pt\").to(device)\n",
    "out = model.generate(**enc, max_new_tokens=30, do_sample=True, temperature=0.7, top_p=0.9,\n",
    "                     pad_token_id=tok.eos_token_id)\n",
    "print(tok.decode(out[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94c38cff",
   "metadata": {},
   "source": [
    "## 5. Perplexity as a tool for text analysis\n",
    "\n",
    "From notebook 4, the perplexity is $\\exp(\\text{mean cross-entropy})$. With a pretrained LM this gives us a fluency meter for any text:\n",
    "\n",
    "$$\\text{PPL}(x) = \\exp\\Big(-\\frac{1}{n-1}\\sum_{t=2}^{n}\\log P(x_t \\mid x_{<t})\\Big)$$\n",
    "\n",
    "If we pass `labels=input_ids`, Hugging Face computes exactly this cross-entropy. It shifts the labels by one position internally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1747af05",
   "metadata": {},
   "outputs": [],
   "source": [
    "def perplexity(text):\n",
    "    enc = tok(text, return_tensors=\"pt\").to(device)\n",
    "    with torch.no_grad():\n",
    "        loss = model(**enc, labels=enc[\"input_ids\"]).loss\n",
    "    return math.exp(loss.item())\n",
    "\n",
    "texts = [\n",
    "    \"The cat sat on the mat and looked out of the window.\",\n",
    "    \"Window the of out looked and mat the on sat cat the.\",\n",
    "    \"Colorless green ideas sleep furiously.\",\n",
    "    \"To be, or not to be, that is the question.\",\n",
    "    \"Sein oder Nichtsein, das ist hier die Frage.\",\n",
    "    \"xq zvb plmr tkw oiu.\",\n",
    "]\n",
    "for t in texts:\n",
    "    print(f\"PPL {perplexity(t):10.1f}   {t}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "271d81e3",
   "metadata": {},
   "source": [
    "A few things to notice:\n",
    "\n",
    "* Word order matters a lot. The same words in a scrambled order get a much higher perplexity.\n",
    "* Grammatical nonsense lands somewhere in between.\n",
    "* Famous text gets a much lower perplexity than ordinary sentences, because the model has partly memorised it. This is the basic idea behind detecting **benchmark contamination**: if a model finds a test set suspiciously easy to predict, it may have seen it during training.\n",
    "* An English-only model finds German surprising, even when the German is perfectly fluent.\n",
    "\n",
    "## 6. Where exactly is the model surprised?\n",
    "\n",
    "The **surprisal** of a token in bits is $s_t = -\\log_2 P(x_t \\mid x_{<t})$. The average surprisal over a text is $\\log_2$ of its perplexity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed050f10",
   "metadata": {},
   "outputs": [],
   "source": [
    "def token_surprisal(text):\n",
    "    ids = tok(text, return_tensors=\"pt\")[\"input_ids\"].to(device)\n",
    "    with torch.no_grad():\n",
    "        logp = F.log_softmax(model(ids).logits[0, :-1], dim=-1)    # predictions for tokens 2..n\n",
    "    s = -logp.gather(1, ids[0, 1:, None]).squeeze(1) / math.log(2)\n",
    "    return tok.convert_ids_to_tokens(ids[0, 1:]), s.cpu()\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(13, 3.5), sharey=True)\n",
    "for ax, t in zip(axes, [\"I take my coffee with milk and sugar.\", \"I take my coffee with milk and concrete.\"]):\n",
    "    toks, s = token_surprisal(t)\n",
    "    ax.bar(range(len(toks)), s)\n",
    "    ax.set_xticks(range(len(toks))); ax.set_xticklabels([x.replace(\"Ġ\", \" \") for x in toks], rotation=45)\n",
    "    ax.set_title(t); ax.set_ylabel(\"surprisal (bits)\")\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d22a46a5",
   "metadata": {},
   "source": [
    "## 7. Zero-shot classification by comparing likelihoods\n",
    "\n",
    "No training at all. We write the task as text and check which continuation the model finds more likely:\n",
    "\n",
    "$$\\hat{y} = \\arg\\max_{y \\in \\{\\text{positive}, \\text{negative}\\}} \\log P(y \\mid \\text{prompt}).$$\n",
    "\n",
    "Small base models often prefer one label no matter what the input is. A well-known fix, called contextual calibration, subtracts the score each label gets for an **empty** review:\n",
    "$\\;\\hat{y} = \\arg\\max_y \\big[\\log P(y \\mid \\text{prompt}) - \\log P(y \\mid \\text{empty prompt})\\big]$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b45ac816",
   "metadata": {},
   "outputs": [],
   "source": [
    "reviews = [\n",
    "    (\"An absolute delight from start to finish, I loved every minute.\", \"positive\"),\n",
    "    (\"The acting was wooden and the plot made no sense.\", \"negative\"),\n",
    "    (\"A beautiful, moving story with a fantastic soundtrack.\", \"positive\"),\n",
    "    (\"Two hours of my life I will never get back.\", \"negative\"),\n",
    "    (\"Funny, warm and surprisingly clever.\", \"positive\"),\n",
    "    (\"Boring, predictable and far too long.\", \"negative\"),\n",
    "    (\"I would happily watch it again tomorrow.\", \"positive\"),\n",
    "    (\"The worst film I have seen this year.\", \"negative\"),\n",
    "]\n",
    "template = 'Review: \"{}\"\\nSentiment (positive or negative):'\n",
    "\n",
    "def logprob_of_continuation(prompt, continuation):\n",
    "    p_ids = tok(prompt, return_tensors=\"pt\")[\"input_ids\"]\n",
    "    c_ids = tok(continuation, return_tensors=\"pt\")[\"input_ids\"]\n",
    "    ids = torch.cat([p_ids, c_ids], dim=1).to(device)\n",
    "    with torch.no_grad():\n",
    "        logp = F.log_softmax(model(ids).logits[0, :-1], dim=-1)\n",
    "    lp = logp.gather(1, ids[0, 1:, None]).squeeze(1)\n",
    "    return lp[-c_ids.shape[1]:].sum().item()               # only the continuation's tokens\n",
    "\n",
    "labels = [\"positive\", \"negative\"]\n",
    "prior = {lab: logprob_of_continuation(template.format(\"\"), \" \" + lab) for lab in labels}   # the bias\n",
    "\n",
    "n_raw = n_cal = 0\n",
    "print(f\"{'raw':10s}{'calibrated':12s}{'true':10s}review\")\n",
    "for text, gold in reviews:\n",
    "    s = {lab: logprob_of_continuation(template.format(text), \" \" + lab) for lab in labels}\n",
    "    raw = max(labels, key=lambda lab: s[lab])\n",
    "    cal = max(labels, key=lambda lab: s[lab] - prior[lab])\n",
    "    n_raw += raw == gold; n_cal += cal == gold\n",
    "    print(f\"{raw:10s}{cal:12s}{gold:10s}{text}\")\n",
    "print(f\"\\naccuracy  raw {n_raw}/{len(reviews)}   calibrated {n_cal}/{len(reviews)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "266354c7",
   "metadata": {},
   "source": [
    "## 8. Sentence embeddings and semantic search\n",
    "\n",
    "An encoder model turns a sentence into one vector. Here we take the mean of the token vectors. Sentences with a similar meaning point in a similar direction, which we measure with the **cosine similarity**\n",
    "\n",
    "$$\\cos(\\mathbf{u}, \\mathbf{v}) = \\frac{\\mathbf{u}\\cdot\\mathbf{v}}{\\lVert\\mathbf{u}\\rVert\\,\\lVert\\mathbf{v}\\rVert}.$$\n",
    "\n",
    "This is the retrieval part of retrieval-augmented generation (RAG)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f34ac0dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "etok = AutoTokenizer.from_pretrained(EMB_NAME)\n",
    "emodel = AutoModel.from_pretrained(EMB_NAME).to(device).eval()\n",
    "\n",
    "def embed(texts):\n",
    "    enc = etok(texts, padding=True, truncation=True, return_tensors=\"pt\").to(device)\n",
    "    with torch.no_grad():\n",
    "        h = emodel(**enc).last_hidden_state                 # (B, T, d)\n",
    "    mask = enc[\"attention_mask\"].unsqueeze(-1).float()      # ignore padding tokens\n",
    "    e = (h * mask).sum(dim=1) / mask.sum(dim=1)             # mean pooling\n",
    "    return F.normalize(e, dim=-1)                           # unit length -> dot product = cosine\n",
    "\n",
    "docs = [\n",
    "    \"Submit your training script to the cluster with sbatch.\",\n",
    "    \"Request a GPU by adding --gres=gpu:1 to your job.\",\n",
    "    \"Use scancel to stop a job that is stuck.\",\n",
    "    \"Knead the dough for ten minutes and let it rise.\",\n",
    "    \"Bake the bread at 220 degrees for thirty minutes.\",\n",
    "    \"The striker scored twice in the second half.\",\n",
    "    \"The match ended in a draw after extra time.\",\n",
    "]\n",
    "E = embed(docs)\n",
    "S = (E @ E.T).cpu()\n",
    "plt.figure(figsize=(6, 5)); plt.imshow(S, cmap=\"viridis\", vmin=0, vmax=1); plt.colorbar()\n",
    "plt.xticks(range(len(docs)), [f\"d{i}\" for i in range(len(docs))]); plt.yticks(range(len(docs)), [d[:40] for d in docs])\n",
    "plt.title(\"cosine similarity\"); plt.show()\n",
    "\n",
    "query = \"How do I get a graphics card for my experiment?\"\n",
    "scores = (embed([query]) @ E.T)[0].cpu()\n",
    "print(\"query:\", query)\n",
    "for i in scores.argsort(descending=True)[:3].tolist():\n",
    "    print(f\"  {scores[i]:.3f}  {docs[i]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d04b4e0",
   "metadata": {},
   "source": [
    "## 9. (Optional) An instruction-tuned chat model\n",
    "\n",
    "Chat models are base models that were fine-tuned on conversations. They expect a **chat template**, which uses special tokens to mark who is speaking. Here is the raw text the model actually gets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7808f112",
   "metadata": {},
   "outputs": [],
   "source": [
    "ctok = AutoTokenizer.from_pretrained(CHAT_NAME)\n",
    "cmodel = AutoModelForCausalLM.from_pretrained(CHAT_NAME)\n",
    "cmodel = cmodel.to(device=device, dtype=torch.bfloat16 if device == \"cuda\" else torch.float32).eval()\n",
    "\n",
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": \"You are a concise, helpful assistant.\"},\n",
    "    {\"role\": \"user\", \"content\": \"Explain self-attention to a 10-year-old in three sentences.\"},\n",
    "]\n",
    "chat_text = ctok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "print(chat_text)\n",
    "\n",
    "enc = ctok(chat_text, return_tensors=\"pt\").to(device)\n",
    "out = cmodel.generate(**enc, max_new_tokens=120, do_sample=True, temperature=0.7, top_p=0.9)\n",
    "print(ctok.decode(out[0, enc[\"input_ids\"].shape[1]:], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "464782d2",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "### Exercises\n",
    "\n",
    "1. **Your language.** Add a sentence in a language you speak to `sentences`. Which tokenizer treats it most fairly? Roughly how much more would a 1,000-word document cost than in English?\n",
    "2. **Find memorised text.** Find a sentence of 15 or more words with a GPT-2 perplexity below 5. Hint: try famous speeches, legal boilerplate or the opening of a constitution. Why does this matter when LLMs are tested on public benchmarks?\n",
    "3. **Decoding.** Generate 5 samples with `temperature=1.5` and 5 with `top_p=0.5`. Which settings would you use for code, for poems and for factual answers?\n",
    "4. **Prompt sensitivity.** Change only the wording of `template` in section 7, for example to `\"The movie was\"`. How much does the accuracy change? Try the same with the chat model from section 9.\n",
    "5. **Search across languages.** Add German or Hindi sentences to `docs`. Does `all-MiniLM-L6-v2` still find them? Try again with `sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2`.\n",
    "\n",
    "### Running it on the cluster\n",
    "\n",
    "`hpc_lab/06_hf_generate` runs a chat model over a whole file of prompts as a GPU batch job (`sbatch generate.sbatch`) and prints the tokens per second. Compare an A100 with this Colab session."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}