{"cells":[{"cell_type":"markdown","metadata":{"id":"bD24MV9XpO00"},"source":["# Bigram Model\n","\n","<table align=\"left\">\n","  <td>\n","    <a href=\"https://colab.research.google.com/github/SkillSurf/introduction_genAI/blob/main/notebooks/S1/Session_1_1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n","  </td>\n","</table>\n","\n","<br />\n","<br />\n","\n","This is one of the simplest generative models. It predicts the next character (letter/symbol) based only on the 1 past character. This is like an extremely simple mobile keyboard auto-complete.\n","\n","\n","#### Tokens & Vocabulary\n","\n","The text is first \"tokenized\" by simply assigning the numbers 0-64 to the uppercase & lowercase letters of the English alphabet and punctuation symbols. This mapping is called vocabulary. Since there are 65 such unique tokens, the vocabulary size is 65.\n","\n","\n","#### Embedding Table\n","\n","The model has only one layer, called an \"embedding table\", that acts as a trainable lookup table. The embedding table is a (65 x 65) sized matrix.\n","\n","Let's say the past token is character \"C\" (uppercase C), which is assigned 15 in the vocabulary. When 15 is given as the input, the embedding table gives the 15th row, a vector of legnth 65 as the output. Each of the 65 elements in this output corresponds to the probability of the next token. If the 13th element has highest value, the output is 13 = \"A\".\n","\n","### Training\n","\n","We train on the \"Tiny Shakespeare\" dataset, that contains 40,000 lines of Shakespeare plays. Given each token, we let the model guess the next token and then calculate the cross-entropy loss: $-Σ t_i \\log(y_i)$, where $t_i$ is the ground truth (actual next token from dataset), and $y_i$ is the token predicted by the model.\n","\n","### Results\n","\n","We can see the loss reducing with training iterations, and the resulting model generates text that somewhat resembles actual text, not pure random strings. But the resulting text is incomprehensible, since this model is too simple.\n","\n","### Limitation\n","\n","This model only looks at a single past token to predict the next token, which makes it extremely short-sighted. This is equal to context window legnth of 1. We will improve on this in the next models."]},{"cell_type":"code","execution_count":1,"metadata":{"executionInfo":{"elapsed":2024,"status":"ok","timestamp":1717896185121,"user":{"displayName":"Abarajithan Gnaneswaran","userId":"15065416675145289008"},"user_tz":420},"id":"Hr63nefWmXIE"},"outputs":[],"source":["import torch\n","import torch.nn as nn\n","from torch.nn import functional as F"]},{"cell_type":"markdown","metadata":{"id":"ngT-59G5m2TM"},"source":["### Hyperparameters"]},{"cell_type":"code","execution_count":2,"metadata":{"executionInfo":{"elapsed":3,"status":"ok","timestamp":1717896185121,"user":{"displayName":"Abarajithan Gnaneswaran","userId":"15065416675145289008"},"user_tz":420},"id":"6kqCFiSLmztM"},"outputs":[],"source":["B = 32 # batch size: how many independent sequences will we process in parallel?\n","T = 1  # time: what is the maximum context length for predictions?\n","max_iters = 3000\n","eval_interval = 300\n","learning_rate = 1e-2\n","device = 'cuda' if torch.cuda.is_available() else 'cpu'\n","eval_iters = 200"]},{"cell_type":"markdown","metadata":{"id":"uImQJNS8ojj1"},"source":["### Dataset"]},{"cell_type":"code","execution_count":3,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":454,"status":"ok","timestamp":1717896185573,"user":{"displayName":"Abarajithan Gnaneswaran","userId":"15065416675145289008"},"user_tz":420},"id":"XsK-a5jQnMkl","outputId":"2b0afdae-8968-4d4c-f904-6891f6c1e78c"},"outputs":[{"name":"stdout","output_type":"stream","text":["--2024-06-09 01:23:03--  https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\n","Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.111.133, 185.199.110.133, ...\n","Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n","HTTP request sent, awaiting response... 200 OK\n","Length: 1115394 (1.1M) [text/plain]\n","Saving to: ‘input.txt.2’\n","\n","\rinput.txt.2           0%[                    ]       0  --.-KB/s               \rinput.txt.2         100%[===================>]   1.06M  --.-KB/s    in 0.02s   \n","\n","2024-06-09 01:23:03 (43.7 MB/s) - ‘input.txt.2’ saved [1115394/1115394]\n","\n"]}],"source":["!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"]},{"cell_type":"code","execution_count":4,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":4,"status":"ok","timestamp":1717896185573,"user":{"displayName":"Abarajithan Gnaneswaran","userId":"15065416675145289008"},"user_tz":420},"id":"bfzVx8bHm1IG","outputId":"42f0bf01-a84d-4001-9088-7465f6b40160"},"outputs":[{"name":"stdout","output_type":"stream","text":["vocab_size: 65\n","vocabulary: \n"," !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n"]}],"source":["torch.manual_seed(1337)\n","\n","with open('input.txt', 'r', encoding='utf-8') as f:\n","    text = f.read()\n","\n","# here are all the unique characters that occur in this text\n","chars = sorted(list(set(text)))\n","vocab_size = len(chars)\n","# create a mapping from characters to integers\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] # encoder: take a string, output a list of integers\n","decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string\n","\n","chars_str = ''.join(chars)\n","print(f'vocab_size: {vocab_size}')\n","print(f'vocabulary: {chars_str}')\n"]},{"cell_type":"code","execution_count":5,"metadata":{"executionInfo":{"elapsed":719,"status":"ok","timestamp":1717896186290,"user":{"displayName":"Abarajithan Gnaneswaran","userId":"15065416675145289008"},"user_tz":420},"id":"fL_iTO4vm7vZ"},"outputs":[],"source":["# Train and test splits\n","data = torch.tensor(encode(text), dtype=torch.long)\n","n = int(0.9*len(data)) # first 90% will be train, rest val\n","train_data = data[:n]\n","val_data = data[n:]\n","\n","\n","# data loading\n","def get_batch(split):\n","    # generate a small batch of data of inputs x and targets y\n","    data = train_data if split == 'train' else val_data\n","    ix = torch.randint(len(data) - T, (T,))\n","    x = torch.stack([data[i:i+T] for i in ix])\n","    y = torch.stack([data[i+1:i+T+1] for i in ix])\n","    x, y = x.to(device), y.to(device)\n","    return x, y\n"]},{"cell_type":"markdown","metadata":{"id":"IpNx_DY-ooR2"},"source":["### Model"]},{"cell_type":"code","execution_count":6,"metadata":{"executionInfo":{"elapsed":3,"status":"ok","timestamp":1717896186291,"user":{"displayName":"Abarajithan Gnaneswaran","userId":"15065416675145289008"},"user_tz":420},"id":"l6bQEIwKoNDd"},"outputs":[],"source":["class BigramLanguageModel(nn.Module):\n","\n","    def __init__(self):\n","        super().__init__()\n","        # each token directly reads off the logits for the next token from a lookup table\n","        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size) # for every possible token, weights for next token\n","\n","    def forward(self, idx, targets=None):\n","\n","        '''\n","        B - batch                   # of independant vectors processed\n","        T - time/block/context      # of tokens in a context\n","        C - channels/dimensionality\n","        '''\n","\n","        # idx and targets are both (B,T) tensor of integers\n","        logits = self.token_embedding_table(idx) # (B,T)\n","\n","        if targets is None:\n","            loss = None\n","        else:\n","            B, T, C = logits.shape # C = vocab_size\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):                        # idx is (B, T) array of indices in the current context\n","            logits, loss = self(idx)                           # get the predictions\n","            logits = logits[:, -1, :]                          # (B,T,C) -> (B, C)\n","            probs = F.softmax(logits, dim=-1)                  # (B, C)\n","            idx_next = torch.multinomial(probs, num_samples=1) # sample from the distribution acc to prob (B, 1)\n","            idx = torch.cat((idx, idx_next), dim=1)            # New idx is concat (B, T+1)\n","        return idx\n","\n","model = BigramLanguageModel()\n","m = model.to(device)"]},{"cell_type":"code","execution_count":7,"metadata":{"executionInfo":{"elapsed":2674,"status":"ok","timestamp":1717896188963,"user":{"displayName":"Abarajithan Gnaneswaran","userId":"15065416675145289008"},"user_tz":420},"id":"z8UZVObFoQUS"},"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","            logits, loss = model(X, Y)\n","            losses[k] = loss.item()\n","        out[split] = losses.mean()\n","    model.train()\n","    return out\n","\n","optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n"]},{"cell_type":"markdown","metadata":{"id":"AR2Vnvf9orK2"},"source":["### Training"]},{"cell_type":"code","execution_count":8,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":5190,"status":"ok","timestamp":1717896194147,"user":{"displayName":"Abarajithan Gnaneswaran","userId":"15065416675145289008"},"user_tz":420},"id":"Vd1wg5JloS5N","outputId":"726801c9-b8cf-4304-9a17-bd3c777e50bf"},"outputs":[{"name":"stdout","output_type":"stream","text":["step 0: train loss 4.7220, val loss 4.7486\n","step 300: train loss 4.1878, val loss 4.3036\n","step 600: train loss 3.9232, val loss 4.0731\n","step 900: train loss 3.7221, val loss 3.7067\n","step 1200: train loss 3.3978, val loss 3.3778\n","step 1500: train loss 3.1350, val loss 3.1010\n","step 1800: train loss 3.1643, val loss 3.1060\n","step 2100: train loss 2.9825, val loss 3.1367\n","step 2400: train loss 3.0083, val loss 2.9184\n","step 2700: train loss 2.8582, val loss 2.9125\n"]}],"source":["for iter in range(max_iters):\n","    if iter % eval_interval == 0:   # every once in a while evaluate the loss on train and val sets\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')     # sample a batch of data\n","\n","    # evaluate the loss\n","    logits, loss = model(xb, yb)\n","    optimizer.zero_grad(set_to_none=True)\n","    loss.backward()\n","    optimizer.step()"]},{"cell_type":"markdown","metadata":{"id":"-GioBti5otI8"},"source":["### Inference"]},{"cell_type":"code","execution_count":9,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":452,"status":"ok","timestamp":1717896194596,"user":{"displayName":"Abarajithan Gnaneswaran","userId":"15065416675145289008"},"user_tz":420},"id":"5PUrGWP-oW7k","outputId":"879af13f-a980-4185-b74f-f9e2b06fac90"},"outputs":[{"name":"stdout","output_type":"stream","text":["\n","S!Emh.. Acfomvid bno.:CUf:d:ghelofad,SVY!ugdathe t ypcxFbyfofineQbmu I&yene d tRFZGNGoufMkHxKj?s yvndinEUzrierGlo m mend.\n","PER:L;ves sBjNMNWWP, ybre,\n","TAdbrengotonol:LbHWz?YCHxrelo in tingisof s, nWbs!oulu y he i, ndanves th sc!$RrvirtR\n","Nf?e, dosxToutimsREs. fad, avYW:CavwivoTxEkendouZ?3KH--RI, ase 3!aists.jxjthsthandt tt'godan: IU!zgu:ztgmno VZYog t f r toitoulpqp CHUDMkXLvOvenXme:gxUGTO:Ml!BNNUFinds h ile IBave; bymeare.\n","Mlou thagis thhe, houXCCoingmof to lyZLCcqbye sisimanereagCad rS&ouRorar;ng\n"]}],"source":["context = torch.zeros((1, 1), dtype=torch.long, device=device)  # start with '\\n' as seed\n","out_ints = m.generate(context, max_new_tokens=500)[0].tolist() # output list of ints\n","print(decode(out_ints))"]}],"metadata":{"colab":{"authorship_tag":"ABX9TyODG3Ra9eU8OuhtThKpgm5h","provenance":[]},"kernelspec":{"display_name":"Python 3","name":"python3"},"language_info":{"name":"python"}},"nbformat":4,"nbformat_minor":0}