{
 "cells": [
  {
   "id": "cell-00",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Base model workshop\n",
    "\n",
    "In this notebook, you send prompts to Llama 3.1 8B, a Base model, through the ACS Infra API. By the end, you will have sampled it at several temperatures, read its token probabilities and its Activations, and steered it to talk like a pirate.\n",
    "\n",
    "Each section after Set up has a cell that you run, then an exercise that you change and run again. Each section names the page of the ACS Infra docs that explains it.\n",
    "\n",
    "You need `ACS_API_BASE` and an API key, as the Quickstart sets them up, and Python 3.14 for `compression.zstd`, NumPy, and the `openai` SDK. The docs' recipe Run the workshop notebook says how to start the notebook with Python, NumPy, and the SDK. Run the cells in order, from the top."
   ]
  },
  {
   "id": "cell-01",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up\n",
    "\n",
    "The next cell reads the API's address from `ACS_API_BASE` and your API key from `ACS_API_KEY`, as the Quickstart sets them up. If `ACS_API_KEY` is not set, the cell asks for your key, and the notebook does not show it.\n",
    "\n",
    "The cell sets two names that the later cells use. `MODEL` is the Base model of every request. `LAYER` is the Layer that the last section steers, near the middle of the model's 32 Layers, 0 to 31.\n",
    "\n",
    "The cell also defines `decode`, the helper from Read Activations, word for word. Then it reads the model's facts with `GET /v1/models/{model}`, and prints its `id`, and the `max_model_len`, `n_layers`, and `hidden_size` in its `capabilities`."
   ]
  },
  {
   "id": "cell-02",
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import base64\n",
    "import getpass\n",
    "import os\n",
    "from compression import zstd\n",
    "\n",
    "import numpy as np\n",
    "from openai import OpenAI\n",
    "\n",
    "MODEL = \"llama-8b\"\n",
    "LAYER = 15\n",
    "\n",
    "if \"ACS_API_KEY\" not in os.environ:\n",
    "    os.environ[\"ACS_API_KEY\"] = getpass.getpass(\"API key: \")\n",
    "\n",
    "client = OpenAI(base_url=os.environ[\"ACS_API_BASE\"], api_key=os.environ[\"ACS_API_KEY\"])\n",
    "\n",
    "\n",
    "def decode(tensor):\n",
    "    raw = zstd.decompress(base64.b64decode(tensor[\"data\"]))\n",
    "    array = np.frombuffer(raw, dtype=tensor[\"dtype\"]).reshape(tensor[\"shape\"])\n",
    "    if tensor[\"original_dtype\"] == \"torch.bfloat16\":\n",
    "        array = (array.view(np.uint16).astype(np.uint32) << 16).view(np.float32)\n",
    "    return array\n",
    "\n",
    "\n",
    "model = client.models.retrieve(MODEL)\n",
    "capabilities = model.capabilities\n",
    "print(model.id, capabilities[\"max_model_len\"], capabilities[\"n_layers\"], capabilities[\"hidden_size\"])"
   ]
  },
  {
   "id": "cell-03",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Send your first prompt\n",
    "\n",
    "This cell sends the prompt `\"The sky is\"`, asks for 2 tokens and the log probability of each, and prints the text and the log probabilities.\n",
    "\n",
    "A Base model continues your text, and the API applies no chat template. The Quickstart explains each field of the Completion, and Base models explains why a Base model continues text.\n",
    "\n",
    "The first Completion to a cold model waits while the model starts. Cold starts and long requests explains the wait."
   ]
  },
  {
   "id": "cell-04",
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "completion = client.completions.create(model=MODEL, prompt=\"The sky is\", logprobs=1, max_tokens=2)\n",
    "print(completion.choices[0].text)\n",
    "print(completion.choices[0].logprobs.token_logprobs)"
   ]
  },
  {
   "id": "cell-05",
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exercise: write the start of a text that you want the model to continue, such as the first\n",
    "# line of a recipe, a letter, or a news story. A question can go on with more questions instead of an answer.\n",
    "# Then raise max_tokens, and run the cell a few times.\n",
    "my_prompt = \"Ingredients for a lemon cake:\"\n",
    "completion = client.completions.create(model=MODEL, prompt=my_prompt, max_tokens=8)\n",
    "print(repr(completion.choices[0].text))"
   ]
  },
  {
   "id": "cell-06",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample at different temperatures\n",
    "\n",
    "This cell sends the same prompt at three temperatures, two times each, with the same `seed`. `temperature: 0` takes the likeliest token each time, and a higher temperature samples less likely tokens more often. The same request with the same `seed` gets the same tokens again, except where vLLM's batching changes one.\n",
    "\n",
    "Sample continuations explains both fields, in Set the sampling fields and Get the same tokens again."
   ]
  },
  {
   "id": "cell-07",
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for temperature in [0, 0.7, 1.5]:\n",
    "    for _ in range(2):\n",
    "        completion = client.completions.create(\n",
    "            model=MODEL, prompt=\"The sky is\", max_tokens=8, temperature=temperature, seed=7\n",
    "        )\n",
    "        print(temperature, repr(completion.choices[0].text))"
   ]
  },
  {
   "id": "cell-08",
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exercise: ask for 5 choices in one request with n=5, without a seed, at a low and a high\n",
    "# temperature. Count the different texts at each temperature. Then add top_p=0.9, and compare.\n",
    "for temperature in [0.3, 1.5]:\n",
    "    completion = client.completions.create(\n",
    "        model=MODEL, prompt=\"The sky is\", max_tokens=8, n=5, temperature=temperature\n",
    "    )\n",
    "    texts = [choice.text for choice in completion.choices]\n",
    "    print(temperature, len(set(texts)), texts)"
   ]
  },
  {
   "id": "cell-09",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read token probabilities\n",
    "\n",
    "This cell reads log probabilities in two ways. `logprobs: 5` gives the log probability of each generated token, and the 5 likeliest tokens at each generated position. `prompt_logprobs: 5` gives the same for each token of your prompt. The `openai` SDK does not know `prompt_logprobs`, so the cell sends it in `extra_body`.\n",
    "\n",
    "In `prompt_logprobs`, position 0 is `None`, because no token comes before it. On Llama 3.1 8B, position 0 is the BOS token, which Inference adds in front of a text prompt. Each later position maps a token id to the token's `logprob`, its `rank`, and its text in `decoded_token`.\n",
    "\n",
    "Read token probabilities explains each field."
   ]
  },
  {
   "id": "cell-10",
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "completion = client.completions.create(model=MODEL, prompt=\"The sky is\", max_tokens=2, logprobs=5)\n",
    "logprobs = completion.choices[0].logprobs\n",
    "for token, logprob, top in zip(logprobs.tokens, logprobs.token_logprobs, logprobs.top_logprobs):\n",
    "    print(repr(token), logprob, top)\n",
    "\n",
    "completion = client.completions.create(\n",
    "    model=MODEL, prompt=\"The sky is\", max_tokens=2, extra_body={\"prompt_logprobs\": 5}\n",
    ")\n",
    "for position, tokens in enumerate(completion.choices[0].prompt_logprobs):\n",
    "    if tokens is None:\n",
    "        print(position, None)\n",
    "        continue\n",
    "    for token_id, token in tokens.items():\n",
    "        print(position, token_id, repr(token[\"decoded_token\"]), token[\"rank\"], token[\"logprob\"])"
   ]
  },
  {
   "id": "cell-11",
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exercise: which text does the model find likelier? With prompt_logprobs 0, each position holds\n",
    "# only your own token, so the sum from position 1 on is the log-likelihood of your text.\n",
    "# Add texts of your own, and compare a likely text with a surprising one.\n",
    "for text in [\"The sky is blue.\", \"The sky is green.\"]:\n",
    "    completion = client.completions.create(\n",
    "        model=MODEL, prompt=text, max_tokens=1, extra_body={\"prompt_logprobs\": 0}\n",
    "    )\n",
    "    positions = completion.choices[0].prompt_logprobs[1:]\n",
    "    print(repr(text), sum(next(iter(tokens.values()))[\"logprob\"] for tokens in positions))"
   ]
  },
  {
   "id": "cell-12",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read Activations\n",
    "\n",
    "This cell reads the residual stream at Layers 0 and 2 with `output_residual_stream`, and decodes it with `decode`. The `openai` SDK does not know the field, so the cell sends it in `extra_body`. The cell checks the array before it uses it. Then it prints the shape, and the first 4 numbers of the first token at each Layer.\n",
    "\n",
    "The shape is Layers, then tokens, then the model's hidden size. The tokens are the prompt tokens and the generated tokens, except the last generated token. So the array has `usage.prompt_tokens + usage.completion_tokens - 1` tokens.\n",
    "\n",
    "Read Activations explains the encoding and the checks. The residual stream explains what the numbers are."
   ]
  },
  {
   "id": "cell-13",
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "completion = client.completions.create(\n",
    "    model=MODEL, prompt=\"The sky is\", max_tokens=2, extra_body={\"output_residual_stream\": [0, 2]}\n",
    ")\n",
    "residual_stream = decode(completion.activations[\"residual_stream\"])\n",
    "\n",
    "layers, tokens, hidden_size = residual_stream.shape\n",
    "usage = completion.usage\n",
    "assert layers == 2, residual_stream.shape\n",
    "assert tokens == usage.prompt_tokens + usage.completion_tokens - 1, residual_stream.shape\n",
    "assert np.isfinite(residual_stream).all(), \"a number is not finite\"\n",
    "print(layers, tokens, hidden_size)\n",
    "print(residual_stream[:, 0, :4].tolist())"
   ]
  },
  {
   "id": "cell-14",
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exercise: read every Layer with output_residual_stream set to True, and print the norm of each\n",
    "# token's row at each Layer. At which Layers does the norm grow? Compare the BOS token, in\n",
    "# column 0, with the other tokens. Then try a longer prompt: a read of every Layer takes a short\n",
    "# one, as the size caps of Completions say.\n",
    "completion = client.completions.create(\n",
    "    model=MODEL, prompt=\"The sky is\", max_tokens=1, extra_body={\"output_residual_stream\": True}\n",
    ")\n",
    "residual_stream = decode(completion.activations[\"residual_stream\"])\n",
    "norms = np.linalg.norm(residual_stream, axis=-1)\n",
    "for layer, row in enumerate(norms):\n",
    "    print(layer, row.round(1))"
   ]
  },
  {
   "id": "cell-15",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Steer the model to talk like a pirate\n",
    "\n",
    "This cell builds a Steering vector from two sets of prompts with the same meanings, one in a pirate style and one plain. It does four steps:\n",
    "\n",
    "1. Read Layer `LAYER` of each prompt with `max_tokens: 1`.\n",
    "2. Take the mean over the tokens of each prompt, without the BOS token.\n",
    "3. Subtract the mean of the plain set from the mean of the pirate set. The difference is the contrast vector, with one number for each dimension of the hidden size.\n",
    "4. Add the contrast at Layer `LAYER` to one prompt at the scales 0, 4, and 8, at `temperature: 0`. Scale 0 adds nothing, so it is the baseline.\n",
    "\n",
    "The cell sends the vector in `apply_steering_vectors`, in `extra_body`. Compare each text with the scale-0 text. When the text starts to repeat phrases or break into fragments, the scale is too large.\n",
    "\n",
    "Steer with a contrast vector explains each step, and Steer a Completion lists the fields of a Steering vector."
   ]
  },
  {
   "id": "cell-16",
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mean_over_tokens(prompt):\n",
    "    completion = client.completions.create(\n",
    "        model=MODEL, prompt=prompt, max_tokens=1, extra_body={\"output_residual_stream\": [LAYER]}\n",
    "    )\n",
    "    residual_stream = decode(completion.activations[\"residual_stream\"])\n",
    "    return residual_stream[0, 1:].mean(axis=0)\n",
    "\n",
    "\n",
    "pirate = [\n",
    "    \"Arr, the wind be fair and the sea be calm.\",\n",
    "    \"Ahoy, me hearties, hoist the sails!\",\n",
    "    \"Shiver me timbers, that be a fine ship.\",\n",
    "    \"Avast, ye scallywags, the treasure be ours!\",\n",
    "]\n",
    "plain = [\n",
    "    \"The wind is fair and the sea is calm.\",\n",
    "    \"Hello, friends, raise the sails!\",\n",
    "    \"Goodness, that is a fine ship.\",\n",
    "    \"Stop, you rascals, the treasure is ours!\",\n",
    "]\n",
    "contrast = np.mean([mean_over_tokens(p) for p in pirate], axis=0) - np.mean(\n",
    "    [mean_over_tokens(p) for p in plain], axis=0\n",
    ")\n",
    "print(contrast.shape, round(float(np.linalg.norm(contrast)), 2))\n",
    "\n",
    "for scale in [0, 4, 8]:\n",
    "    completion = client.completions.create(\n",
    "        model=MODEL,\n",
    "        prompt=\"My advice for your first day at a new job is\",\n",
    "        max_tokens=16,\n",
    "        temperature=0,\n",
    "        extra_body={\n",
    "            \"apply_steering_vectors\": [\n",
    "                {\"activations\": [contrast.tolist()], \"layer_indices\": [LAYER], \"scale\": scale}\n",
    "            ]\n",
    "        },\n",
    "    )\n",
    "    print(scale, repr(completion.choices[0].text))"
   ]
  },
  {
   "id": "cell-17",
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exercise: equal texts do not show that a steer ran, so check that your vector arrived. This cell\n",
    "# reads Layers LAYER - 1 and LAYER at scale 0 and at scale 4, with max_tokens 1 and temperature 0.\n",
    "# Layer LAYER - 1 comes before the steer, so it must not move. Layer LAYER must move by 4 times the\n",
    "# contrast at each token, up to bfloat16's rounding, which grows with the size of each number. The\n",
    "# cell prints each token's largest gap as a fraction of that token's largest number. Row 0, the BOS\n",
    "# token, has the largest numbers.\n",
    "# Then try \"norm_match\": True: the norm of each row's difference at Layer LAYER is then about\n",
    "# the scale times the norm of the row at scale 0.\n",
    "reads = {}\n",
    "for scale in [0, 4]:\n",
    "    completion = client.completions.create(\n",
    "        model=MODEL,\n",
    "        prompt=\"My advice for your first day at a new job is\",\n",
    "        max_tokens=1,\n",
    "        temperature=0,\n",
    "        extra_body={\n",
    "            \"output_residual_stream\": [LAYER - 1, LAYER],\n",
    "            \"apply_steering_vectors\": [\n",
    "                {\"activations\": [contrast.tolist()], \"layer_indices\": [LAYER], \"scale\": scale}\n",
    "            ],\n",
    "        },\n",
    "    )\n",
    "    reads[scale] = decode(completion.activations[\"residual_stream\"])\n",
    "difference = reads[4] - reads[0]\n",
    "print(LAYER - 1, np.abs(difference[0]).max(axis=-1).round(3))\n",
    "gap = np.abs(difference[1] - 4 * contrast).max(axis=-1)\n",
    "print(LAYER, (gap / np.abs(reads[0][1]).max(axis=-1)).round(4))"
   ]
  },
  {
   "id": "cell-18",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "The residual stream explains why one scale does not fit every Layer, and what `norm_match` does. Run a Job reads the Activations of many prompts in the background. Score a text matches each token by its id."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}