bigram.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bigram\n",
    "\n",
    "Bigram are the precursor to transformer models. These are models where only the previous word is used to predict the next word. For example, `\"I am a man\"` when you are training the `am` it only gets `I` as input. Same with `a` it only gets `am` as input.\n",
    "\n",
    "Thi is known as the Markov assumption, where the probability of a current word depends only on the previous word.\n",
    "\n",
    "They are calculated by analyzing a large text corpus and counting the frequency of each bigram. This data is used to estimate the conditional probailites P(word 2| word 1)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "torch.manual_seed(1337)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "batch_size = 32\n",
    "block_size = 8\n",
    "max_iters = 3000 \n",
    "eval_interval = 300\n",
    "learning_rate = 1e-2\n",
    "device = \"mps\"\n",
    "eval_iters = 200\n",
    "n_embd = 384 \n",
    "n_head = 6\n",
    "n_layer = 6\n",
    "dropout = 0.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"./data/karpathy/input.txt\", \"r\", encoding=\"utf-8\") as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# char tokenizer\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "\n",
    "stoi = {ch:i for i,ch in enumerate(chars)}\n",
    "itos = {i:ch for i,ch in enumerate(chars)}\n",
    "encode = lambda s: [stoi[c] for c in s]\n",
    "decode = lambda l: ''.join([itos[i] for i in l])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train and test splits\n",
    "data = torch.tensor(encode(text), dtype=torch.long)\n",
    "n = int(0.9*len(data))\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_batch(split):\n",
    "    data = train_data if split == \"train\" else val_data\n",
    "    ix = torch.randint(len(data) - block_size, (batch_size,))\n",
    "    x = torch.stack([data[i:i+block_size] for i in ix])\n",
    "    y = torch.stack([data[i+1:i+block_size + 1] for i in ix])\n",
    "    x,y = x.to(device), y.to(device)\n",
    "    return x,y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BigramLanguageModel(nn.Module):\n",
    "    def __init__(self, vocab_size):\n",
    "        super().__init__()\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        logits = self.token_embedding_table(idx)\n",
    "\n",
    "        if targets is None: \n",
    "            loss = None\n",
    "        else: \n",
    "            B,T,C = logits.shape\n",
    "            logits = logits.view(B*T, C)\n",
    "            targets = targets.view(B*T)\n",
    "            loss = F.cross_entropy(logits, targets)\n",
    "        \n",
    "        return logits, loss\n",
    "    \n",
    "    def generate(self, idx, max_new_tokens):\n",
    "        for _ in range(max_new_tokens):\n",
    "            logits, loss = self(idx)\n",
    "            logits = logits[:, -1, :] # becomes (B, C)\n",
    "            probs = F.softmax(logits, dim=-1) # (B, C)\n",
    "            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
    "            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)\n",
    "        return idx\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = BigramLanguageModel(vocab_size)\n",
    "m = model.to(device)\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def estimate_loss():\n",
    "    out = {}\n",
    "    model.eval()\n",
    "    for split in [\"train\", \"val\"]:\n",
    "        losses = torch.zeros(eval_iters)\n",
    "        for k in range(eval_iters):\n",
    "            X, Y = get_batch(split)\n",
    "            _, loss = model(X,Y)\n",
    "            losses[k] = loss.item()\n",
    "        out[split] = losses.mean()\n",
    "    model.train()\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for iter in range(max_iters):\n",
    "    if iter % eval_interval == 0:\n",
    "        losses = estimate_loss()\n",
    "        print(f\"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
    "    \n",
    "    xb, yb = get_batch(\"train\")\n",
    "\n",
    "    logits, loss = m(xb,yb)\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    loss.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context = torch.zeros((1,1), dtype=torch.long, device=device)\n",
    "print(decode(model.generate(context, max_new_tokens=500)[0].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(1337)\n",
    "B, T, C = 4,8,32\n",
    "x = torch.randn(B, T, C)\n",
    "\n",
    "head_size = 16\n",
    "key = nn.Linear(C, head_size, bias=False)\n",
    "query = nn.Linear(C, head_size, bias=False)\n",
    "value = nn.Linear(C, head_size, bias=False)\n",
    "\n",
    "# no communication has happened yet\n",
    "k = key(x)\n",
    "q = query(x)\n",
    "\n",
    "# communication\n",
    "wei = q @ k.transpose(-2,-1)\n",
    "\n",
    "tril = torch.tril(torch.ones(T, T))\n",
    "#wei = torch.zeros((T, T))\n",
    "# In encoder blocks you would delete this. Only in decoder blocks would you want the autoregressive format.\n",
    "wei = wei.masked_fill(tril==0, float(\"-inf\"))\n",
    "wei = F.softmax(wei, dim=-1)\n",
    "\n",
    "v = value(x)\n",
    "out = wei @ v\n",
    "#out = wei @ x\n",
    "\n",
    "out.shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}