explain-06.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bdcc1ef7-faa5-48da-8c12-e1321664804c",
   "metadata": {},
   "source": [
    "### Some terms you need to understand for understanding Large Language Models\n",
    "  - **logits**\n",
    "  - softmax\n",
    "  - cross entropy\n",
    "  - one hot encoding\n",
    "  - encoding/embedding\n",
    "  - tokens, tokenizer\n",
    "  - attention\n",
    "    - multi-head attention\n",
    "  - positional encoding\n",
    "    - rotary positional encoding\n",
    "  - transformer\n",
    "  - optimizer\n",
    "  - Important players in the AI Ecosystem\n",
    "  - How to train/fine tune!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5277311-7714-45e0-a1b6-4e29b6819bc0",
   "metadata": {},
   "source": [
    "### Most of these terms came about because of odd historical glitches. We're stuck with them because EVERYBODY uses them."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "881d27e8-3b52-41ba-802e-902f9763f04b",
   "metadata": {},
   "source": [
    "#### Assume you have a list of common names in a simple text file, one name per line. Let's call this file `names.txt`\n",
    "#### Assume you want to use the information in this text file to create more names, \n",
    "#### with the provisio that the new names that you make must \"seem similar\" to the names in the `names.txt`\n",
    "#### One simple way to do this, is to use probabilities!\n",
    "#### Assume that you want to measure $P(l_{j+1} | l_{j})$ e.g. given the letter 'e', what is the likelyhood of some other letter (e.g. 'r' ?)\n",
    "#### In other words, start with a random letter, and randomly assign new letters according to the probability distribution you glean from `names.txt`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9390cd1-393f-46d6-973e-09d74748bf52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# first load up names.txt and split them into each line"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22c4e9a5-52b9-410d-9ce7-eebcaea7bf10",
   "metadata": {},
   "outputs": [],
   "source": [
    "words = open('names.txt', 'r').read().splitlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ced37928-405c-47be-bee7-dba7dd69a217",
   "metadata": {},
   "outputs": [],
   "source": [
    "words[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5213dd7e-c042-4de6-8aad-cfe9f98e8e1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from tinytorch import *\n",
    "from subroutines import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b130133-e32b-4221-8ba9-200cab15bedb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from matplotlib.patches import Ellipse\n",
    "from matplotlib.text import OffsetFrom\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1210323a-d59e-461e-b3c0-edd7a19b44ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = torch.zeros((27, 27), dtype=torch.int32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68456d75-c74d-49da-907d-e8deccfe7a63",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the list of all characters used in 'names.txt' \n",
    "chars = sorted(list(set(''.join(words))))\n",
    "# to make things look prettier, use '.' as the first 'token'\n",
    "stoi = {s:i for i,s in enumerate(['.'] + chars)}\n",
    "# stoi is string to int, itos is int to string \n",
    "itos = {i:s for s,i in stoi.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "824fb09a-3629-440a-9b1a-705ce902d379",
   "metadata": {},
   "outputs": [],
   "source": [
    "# now iterate through all the words, and simply count how many times a given transition occurs\n",
    "# i.e. given 'e', how many times does each of the other letters occur?\n",
    "# The '.' is to signify that we are either starting or ending a name.\n",
    "for w in words:\n",
    "  # pad each name with '.' at the start and end\n",
    "  chs = ['.'] + list(w) + ['.']\n",
    "  # for each pair of characters\n",
    "  for ch1, ch2 in zip(chs, chs[1:]):\n",
    "    ix1 = stoi[ch1]\n",
    "    ix2 = stoi[ch2]\n",
    "    N[ix1, ix2] += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30a5cc62-05a2-4fe7-8aa7-7a4f5353b673",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "plt.figure(figsize=(16,16))\n",
    "plt.imshow(N, cmap='Blues')\n",
    "for i in range(27):\n",
    "    for j in range(27):\n",
    "        chstr = itos[i] + itos[j]\n",
    "        plt.text(j, i, chstr, ha=\"center\", va=\"bottom\", color='gray')\n",
    "        plt.text(j, i, N[i, j].item(), ha=\"center\", va=\"top\", color='gray')\n",
    "plt.axis('off');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d82c7aa2-f30a-4bd2-93ba-c88ece2b350f",
   "metadata": {},
   "source": [
    "### Remember we used 0 as the start/stop marker, i.e. '.'\n",
    "### So N[0] is the counts of all of the times a particular letter appeared as a first letter in a name,\n",
    "#### starting from '.' (no name starts with '.'!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f13d0e57-b230-4902-93ff-e53c54b2d8b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(hprt(N[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a51daeb-0308-4791-84b1-3fa6e5a3b885",
   "metadata": {},
   "outputs": [],
   "source": [
    "N[0,stoi['a']],N[0,stoi['b']],N[0,stoi['c']],N[0,stoi['z']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c776484-1a0d-4496-951e-81fae3fc419d",
   "metadata": {},
   "outputs": [],
   "source": [
    "list(map(lambda x: x.item(), [N[0,stoi['a']],N[0,stoi['b']],N[0,stoi['c']],N[0,stoi['z']]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09ef68d2-217a-4dac-b63f-43e99633ff4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# here is how we convert raw counts to probabilities\n",
    "p = N[0].float()\n",
    "p = p / p.sum()\n",
    "print(hprt(ps=p.shape, p=p))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e6b5e1f-cc15-4189-86d4-1524a537a162",
   "metadata": {},
   "source": [
    "### Notice the zero. in the first column. This is generally bad $-\\infty = log(0)$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb171ec9-bfee-4a4b-8bd4-c254d29f1387",
   "metadata": {},
   "source": [
    "### Below, we're are taking a probabilistic sampling from the array p.\n",
    "#### `torch.multinomial` returns a list of indices (of shape `(num_samples)`) \n",
    "#### `replacement` is by default `False` which means if an index is chosen, it will not be chosen again.\n",
    "#### i.e. `replacement=True` allows the same `letter` (e.g. `index` into `p`)  to be chosen multiple times"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df3c6b9b-2d9b-425f-a2ba-9cc982d3a51a",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(1337)\n",
    "PP = torch.multinomial(p, num_samples=3, replacement=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "287368d3-3031-4f60-8762-e51be0af87cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "PP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b31daf17-28a8-4d05-ac27-0e3afb98b46a",
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx in PP.tolist():\n",
    "  print(f'{idx=} {itos[idx]=}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2dd64e7d-f478-4619-9597-b642b838f0dc",
   "metadata": {},
   "source": [
    "## In LLMs, the probability of the next token (or word) is generated \n",
    "  - Alert! LOGITS and SOFTMAX ahead!\n",
    "## Then we call torch.multinomial to select the next token!\n",
    "### So are we done? Can we just use these probabilities to generate cool looking names?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b50533a-509a-414a-af01-09dcfc6b0b59",
   "metadata": {},
   "outputs": [],
   "source": [
    "g = torch.Generator().manual_seed(2147483647)\n",
    "def genNames(N, count):\n",
    "  for i in range(count):\n",
    "    ix = 0\n",
    "    dst = []\n",
    "    ids = []\n",
    "    while True:\n",
    "      p = N[ix].float()\n",
    "      p = p / p.sum()\n",
    "      ix = torch.multinomial(p, num_samples=1, replacement=True, generator=g).item()\n",
    "      ids.append(ix)\n",
    "      dst.append(itos[ix])\n",
    "      # print(f'{ix=} {itos[ix]=}')\n",
    "      if ix == 0:\n",
    "        break\n",
    "    print(f\"{len(ids)=} name={''.join(dst)}\")\n",
    "genNames(N, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c57e34a5-6194-44c2-bf3e-5f30eda602e7",
   "metadata": {},
   "source": [
    "### They are kind of like names, but, not that great!\n",
    "#### But they are better than generating a random probability matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63ea9c7a-e6e0-41f6-a8d7-c296f8a757c0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "FakeN = torch.ones(27,27,dtype = torch.float32) / 27.0\n",
    "assert N.shape == FakeN.shape\n",
    "print(N[0]); print(FakeN[0])\n",
    "genNames(FakeN, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "373b2d34-e360-448e-a2d3-145b8cdc8dd0",
   "metadata": {},
   "source": [
    "## In LLMs, the probability matrix (much like `p` above) is typically generated by using `Softmax(LOGITS)`\n",
    "  - <https://pytorch.org/docs/stable/generated/torch.nn.Softmax.html#softmax>\n",
    "  - <https://pytorch.org/docs/stable/generated/torch.nn.functional.softmax.html#torch-nn-functional-softmax>\n",
    "  - `torch.nn.Softmax(dim=None)` is defined as $\\displaystyle \\frac{\\displaystyle exp(x_i)}{\\displaystyle \\sum_{j=1}^{j=L} exp(x_j)}$\n",
    "    -  $exp(x)$ is ofcourse $e^x$\n",
    "    -  $x_i$ is the $i$'th element of $x$ i.e. `x[i]`\n",
    "    -  $L$ is the length of the tensor (assuming a 1D tensor)\n",
    "    -  The output always sums to 1.0\n",
    "  - In other words, first, exponentiate each element of the tensor (separately) and divide by the sum of all of the exponentiated elements\n",
    "  - Careful! pytorch-isms are sometimes confusing.\n",
    "    - `nn.Softmax` is a `nn.Module` (think of it as a function).\n",
    "    - When the instance of this module (e.g. function) is \"called\" with a `tensor`, it relays the call (and the argument) to `torch.nn.functional.softmax`\n",
    "    - `nn.functional.softmax` calls out to the `C++` back end, which does the actual work!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2876f25-e2b5-48b6-8828-4a8f924024d4",
   "metadata": {},
   "source": [
    "## So WTF is **logits** ??\n",
    "  - It's similar in concept to the `N` matrix above. But instead of counts, think of it as being something close to $log(probability)$ or $log(counts)$\n",
    "  - Why do we do this? because keeping counts using floats is harder, and it's much easier to start from values in a small range (say 0 to 1.0) for numerical stability\n",
    "  - Think of `logits` as a weird mangling of *log-probability*\n",
    "  -  Historically, the tensor (which is an output of some layer of a neural network), which is fed into the `Softmax()` function is called `logits`\n",
    "  -  Because neural networks like to work with small numbers between -1.0 to 1.0, it's cumbersome to represents counts of events directly.\n",
    "  -  Instead, we pretend that these small values are log(probabilities), so that we can push them into Softmax to turn the logits into a tensor of probabilities"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e43de29-2169-4f5c-917b-cc457620c0e2",
   "metadata": {},
   "source": [
    "### Cross Entropy\n",
    "  - This is used during the training phase to calculate the \"error\", usually called `Loss` \n",
    "  - log(a*b*c) = log(a) + log(b) + log(c)\n",
    "  - GOAL: maxmimize likelyhood of the data w.r.t model parameters (statistically modeling)\n",
    "  - Equivalent to maximizing the log likelyhood (because log is monotonic)\n",
    "  - Equivalent to minimizing the negative log likelyhood\n",
    "  - Equivalent to minimizing the average negative log likelyhood\n",
    "  - Assume you have a int tensor Y which is of shape (B) (set B=1 for simple case)\n",
    "  - Assume you have a logits vector of shape (B, R) where R is the number of tokens (or \"classes\") \n",
    "  - Then, `counts = logits.exp()` (elementwise exponentiation)\n",
    "  - Then, `probabilities = counts/counts.sum(...)` (elementwise dividing counts by the total number of \"events\")\n",
    "  - Assume you have expected answers `Y` which is a tensor of integers, each denoting a specific token\n",
    "  - ``` \n",
    "    def CrossEntropy(logits, Y):\n",
    "      B = Y.size()[0]\n",
    "      counts = logits.exp()  # akin to N above\n",
    "      probs = counts / counts.sum(1, keepdims=True) # this is the result of softmax!\n",
    "      loss = -probs[torch.arange(B), Y].log().mean()\n",
    "    ```\n",
    "  - So what is `probs[torch.arange(B), Y]` doing?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4cfa71e-a122-46bb-9ab6-c1aedec92f84",
   "metadata": {},
   "source": [
    "### one-hot-encoding\n",
    "  - When you have an ordered list of things that you can choose from, you can take advantage of pytorch to specify which item you want.\n",
    "  - Assume you have $R$ elements you can choose from. You can specify which one by using the *one hot* encoding.\n",
    "  - This one hot vector can be used to select the `j`th row or column from a matrix (of shape (R, C) by doing\n",
    "  - `OneHot = F.one_hot(torch.tensor(j), num_classes=R)`\n",
    "  - `SelectedRow = OneHot @ Embed` (  R @ (R,C) --> (1,R) @  (R,C) = (1,C) --> tensor of shape (C)   )\n",
    "  - You can select the `j`th column by `Embed @ OneHot` ( (R,C) @ (C) --> (R,C) @ (C,1) = (R,1) --> tensor of shape (R) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bd96a73-bf35-4220-9ce9-8e7f8c604efd",
   "metadata": {},
   "outputs": [],
   "source": [
    "AR =  torch.arange(0, 5)\n",
    "OH = F.one_hot(AR % 3)\n",
    "OH2 = F.one_hot(AR % 3, num_classes=5)\n",
    "print(vprt(AR=AR, Results=hprt(OH=OH, OH2=OH2)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ddfbad8-7232-407f-86d3-38618738f4bf",
   "metadata": {},
   "source": [
    "### Embedding\n",
    "  - <https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html#torch.nn.Embedding>\n",
    "  - In LLMs information about a *token* (think of it as a word, for now) is kept in a large 1-D tensor.\n",
    "  - Think of it as either a line from the origin pointing out to an n-dimensional space, or a point in an n-dimensional space\n",
    "  - LLMs typically have anywhere from 30000 to ~120K tokens (words, or some similiar concept) that it knows about\n",
    "  - Resulting in a *Embedding* (a tensor) of shape `(nVocab, nEmbed)`\n",
    "    - `Embedding` matrix is typically in row-major memory layout\n",
    "  - `nEmbed` is typically at least 768 (for GPT2, BERT, etc..), 4096 for LLama(1,2,3)\n",
    "  - `nVocab` is usually at least 32000"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a1bfc08-1751-4e63-8311-0fb7eedfca0a",
   "metadata": {},
   "source": [
    "### The Role of the Tokenizer in LLMs\n",
    "  - Tokenizer turns a sentence into a sequence of integer ids \n",
    "  - This integer is then used to select the embedding tensor to get the `embedding vector` for that token\n",
    "  - Each word (e.g. token) in the input sentence is thus transformed into its embedding vector.\n",
    "  - These vectors (which represents the input words) is then fed into the neural network!\n",
    "  - There are complexities on how many tokens are fed in to the NN at once (more on that later!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a6f9ed8-b372-4b23-9008-f801d1686738",
   "metadata": {},
   "outputs": [],
   "source": [
    "nVocab, nEmbed  = (5,2)\n",
    "Vocab = ['a', 'b', 'c', 'd', '.' ] # this is a really stupid language!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58647867-f964-445c-9578-862ce56ae48d",
   "metadata": {},
   "outputs": [],
   "source": [
    "Embed1 = torch.arange(nEmbed*nVocab).view(nVocab,-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfc411f9-83bb-481e-a71d-d2e403427a05",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(hprt(Embed1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1ef16d7-6b36-491d-8148-e43c5a7f9820",
   "metadata": {},
   "outputs": [],
   "source": [
    "nBatch=3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dfea2c2-5bc0-4943-9a56-c3df84e1443e",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(1337)\n",
    "Embed = torch.randn((nVocab, nEmbed))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad2bdf2d-78a2-4bd5-a466-bb2783729acf",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(hprt(Embed))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de9a0a7b-1a11-4bb0-8aa5-b5ea8ac5cdd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dumTok(inp):\n",
    "  return torch.tensor([ stoi[k] for k in inp ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "147e851f-e7ec-42c3-8bb0-e4b4d754ec1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "Inp = dumTok('abc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ebc983b-e6e7-4ca1-8211-056403599d16",
   "metadata": {},
   "outputs": [],
   "source": [
    "Inp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80f1276c-f1c3-4b1c-a264-c3c4ac1b8df1",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = Embed[Inp]\n",
    "# the above accessing Embed by [Inp[0]], [Inp[1]], [Inp[2]]\n",
    "# i.e. row[Inp[0]], row[Inp[1]], row[Inp[2]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "348d3f75-bdd4-48d1-a188-fcbbb46fa7c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(printTensors(inp=Inp.view(-1,1), Embed=Embed, x=x))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26ac6f63-f249-4182-9069-54dcc12a02ed",
   "metadata": {},
   "source": [
    "## The above vector `x` is typically fed into the neural network"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96034f60-71a8-4288-b0b8-3933030eebe4",
   "metadata": {},
   "source": [
    "## Now that we've seen what goes INTO a LLM, let's do a simple example of what happens at the output end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "180fa69f-1822-40ed-b68b-f6892bf3cc0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(42)\n",
    "logits = torch.randn((nBatch, nVocab))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bed49d13-b311-46d8-923f-1e0a6d2c5a4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(printTensors(logits=logits))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c28450c7-1f5a-42a0-a2bf-00cbddd8cc86",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(printTensors(softmax=F.softmax(logits, dim=1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afd4c1d8-5b86-4892-bbfa-eb8feb4b027c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# lets assume Y is to be 3 for all three cases\n",
    "Y = torch.tensor([3,3,3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5f20f34-30f7-48a3-8d3d-4afc10ac0b3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = F.cross_entropy(logits, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87897d4d-3ce0-4431-9d8c-a932e3ce1125",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1935e3f3-b67c-45f2-870a-e10cc0dcf957",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts = logits.exp()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ed367aa-7f8d-4f2f-9549-c300d1ed10b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "probs = counts / counts.sum(1,keepdims=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6e55837-1fba-45c4-9e31-8150d37a1894",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(hprt(probs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64b55434-9840-41b8-b042-52140862b57f",
   "metadata": {},
   "outputs": [],
   "source": [
    "CC = torch.arange(nBatch)\n",
    "ACC = probs[CC, Y]\n",
    "## the above It is accessing probs at [CC[0], Y[0]], [CC[1], Y[1]], [CC[2], Y[2]]\n",
    "## and generaring a tensor of shape (nBatch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b192c1eb-95e1-4e83-80aa-7cc1b4353384",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(printTensors(CC=CC,ACC=ACC))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d8825ac-e145-4e1f-bb3c-3b6fefcfcf87",
   "metadata": {},
   "outputs": [],
   "source": [
    "ACC.log()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a69bac6-af1e-429c-90d9-df8b501bf8fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "ACC.log().mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbc2de73-06cb-4d75-ae63-b389aa0acabd",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss2 = -ACC.log().mean().item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb59f9e2-0752-4e5c-b742-91a9bdd62da1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# floating point roundoff between F.cross_entropy() and our test\n",
    "print(loss, loss2, (loss -loss2).item())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "504eae10-050d-46c3-bd77-0b4d4e89b8cb",
   "metadata": {},
   "source": [
    "## Plot of the -log(probability) from 0.0001 to 1.0\n",
    "### loss of 1.0 is at probility of $\\frac{1}{e}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb9aec47-b5b9-411f-9bc4-4d37173f177a",
   "metadata": {},
   "outputs": [],
   "source": [
    "delta=0.0001\n",
    "xs = np.arange(delta, 1.0, delta)\n",
    "ys = -np.log(xs)\n",
    "fig, ax = plt.subplots()\n",
    "line = ax.plot(xs, ys)\n",
    "#ax.axis('equal')\n",
    "ax.grid(True, which='both')\n",
    "plt.title('Probability vs -log(P)')\n",
    "plt.ylabel('-log(probability)')\n",
    "plt.xlabel('probability 0.0 - 1.0')\n",
    "ax.axhline(y=1.0, color='y')\n",
    "ax.axvline(x=1/math.e, color='red')\n",
    "ax.axhline(y=0, color='g')\n",
    "ax.axvline(x=0, color='g')\n",
    "ax.annotate(f'loss=1.0', xy=(1/math.e, 1.0),  xycoords='data', xytext=(0.2, 0.5), textcoords='axes fraction',\n",
    "            arrowprops=dict(facecolor='black', shrink=0.05,width=0.1,headwidth=5.0),\n",
    "            horizontalalignment='right',\n",
    "            verticalalignment='bottom')\n",
    "ax.annotate(f'x=1/{math.e:0.4f}={1/math.e:0.4f}', xy=(1/math.e, 1.0),  xycoords='data', xytext=(.9, 0.5), textcoords='axes fraction',\n",
    "            arrowprops=dict(facecolor='black', shrink=0.05,width=0.1, headwidth=5.0),\n",
    "            horizontalalignment='right',\n",
    "            verticalalignment='bottom')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "703174c4-a3e0-4519-ac25-c6d145292be1",
   "metadata": {},
   "source": [
    "### So why cross entropy?\n",
    "  - LLMs are fed in one or more tensors picked from the embedding matrix.\n",
    "  - It then produces the ouput tensor called `logits`\n",
    "  - LLMs have a number of Tokens that it can produce at any given step (classification!)\n",
    "  - to pick one, it calls `multinomial(softmax(logits), num_classes=nVocab, replace=True)`\n",
    "  - `Softmax` returns a tensor of probabilities, ranging from (0 .. 1.0)\n",
    "  - `multinomial` randomly picks a token from the probability tensor (which is of shape (nBatch, nVocab)\n",
    "### Characteristics of Cross Entropy\n",
    "  - Arbitrarily assigns $P(\\displaystyle \\frac{1}{e})$ as the nominal `1.0` loss\n",
    "  - Softmax takes care of normalizing the logits.\n",
    "  - Is the loss related to which class has the max probability? (no, not really)\n",
    "  - Does not work well if number of classes are few!\n",
    "  - IMPORTANT! We don't have the **TRUE\"** (or desired!) probability distribution! (i.e. the *expected* probability distribution.\n",
    "  - All we have is the **ACTUAL** probability distribution (from the logits!)\n",
    "  - Cross Entropy Loss measures the overall likelyhood of the $i$th token being the correct value.\n",
    "    - It doesn't take into account things like whether the $i$th generated token was the most likely (i.e. had the highest probability)\n",
    "### Strange issue with LLMs\n",
    "  - Because of multinomial, the NN can 'accidentally' choose either the right answer or the `nVocab-1` *wrong* answers at each time step.\n",
    "  - However! Not always the case that choosing the \"best\" token at each step always results in the best answer!\n",
    "  - In this case, CrossEntropyLoss is the best we can do!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28ecf670-4d67-44b2-95f1-4a6ea89ae0d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# simple case, only one row in logits\n",
    "aa = torch.tensor([[ 0, 1, 1, 0, 0.0]])\n",
    "aaP = F.softmax(aa, dim=1)\n",
    "bb = torch.tensor([0]).long()\n",
    "cc = F.cross_entropy(aa, bb)\n",
    "print(printTensors(CE=cc, softmax=aaP, log=aaP.log())) # sum=aaP.sum(), "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4ae1583-179f-4f19-9bd4-003a701a4495",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example of target with multiple rows in logits\n",
    "torch.manual_seed(86)\n",
    "logits = torch.randn(nBatch, nVocab, requires_grad=True)\n",
    "target = torch.empty(nBatch, dtype=torch.long).random_(5)\n",
    "print(hprt(target=target,logits=logits, softmax=F.softmax(logits, dim=1)))\n",
    "loss11 = F.cross_entropy(logits, target)\n",
    "SV=F.softmax(logits, dim=1)\n",
    "Sel=SV[torch.arange(nBatch), target]\n",
    "print(printTensors(loss=loss11, loss2=-Sel.log().mean(), Sel=Sel))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13c9da04-6613-422a-b28a-6ef3ba5471cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "-np.log(1/2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06235620-dc03-45e0-9aac-fb0a3238058f",
   "metadata": {},
   "outputs": [],
   "source": [
    "-math.log(1/3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d79b706e-435c-4b98-b844-36b8d9fc803f",
   "metadata": {},
   "outputs": [],
   "source": [
    "-math.log(1/4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd9b8a17-74a8-4654-b7b2-3f8dad448d5a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}