WAIR_Alpaca.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "T4",
      "gpuClass": "premium",
      "toc_visible": true,
      "private_outputs": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "premium"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "To open in Google Colab click here:"
      ],
      "metadata": {
        "id": "4-3gG_velqjw"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Women in AI & Robotics\n",
        "\n",
        "## Text Generation with Alpaca\n",
        "\n",
        "### October 2023\n",
        "\n",
        "This notebook explains how to make use of the capabilities of large language models through some of the interfaces they offer.\n",
        "\n",
        "The language model we will use in this notebook is [Alpaca-LoRA](https://github.com/tloen/alpaca-lora/). This is because it is a very small model, capable to run on a free tier colab (as long as you get access to a GPU). It has also been finetuned for use as an assistant, which makes our job easier as we will not need to spend much time prompt engineering."
      ],
      "metadata": {
        "id": "lZPXdWTb9fg0"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Resource: https://colab.research.google.com/github/eemlcommunity/PracticalSessions2023/blob/main/NLP/NLP_tutorial_Alpaca.ipynb"
      ],
      "metadata": {
        "id": "8SitYjlnfdIZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Install packages\n",
        "!pip install bitsandbytes\n",
        "!pip install -q sentencepiece\n",
        "!pip install -q git+https://github.com/huggingface/transformers@v4.30.2\n",
        "!pip install -q git+https://github.com/huggingface/peft.git\n"
      ],
      "metadata": {
        "id": "X_pz8MuY84Qh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Python imports\n",
        "\n",
        "from peft import PeftModel\n",
        "from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig\n",
        "\n",
        "import functools\n",
        "import textwrap\n",
        "import numpy as np\n",
        "from scipy.special import softmax\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "from google.colab.output import eval_js"
      ],
      "metadata": {
        "id": "CiPz7mWjJVFr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# load the model\n",
        "\n",
        "tokenizer = LlamaTokenizer.from_pretrained(\"decapoda-research/llama-7b-hf\")\n",
        "model = LlamaForCausalLM.from_pretrained(\n",
        "    \"decapoda-research/llama-7b-hf\",\n",
        "    load_in_8bit=True,\n",
        "    device_map=\"auto\",\n",
        ")\n",
        "model = PeftModel.from_pretrained(model, \"tloen/alpaca-lora-7b\")\n",
        "\n",
        "eval_js('google.colab.output.setIframeHeight(\"250\")')"
      ],
      "metadata": {
        "id": "VucO3HSMoJkz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# A Minimalist example for querying an Alpaca model.\n",
        "\n",
        "Let's make a simple question answerer. You could test it with the question: `\"What is the first name of Einstein?\"` But feel free to be creative.\n",
        "\n",
        "If the below prints some text as an answer, it means we can correctly run the Alpaca model (even though the text that is actually sampled might not make much sense for now)."
      ],
      "metadata": {
        "id": "zfxnnriCHdze"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "generation_config = GenerationConfig()\n",
        "prompt = input(\"Enter a question here: \")\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "input_ids = inputs[\"input_ids\"].cuda()\n",
        "generation_output = model.generate(\n",
        "    input_ids=input_ids,\n",
        "    generation_config=generation_config,\n",
        "    return_dict_in_generate=False,\n",
        "    max_new_tokens=32,\n",
        ")\n",
        "answer = tokenizer.decode(generation_output[0])\n",
        "print(\"Answer:\", answer)\n"
      ],
      "metadata": {
        "id": "z_qMOndtHLlZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# On the Alpaca model.\n",
        "\n",
        "The Alpaca model used here has been trained on [the Stanford Alpaca dataset](https://github.com/tatsu-lab/stanford_alpaca). We will need to have a look at the data on which the Alpaca model has been trained, in order to make sure our question is sufficiently in-domain.\n",
        "\n",
        "As all Machine Learning model, natural language models are trained to deal with data coming from a specific distribution.\n",
        "\n",
        "Have a look at the link, and see if you can find a good way of formatting our question into something the model has seen during training."
      ],
      "metadata": {
        "id": "6Sj06bOYNhtS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "generation_config = GenerationConfig()\n",
        "question = input(\"Enter a question here: \")\n",
        "prompt = f\"\"\"<FIND OUT WHAT TO PUT HERE>{question}\"\"\"\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "input_ids = inputs[\"input_ids\"].cuda()\n",
        "generation_output = model.generate(\n",
        "    input_ids=input_ids,\n",
        "    generation_config=generation_config,\n",
        "    return_dict_in_generate=False,\n",
        "    max_new_tokens=32,\n",
        ")\n",
        "answer = tokenizer.decode(generation_output[0])\n",
        "print(\"Answer:\", answer)\n"
      ],
      "metadata": {
        "id": "kHbkBEynPika"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "If all is good, you should now start to see a decent output hidden in the answer."
      ],
      "metadata": {
        "id": "h3ID8jxDSLfy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# On tokens\n",
        "Large Language models these days do not operate on a character level, but on the level of so called Tokens. Each token is multiple characters or even words.\n",
        "\n",
        "In the case of Alpaca, tokens are represented as numbers between 0 and 31999. It has a vocab size of 32000.\n",
        "Alpaca has been trained to work with these tokens, not with individual characters.\n",
        "\n",
        "Therefore, the model uses a so called tokenizer to turn tokens into text, and vice versa, to turn text into a list of tokens."
      ],
      "metadata": {
        "id": "lGTvuhWbVPKD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Print the individual tokens of the previous output.\n",
        "print(generation_output[0])\n",
        "print()\n",
        "print(\"_\".join([tokenizer.decode(token) for token in generation_output[0]]))"
      ],
      "metadata": {
        "id": "STybG_opVvfN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(tokenizer.decode(2694))\n",
        "print(tokenizer.decode(5465))\n",
        "print(tokenizer.decode(3337))\n",
        "print(tokenizer.decode(31999))\n",
        "print(tokenizer(\"Einstein Ейнштейн\").input_ids)"
      ],
      "metadata": {
        "id": "fzOisWgfWpTV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Important to understand here, is that there are three special tokens, which do not really map onto text:\n",
        "\n",
        "*   The \"unknown\" token, used to encode characters which are not in the encoding. These are used during training to replace all characters not known to the tokenizer, something which can happen a lot with languages that don't use a latin script.\n",
        "*   The \"BOS\" token, or Beginning Of Sentence. This token indicates that a new sentence has begun. In many models, this also manipulates the attention of the transformer, making sure tokens that come after this token, cannot attend to tokens that came before this token.\n",
        "*   The \"EOS\" token, or End Of Sentence. In many models, this token indicates the sampler to stop sampling. That is convenient, because it makes sure the answer is returned faster and the model does not sample too many useless tokens.\n",
        "\n",
        "In Alpaca, these are the tokens 0, 1 and 2 respectively.\n",
        "\n"
      ],
      "metadata": {
        "id": "W1CXpa0KYuTn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(tokenizer.decode(0))\n",
        "print(tokenizer.decode(1))\n",
        "print(tokenizer.decode(2))"
      ],
      "metadata": {
        "id": "SLDaLUU-YLMC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now with this understanding, can we clean up our answer to only contain the answer, and no longer the text of our original prompt or other additional nonsense?"
      ],
      "metadata": {
        "id": "ElMUS2NfYH8c"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "generation_config = GenerationConfig()\n",
        "question = input(\"Enter a question here: \")\n",
        "prompt = (\n",
        "    f\"Below is an instruction that describes a task. Write a response\"\n",
        "    f\" that appropriately completes the request.\\n\\n\"\n",
        "    f\"### Instruction:\\n{question}\\n\\n### Response:\"\"\"\n",
        ")\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "input_ids = inputs[\"input_ids\"].cuda()\n",
        "generation_output = model.generate(\n",
        "    input_ids=input_ids,\n",
        "    generation_config=generation_config,\n",
        "    return_dict_in_generate=False,\n",
        "    max_new_tokens=64,\n",
        ")\n",
        "answer = tokenizer.decode(generation_output[0][\"<How to index here?>\"])\n",
        "print(\"Answer:\", answer)\n"
      ],
      "metadata": {
        "id": "966pAw8AbUJH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Greedy sampling and the role of the different sampling methods\n",
        "Let's improve the output of our model a bit.\n",
        "\n",
        "You might have noticed that the above already works great for simple questions like, `What is Einstein's first name?`. But if you ask it questions like `Can you write a paragraph about the role of sampling methods in large language models?`, you would notice a problem."
      ],
      "metadata": {
        "id": "1808d8tKSRsO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "generation_config = GenerationConfig()\n",
        "question = input(\"Enter a question here: \")\n",
        "prompt = (\n",
        "    f\"Below is an instruction that describes a task. Write a response\"\n",
        "    f\" that appropriately completes the request.\\n\\n\"\n",
        "    f\"### Instruction:\\n{question}\\n\\n### Response:\"\"\"\n",
        ")\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "input_ids = inputs[\"input_ids\"].cuda()\n",
        "generation_output = model.generate(\n",
        "    input_ids=input_ids,\n",
        "    generation_config=generation_config,\n",
        "    return_dict_in_generate=False,\n",
        "    max_new_tokens=64,\n",
        ")\n",
        "answer = tokenizer.decode(generation_output[0][input_ids.shape[1]:])\n",
        "paragraphed = '\\n'.join(textwrap.wrap(answer))\n",
        "print(\"Answer:\", paragraphed)\n"
      ],
      "metadata": {
        "id": "jT6T6662P7UU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "What is going on here?\n",
        "\n",
        "The reason is that by default, we are sampling with greedy sampling. Every time we sample a new token from our auto-regressive model, we only take the most probable token.\n",
        "\n",
        "Intuitively, that makes a lot of sense. After all, wouldn't the best token to sample be the most probable one? Therefore, wouldn't the most probable token be the best token to sample?\n",
        "\n",
        "Yes! The most probable token usually is a really good token to sample. However, an issue arises when you _only_ take the most probable continuation. Because while that continuation indeed has a high likelihood, but it is not a typical sample from the model.\n",
        "\n",
        "And while at first this might seem counterintuitive, it is a very important notion to keep in mind. When you sample in a high-dimensional space, all samples you will get within the first few bazillion times you sample your distribution, will actually have a relatively low probability of being sampled. At least, compared to the sequence with the highest probability. These 'true' samples are called _typical_ samples.\n",
        "\n",
        "Usually, in high dimensions, *the most probable sample is not typical*."
      ],
      "metadata": {
        "id": "CN8VhkwmeFQr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# We sample 999, 1000-dimensional vectors from a Gaussian with mean=0, stddev=1\n",
        "x_typ = np.random.normal(size=(999, 1000))\n",
        "plt.hist(np.linalg.norm(x_typ, axis=1), density=True, color='cyan', label=\"Typical sample\")\n",
        "# The most probable point to sample from this 1000-dimensional Gaussian, is:\n",
        "x_maxprob = np.zeros(shape=(1, 1000))\n",
        "plt.hist(np.linalg.norm(x_maxprob, axis=1), density=True, color='red', label=\"Most probable sample\")\n",
        "plt.xlabel(\"Norm of each sample\")\n",
        "plt.ylabel(\"pdf\")\n",
        "plt.legend()\n",
        "# As you can see, in high dimensions,\n",
        "# the most probable sample suddenly looks very different from a typical sample!"
      ],
      "metadata": {
        "id": "VCZ90fWviXyf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "It is better to think of high-dimensional Gaussians as a soap-bubble."
      ],
      "metadata": {
        "id": "rhrklHVOmMvA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "norm = functools.partial(np.linalg.norm, axis=-1, keepdims=True)\n",
        "# Let's project our typical samples on a 2D plane, but maintain their norm.\n",
        "x = x_typ[:, :2] / norm(x_typ[:, :2]) * norm(x_typ)\n",
        "plt.scatter(*x.T, marker=',', s=1)\n",
        "plt.title(\"High dimensional typical samples, rotated onto the 2D plane.\")"
      ],
      "metadata": {
        "id": "4U54yKYMkwjv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "If you want further reading on the typicality-problem, I can recommend these two resources:\n",
        "\n",
        "\n",
        "* Ferenc Huszár's [*Gaussian Distributions are Soap Bubbles*](https://www.inference.vc/high-dimensional-gaussian-distributions-are-soap-bubble/)\n",
        "* Sander Dieleman's [*Musings on typicality*](https://sander.ai/2020/09/01/typicality.html)\n",
        "\n"
      ],
      "metadata": {
        "id": "Z7KTAXN1n5TP"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## The solution: true categorical sampling"
      ],
      "metadata": {
        "id": "71odUAvunIuK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "generation_config = GenerationConfig(\n",
        "    do_sample=True  # This enables categorical sampling in our Alpaca model\n",
        ")\n",
        "question = input(\"Enter a question here: \")\n",
        "prompt = (\n",
        "    f\"Below is an instruction that describes a task. Write a response\"\n",
        "    f\" that appropriately completes the request.\\n\\n\"\n",
        "    f\"### Instruction:\\n{question}\\n\\n### Response:\"\"\"\n",
        ")\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "input_ids = inputs[\"input_ids\"].cuda()\n",
        "generation_output = model.generate(\n",
        "    input_ids=input_ids,\n",
        "    generation_config=generation_config,\n",
        "    return_dict_in_generate=False,\n",
        "    max_new_tokens=256,\n",
        ")\n",
        "answer = tokenizer.decode(generation_output[0][input_ids.shape[1]:])\n",
        "paragraphed = '\\n'.join(textwrap.wrap(answer))\n",
        "print()\n",
        "print(\"Answer:\", paragraphed)"
      ],
      "metadata": {
        "id": "boAT1NdeVF70"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "That is much better! The repetition is gone.\n",
        "\n",
        "However, you might notice that the model does not really stick to the topic very well. It tends to \"drift\" away, as later tokens are increasingly influenced by tokens that are sampled by the model, rather than tokens that were in the prompt.\n",
        "\n",
        "In general, we want to have typical samples, but there was something \"nice\" about the greediness too. There are various common approaches to make the samples stick closer to the most likely prediction of the model, while also keeping them typical. We will discuss some of these in the next chapters."
      ],
      "metadata": {
        "id": "cbsrLdn6qwb5"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## A first intermediate approach: changing the temperature"
      ],
      "metadata": {
        "id": "gtNfhAJKtFX7"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can think of changing the temperature of a categorical distribution as similar to what happens when changing the temperature in statistical thermodynamics.\n",
        "\n",
        "If you are not familiar with this framework, you can play with the slider below.\n",
        "\n",
        "In short if you have your original distribution $p(x)$, changing the temperature creates a distribution $p_t(x)=\\frac{p(x)^t}{\\Sigma_X p(x)^t}$. So you take the probability of the initial distribution to the power of the temperature, and then you renormalize."
      ],
      "metadata": {
        "id": "oOsyY9MntMDx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Play with the temperature of a distribution { run: \"auto\" }\n",
        "\n",
        "temp = 1.35 #@param {type:\"slider\", min:0.01, max:3, step:0.01}\n",
        "\n",
        "np.random.seed(317070)\n",
        "x = softmax(np.random.rand(100,))\n",
        "plt.subplot(1,2,1)\n",
        "plt.bar(range(1, 101), x, width=1.)\n",
        "plt.xlabel(\"Original Categories\")\n",
        "plt.ylabel(\"Probability\")\n",
        "\n",
        "x_temp = np.power(x, 1./(temp + 1e-2))\n",
        "x_temp = x_temp / np.sum(x_temp)\n",
        "plt.subplot(1,2,2, sharey=plt.gca())\n",
        "plt.bar(range(1, 101), x_temp, width=1.)\n",
        "plt.gca().get_yaxis().set_visible(False)\n",
        "plt.xlabel(f\"Categories with {temp=}\")\n",
        "plt.ylabel(\"Probability\")\n",
        "print()"
      ],
      "metadata": {
        "id": "EgsmwaXktLhk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "As you can see, the original distribution has a temperature of 1. This was the true sampling case.\n",
        "\n",
        "When you heat up, the distribution gets more and more uniform. In the limit, when $t=\\infty$, you get the uniform distribution. That is quite pointless for our purpose.\n",
        "\n",
        "When you cool off, less likely samples get less likely, while more likely samples get more likely. In the limit of $t=0$, only the most likely sample remains. This is greedy sampling.\n",
        "\n",
        "So by adjusting the temperature, we get to interpolate between our original $t=0$ greedy sampling, and $t=1$ true sampling.\n",
        "\n",
        "Depending on your application, it can be a good idea to set the temperature to 0.8 to get succinct, accurate answers from your large language model. This will remove some of the \"creativity\" and \"wildness\" though. In general, I have seen values ranging from 0.1 to 1.0 being used."
      ],
      "metadata": {
        "id": "XsBTdNQhyNhy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "generation_config = GenerationConfig(\n",
        "    do_sample=True,\n",
        "    temperature=0.8,  # This sets the temperature of our model\n",
        ")\n",
        "question = input(\"Enter a question here: \")\n",
        "prompt = (\n",
        "    f\"Below is an instruction that describes a task. Write a response\"\n",
        "    f\" that appropriately completes the request.\\n\\n\"\n",
        "    f\"### Instruction:\\n{question}\\n\\n### Response:\"\"\"\n",
        ")\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "input_ids = inputs[\"input_ids\"].cuda()\n",
        "generation_output = model.generate(\n",
        "    input_ids=input_ids,\n",
        "    generation_config=generation_config,\n",
        "    return_dict_in_generate=False,\n",
        "    max_new_tokens=256,\n",
        ")\n",
        "answer = tokenizer.decode(generation_output[0][input_ids.shape[1]:])\n",
        "paragraphed = '\\n'.join(textwrap.wrap(answer))\n",
        "print()\n",
        "print(\"Answer:\", paragraphed)"
      ],
      "metadata": {
        "id": "WP5QNukhLWPj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Other approaches: avoiding the tails, and beam sampling."
      ],
      "metadata": {
        "id": "1tNv5VWl1Bov"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "A problem which is hard to illustrate in a colab, is that once in a blue moon you might still get very unlucky and sample a weird token. In order to avoid that problem, there are various methods that only consider the top tokens during sampling. E.g.\n",
        "\n",
        "* top_k sampling: only look at the k most likely tokens\n",
        "* top_p sampling: only look at the most likely tokens up to probability mass p\n",
        "\n",
        "Both these methods discard the long tail of tokens. In general, like lowering the temperature this takes away a bit of the weirdest wackiness of the answers. Unlike lowering the temperature, this keeps most of the creativity still intact.\n",
        "\n"
      ],
      "metadata": {
        "id": "8BTfNaMu1X4y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "generation_config = GenerationConfig(\n",
        "    do_sample=True,\n",
        "    temperature=0.8,\n",
        ")\n",
        "question = input(\"Enter a question here: \")\n",
        "prompt = (\n",
        "    f\"Below is an instruction that describes a task. Write a response\"\n",
        "    f\" that appropriately completes the request.\\n\\n\"\n",
        "    f\"### Instruction:\\n{question}\\n\\n### Response:\"\"\"\n",
        ")\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "input_ids = inputs[\"input_ids\"].cuda()\n",
        "generation_output = model.generate(\n",
        "    input_ids=input_ids,\n",
        "    generation_config=generation_config,\n",
        "    return_dict_in_generate=False,\n",
        "    max_new_tokens=256,\n",
        ")\n",
        "answer = tokenizer.decode(generation_output[0][input_ids.shape[1]:])\n",
        "paragraphed = '\\n'.join(textwrap.wrap(answer))\n",
        "print()\n",
        "print(\"Answer:\", paragraphed)"
      ],
      "metadata": {
        "id": "pxMLlM681Vvt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Beam sampling actually allows us to go the other way, and to be more greedy than greedy sampling!\n",
        "\n",
        "With greedy sampling, we only consider the most likely token at every point in time. However, the most likely token at every point in time is not necessarily going to give us the most likely sequence within the distribution of sequences. We might sample to greedily early on, and so we might miss out on some high probability token later on.\n",
        "\n",
        "In order to mitigate this somewhat, we can keep track of a number of beams through our probability space. At every point during our sampling, we will try to maintain the $N$ most likely token sequences so far, the so called \"beams\".\n",
        "\n"
      ],
      "metadata": {
        "id": "pcl3q8PK22-n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "generation_config = GenerationConfig(\n",
        "    do_sample=True,\n",
        "    num_beams=4,\n",
        ")\n",
        "question = input(\"Enter a question here: \")\n",
        "prompt = (\n",
        "    f\"Below is an instruction that describes a task. Write a response\"\n",
        "    f\" that appropriately completes the request.\\n\\n\"\n",
        "    f\"### Instruction:\\n{question}\\n\\n### Response:\"\"\"\n",
        ")\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "input_ids = inputs[\"input_ids\"].cuda()\n",
        "generation_output = model.generate(\n",
        "    input_ids=input_ids,\n",
        "    generation_config=generation_config,\n",
        "    return_dict_in_generate=False,\n",
        "    max_new_tokens=64,\n",
        ")\n",
        "answer = tokenizer.decode(generation_output[0][input_ids.shape[1]:])\n",
        "paragraphed = '\\n'.join(textwrap.wrap(answer))\n",
        "print()\n",
        "print(\"Answer:\", paragraphed)"
      ],
      "metadata": {
        "id": "QMdzqIgU22TC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "For more explanations on the various ways these sampling methods work:, I would refer to [this excellent blogpost](https://huggingface.co/blog/how-to-generate).\n",
        "\n",
        "In general, finding a good trade-off between all these parameters is a bit of an art, and can strongly depend on your application. How creative do you need the answers to be? How close to the training data does the model need to stay? How important is it to not have a Chinese token slip in in weird corner cases?"
      ],
      "metadata": {
        "id": "mI_3p9hP3xOH"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Chain of Thought reasoning\n",
        "\n",
        "An important technique for improving the quality of answers to questions, is to give the Language Models a little bit of space to reason before they have to give a final answer. You could think about it as using bit of memory to work with before giving an answer to the question. In general, it is observed that this improves the quality of answer to reasoning questions dramatically.\n",
        "\n",
        "Keep in mind that models sample these tokens auto-regressively. If they answer first, they will have to come up with a reason to make that answer plausible. If they start with the first step, that might be easier to infer directly from the question.\n",
        "\n",
        "For example: ask the model `Piotr and Monika have 2 cans. Each can has 5 pierogi. How many pierogi does Monika have?`"
      ],
      "metadata": {
        "id": "PfTyD7mu7ARy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Without chain of thought reasoning\n",
        "generation_config = GenerationConfig(\n",
        "    do_sample=True,\n",
        "    temperature=0.8,\n",
        ")\n",
        "question = input(\"Enter a question here: \")\n",
        "prompt = (\n",
        "    f\"Below is an instruction that describes a task. Write a response\"\n",
        "    f\" that appropriately completes the request.\\n\\n\"\n",
        "    f\"### Instruction:\\n{question}\\n\\n### Response:\"\"\"\n",
        ")\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "input_ids = inputs[\"input_ids\"].cuda()\n",
        "generation_output = model.generate(\n",
        "    input_ids=input_ids,\n",
        "    generation_config=generation_config,\n",
        "    return_dict_in_generate=False,\n",
        "    max_new_tokens=32,\n",
        ")\n",
        "answer = tokenizer.decode(generation_output[0][input_ids.shape[1]:])\n",
        "paragraphed = '\\n'.join(textwrap.wrap(answer))\n",
        "print()\n",
        "print(\"Answer:\", paragraphed)"
      ],
      "metadata": {
        "id": "bqMRoPW7MvFR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title With chain of thought reasoning\n",
        "\n",
        "# First, we generate a reasoning.\n",
        "generation_config = GenerationConfig(\n",
        "    do_sample=True,\n",
        "    temperature=1.0,\n",
        ")\n",
        "\n",
        "prompt = (\n",
        "    f\"Below is an instruction that describes a task. Write a response \"\n",
        "    f\"that appropriately completes the request.\\n\\n\"\n",
        "    f\"### Instruction:\\nWrite a paragraph to answer the following question. \"\n",
        "    f\"Think step by step. {question}\\n\\n\"\n",
        "    f\"### Response:\\n\"\"\"\n",
        ")\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "input_ids = inputs[\"input_ids\"].cuda()\n",
        "generation_output = model.generate(\n",
        "    input_ids=input_ids,\n",
        "    generation_config=generation_config,\n",
        "    return_dict_in_generate=False,\n",
        "    max_new_tokens=256,\n",
        ")\n",
        "answer = tokenizer.decode(generation_output[0][input_ids.shape[1]:])\n",
        "paragraphed = '\\n'.join(textwrap.wrap(answer))\n",
        "print()\n",
        "print(\"Reasoning:\", paragraphed)\n",
        "\n",
        "# Then, using this reasoning, we ask the model for a final answer.\n",
        "prompt = (\n",
        "    f\"Below is an instruction that describes a task. Write a response \"\n",
        "    f\"that appropriately completes the request.\\n\\n\"\n",
        "    f\"### Instruction:\\nAnswer the following question '{question}' given these steps. {answer}\"\n",
        "    f\"\\n\\n### Response:\"\"\"\n",
        ")\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "input_ids = inputs[\"input_ids\"].cuda()\n",
        "generation_output = model.generate(\n",
        "    input_ids=input_ids,\n",
        "    generation_config=generation_config,\n",
        "    return_dict_in_generate=False,\n",
        "    max_new_tokens=32,\n",
        ")\n",
        "final_answer = tokenizer.decode(generation_output[0][input_ids.shape[1]:])\n",
        "paragraphed = '\\n'.join(textwrap.wrap(final_answer))\n",
        "print()\n",
        "print(\"Answer:\", paragraphed)\n",
        "\n",
        "# <Insert here a step where you ask the model to verify its answer>"
      ],
      "metadata": {
        "id": "SKk-iLxrSadj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Exercise: Can you add a step asking if the model thinks the answer is correct?\n",
        "\n",
        "I do want to note that Alpaca model used for this colab is not great at Chain of Thought reasoning. It is probably too small."
      ],
      "metadata": {
        "id": "mIkofkPPZ-YY"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Using large language models as an interface for algorithms\n",
        "\n",
        "Above, we were using large language models to generate us some text. However, most of the time we don't actually want the model to give us an answer that is understandable to humans. We want the model to give us an answer that can be used by the rest of our computer program. We only need the model to process text as input, and we don't actually want to use text as output.\n",
        "\n",
        "For instance, we might use a language model to tell us if:\n",
        "\n",
        "* Is this datapoint an outlier compared to these other datapoints?\n",
        "* Is this piece of text saying the same as that piece of text?\n",
        "* From this pre-considered list of categories, which category would you say this datapoint is?\n",
        "* Some other model has generated me this piece of text. Do you think it is any good? (See the last exercise in Chain of Thought reasoning.)\n",
        "\n"
      ],
      "metadata": {
        "id": "nkSdq34954A5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title We could sample a yes/no answer\n",
        "generation_config = GenerationConfig(\n",
        "    do_sample=True,\n",
        ")\n",
        "question = \"Is looking at the logits of answers a good way to extract the knowledge of large language models?\"\n",
        "prompt = (\n",
        "    f\"Below is an instruction that describes a task. Write a response\"\n",
        "    f\" that appropriately completes the request.\\n\\n\"\n",
        "    f\"### Instruction:\\n{question}\\n\\n### Response:\\n\"\"\"\n",
        ")\n",
        "for i in range(10):\n",
        "  inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "  input_ids = inputs[\"input_ids\"].cuda()\n",
        "  generation_output = model.generate(\n",
        "      input_ids=input_ids,\n",
        "      generation_config=generation_config,\n",
        "      return_dict_in_generate=False,\n",
        "      max_new_tokens=16,\n",
        "  )\n",
        "  answer = tokenizer.decode(generation_output[0][input_ids.shape[1]:])\n",
        "  paragraphed = '\\n'.join(textwrap.wrap(answer))\n",
        "  print()\n",
        "  print(\"Answer:\", paragraphed)\n",
        "  if \"Yes\" in answer:\n",
        "    print(\"True\")\n",
        "  elif \"No\" in answer:\n",
        "    print(\"False\")\n",
        "  else:\n",
        "    print(\"?\")"
      ],
      "metadata": {
        "id": "BYQzd9Ky53XF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "As you can see, the model is still sampling stochastically. It has not completely made up its mind as to the correct answer of this question. It is also not always using the exact right words for us to do the conversion from a text answer to True/False in the end.\n",
        "\n",
        "A solution to this problem is to instead of looking at the samples of the model, looking at the logits of the model directly.\n"
      ],
      "metadata": {
        "id": "WX30J5JUAxaH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "question = \"Is looking at the logits of answers a good way to extract the knowledge of large language models?\"\n",
        "prompt = (\n",
        "    f\"Below is an instruction that describes a task. Write a response\"\n",
        "    f\" that appropriately completes the request.\\n\\n\"\n",
        "    f\"### Instruction:\\n{question}\\n\\n### Response:\\n\"\"\"\n",
        ")\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "input_ids = inputs[\"input_ids\"].cuda()\n",
        "\n",
        "res = model.forward(input_ids=input_ids,)\n",
        "print(\"The shape of the inputs is:\", input_ids.shape)\n",
        "print(\"The shape of the logits is:\", res.logits.shape)\n",
        "# This result of the forward pass has three dimensions.\n",
        "# * A batch dimension\n",
        "# * A sequence dimension (the input had this many tokens)\n",
        "# * A vocab dimension (the model can choose between this many tokens)\n",
        "next_token_logits = res.logits[0, -1, :].cpu()\n",
        "next_token_prob = softmax(next_token_logits, axis=-1)\n",
        "plt.plot(range(len(next_token_prob)), next_token_prob)\n",
        "plt.xlabel(\"Token index\")\n",
        "plt.ylabel(\"Probability\")\n",
        "plt.show()\n"
      ],
      "metadata": {
        "id": "gLEqyM6zO8Ui"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "As we can see, there are 2 really probable tokens, but there are many more tokens that appear to have some probability mass."
      ],
      "metadata": {
        "id": "uzCNSGM4FGH3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Print the highest probability tokens\n",
        "# Take the 20 most probable tokens which come after the question we asked.\n",
        "highest_prob_tokens = np.argsort(next_token_prob, axis=0)[-1:-50:-1]\n",
        "# Print those tokens\n",
        "print([(tokenizer.decode(i), i) for i in highest_prob_tokens])"
      ],
      "metadata": {
        "id": "FadIbZQEEYSx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here is a part of the annoyance. There are many variations of yes and no available as a token. We should process all of these alternative spellings."
      ],
      "metadata": {
        "id": "lE9C9HPGPOwM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "yes_tokens = [<Find the various tokens which mean yes>]\n",
        "no_tokens = [<Find the various tokens which mean no>]\n",
        "\n",
        "print([(i, tokenizer.decode(i)) for i in yes_tokens])\n",
        "print([(i, tokenizer.decode(i)) for i in no_tokens])\n",
        "\n",
        "yes_probability = np.sum([next_token_prob[i] for i in yes_tokens])\n",
        "no_probability = np.sum([next_token_prob[i] for i in no_tokens])\n",
        "\n",
        "print(f\"Yes probability: {yes_probability}\")\n",
        "print(f\"No probability: {no_probability}\")"
      ],
      "metadata": {
        "id": "1PWALWwUITKW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title The final solution\n",
        "\n",
        "question = input(\"Enter a question here: \")\n",
        "prompt = (\n",
        "    f\"Below is an instruction that describes a task. Write a response\"\n",
        "    f\" that appropriately completes the request.\\n\\n\"\n",
        "    f\"### Instruction:\\n{question}\\n\\n### Response:\\n\"\"\"\n",
        ")\n",
        "\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
        "input_ids = inputs[\"input_ids\"].cuda()\n",
        "res = model.forward(input_ids=input_ids,)\n",
        "next_token_logits = res.logits[0, -1, :].cpu()\n",
        "next_token_prob = softmax(next_token_logits, axis=-1)\n",
        "yes_tokens = [8241, 3582, 21143, 5574, 5574, 3009, 29979, 3869, 20652]\n",
        "no_tokens = [3782, 1217, 6632, 8824, 4541, 29940]\n",
        "yes_probability = np.sum([next_token_prob[i] for i in yes_tokens])\n",
        "no_probability = np.sum([next_token_prob[i] for i in no_tokens])\n",
        "if yes_probability > no_probability:\n",
        "  print(\"True\")\n",
        "else:\n",
        "  print(\"False\")\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "n6a-JlKtPvKG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Conclusion\n",
        "\n",
        "I hope I could convince you that Large Language models are fairly flexible things, and depending on your task or application, might be used as an alternative to training a model yourself.\n",
        "\n",
        "Note that the tiny Alpaca model used in this Colab does not have a great quality. In general, I would recommend using a larger language model, but I do hope I have convinced you that even this tiny language model is a great tool to have under your belt. It can be an enormous help at tasks like data cleaning.\n",
        "\n",
        "## The future\n",
        "\n",
        "Finally, it is my hope that models like these will be able to generate their own datasets on which they can finetune and get better. I expect that people will find out how to use Chain of Thought reasoning to generate a new dataset to finetune the model on, and that by finetuning on this data the model will improve. This process can then be repeated, which one day might bring us a language model with superhuman intelligence. That model would be the closest humanity can hope to get to an oracle.\n",
        "\n",
        "How to exactly do that is still an open question, but I reckon it is one we are close to answering."
      ],
      "metadata": {
        "id": "tGjpxbTgYBeI"
      }
    }
  ]
}