decoder_gpt.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([1.], device='mps:0')\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn \n",
    "from torch.nn import functional as F\n",
    "\n",
    "\n",
    "\n",
    "if torch.backends.mps.is_available():\n",
    "    mps_device = torch.device(\"mps\")\n",
    "    x = torch.ones(1, device=mps_device)\n",
    "    print (x)\n",
    "else:\n",
    "    print (\"MPS device not found.\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "### paramters\n",
    "block_size =64\n",
    "batch_size =128\n",
    "max_iterations = 1000\n",
    "learning_rate =3e-4\n",
    "eval_iters = 250\n",
    "n_embd = 384\n",
    "n_head = 4\n",
    "n_layer =4\n",
    "dropout = 0.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "char =\"\"\n",
    "\n",
    "with open('wizard_of_oz.txt','r',encoding='utf-8') as f:\n",
    "    text = f.read()\n",
    "    chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "### character tokenizer from scratch \n",
    "string_to_int ={ch:i for i, ch in enumerate(chars)} ## mapping aka dicts\n",
    "int_to_string ={i:ch for i,ch in enumerate(chars)}\n",
    "encode = lambda s: [string_to_int[c] for c in s]\n",
    "decode = lambda l: ''.join([int_to_string[i] for i in l])\n",
    "## using pytorch tensors to encode the entire text\n",
    "data =torch.tensor(encode(text), dtype= torch.long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inputs:\n",
      " tensor([[50, 67, 64,  ..., 65, 54, 50],\n",
      "        [53,  1, 64,  ...,  1, 58, 63],\n",
      "        [56,  1, 70,  ..., 72, 50, 68],\n",
      "        ...,\n",
      "        [53,  1, 50,  ..., 58, 54, 53],\n",
      "        [53,  1, 57,  ..., 57, 54, 62],\n",
      "        [ 6,  1, 69,  ..., 55, 70, 69]], device='mps:0')\n",
      "traget:\n",
      " tensor([[67, 64, 70,  ..., 54, 50, 67],\n",
      "        [ 1, 64, 63,  ..., 58, 63, 66],\n",
      "        [ 1, 70, 65,  ..., 50, 68,  1],\n",
      "        ...,\n",
      "        [ 1, 50, 63,  ..., 54, 53,  1],\n",
      "        [ 1, 57, 50,  ..., 54, 62,  1],\n",
      "        [ 1, 69, 57,  ..., 70, 69, 70]], device='mps:0')\n"
     ]
    }
   ],
   "source": [
    "n =int(0.8*len(data))\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]\n",
    "\n",
    "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",
    "    #print(ix)\n",
    "    x = torch.stack([data[d:d+block_size] for d in ix ])\n",
    "    y = torch.stack([data[d+1:d+block_size+1] for d in ix])\n",
    "    x,y = x.to(mps_device),y.to(mps_device)\n",
    "    return x,y\n",
    "\n",
    "x,y = get_batch('train')\n",
    "print(f\"inputs:\\n {x}\")\n",
    "print(f\"traget:\\n {y}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Head(nn.Module):\n",
    "    \"\"\"one head self attention\"\"\"\n",
    "    def __init__(self, head_size):\n",
    "        super().__init__()\n",
    "        self.key = nn.Linear(n_embd,head_size,bias=False)\n",
    "        self.query = nn.Linear(n_embd,head_size, bias=False)\n",
    "        self.value = nn.Linear(n_embd,head_size, bias=False)\n",
    "        self.register_buffer('tril', torch.tril(torch.ones(block_size,block_size))) ## register in model state can be removed but will take longer to train\n",
    "\n",
    "        self.dropout= nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # input of size (batch, time-step, channels)  \n",
    "        # output of size (batch, time-step, head size)\n",
    "        B,T,C = x.shape\n",
    "        k = self.key(x)   # (B,T,hs)\n",
    "        q = self.query(x) # (B,T,hs)\n",
    "        # compute attention scores (\"affinities\")\n",
    "        wei = q @ k.transpose(-2,-1) * k.shape[-1]**-0.5 # (B, T, hs) @ (B, hs, T) -> (B, T, T) ### scaling mathhhhh\n",
    "        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)\n",
    "        wei = F.softmax(wei, dim=-1) # (B, T, T)\n",
    "        wei = self.dropout(wei)\n",
    "        # perform the weighted aggregation of the values\n",
    "        v = self.value(x) # (B,T,hs)\n",
    "        out = wei @ v # (B, T, T) @ (B, T, hs) -> (B, T, hs)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    '''multiple heads of self-attenton parralel '''\n",
    "    def __init__(self, num_heads,head_size):\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)]) ## heads in parallel \n",
    "        self.proj = nn.Linear(head_size * num_heads, n_embd) ## project head size ## changeable to not throw dimintionality errors\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self,x):\n",
    "        out = torch.cat([h(x) for h in self.heads], dim = -1) ## \n",
    "        out = self.dropout(self.proj(out))\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    '''linear layer followed by a non-linearity'''\n",
    "\n",
    "    def __init__(self, n_embd):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(n_embd,4*n_embd),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(4*n_embd,n_embd),\n",
    "            nn.Dropout(dropout),\n",
    "\n",
    "        )\n",
    "    def forward(self,x):\n",
    "        return self.net(x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    \"\"\" Transformer block: communications followed by computation \"\"\"\n",
    "\n",
    "    def __init__(self,n_embd,n_head):\n",
    "        ## n_embd: embedding dimnsion , N_head : the number of heads we'd like\n",
    "        super().__init__()\n",
    "        head_size  =n_embd // n_head\n",
    "        self.sa = MultiHeadAttention(n_head, head_size) ## self attention\n",
    "        self.ffwd = FeedForward(n_embd) \n",
    "        self.ln1 = nn.LayerNorm(n_embd) ## post norm \n",
    "        self.ln2 = nn.LayerNorm(n_embd)\n",
    "    def forward (self,x):\n",
    "        y = self.sa(x)\n",
    "        x = self.ln1(x+y)\n",
    "        y= self.ffwd(x)\n",
    "        x = self.ln2(x + y)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "### neural network \n",
    "class GPTModel(nn.Module):\n",
    "    def __init__(self, vocab_size):\n",
    "        super().__init__()\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)\n",
    "        self.position_embedding_table = nn.Embedding(block_size, n_embd)\n",
    "        self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])\n",
    "        self.ln_f = nn.LayerNorm(n_embd) # final layer norm\n",
    "        self.lm_head = nn.Linear(n_embd, vocab_size)\n",
    "        ### std used by professinoals mostly understand it as much as you can\n",
    "        ## basically is it weight and bias initalizing using std \n",
    "        ## for trainig to converge better\n",
    "        self.apply(self._init_weights)\n",
    "\n",
    "    def _init_weights(self, module):\n",
    "        if isinstance(module, nn.Linear):\n",
    "            torch.nn.init.normal_(module.weight, mean = 0.0, std =0.02)\n",
    "            if module.bias is not None:\n",
    "                torch.nn.init.zeros_(module.bias)\n",
    "        elif isinstance(module, nn.Embedding):\n",
    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "\n",
    "\n",
    "\n",
    "    def forward(self, index, target= None):\n",
    "        B, T = index.shape\n",
    "\n",
    "        ## idx and targets are both (B,T) tensor of integers\n",
    "        tok_emb = self.token_embedding_table(index) # token_emedding_table\n",
    "        pos_emb = self.position_embedding_table(torch.arange(T, device =mps_device)) ## how long is T then give a number here incdies 8\n",
    "        x = tok_emb + pos_emb ## adding them together \n",
    "        x = self.blocks(x)\n",
    "        x = self.ln_f(x)\n",
    "        logits = self.lm_head(x)\n",
    "        \n",
    "\n",
    "        if target is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            B, T, C =logits.shape ## batch / time / \n",
    "            logits = logits.view(B*T, C)\n",
    "            target = target.view(B*T) ## reshaping bc cross entropy \n",
    "            loss = F.cross_entropy(logits, target)\n",
    "        return logits, loss\n",
    "    \n",
    "    def generate(self, index, max_new_tokens):\n",
    "        ## starting here the index is (B,T) array of indices \n",
    "        for _ in range(max_new_tokens):\n",
    "            ## getting predictions\n",
    "            logits, loss = self.forward(index)\n",
    "            logits = logits[:,-1,:] ## to get previous step aw last since bigram model\n",
    "            probs = F.softmax(logits, dim=-1) ## softmax function\n",
    "            index_next = torch.multinomial(probs, num_samples=1)#sample from distribution(B,1)\n",
    "            index = torch.cat((index, index_next), dim=1) ## adding the prev and next index matrix\n",
    "        return index\n",
    "    \n",
    "model = GPTModel(vocab_size)\n",
    "m = model.to(mps_device)\n",
    "\n",
    "# context = torch.zeros((1,1), dtype=torch.long, device=mps_device)\n",
    "# generated_chars = decode(m.generate(context, max_new_tokens=500)[0].tolist())\n",
    "# print(generated_chars)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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",
    "            logits, loss = model(X,Y)\n",
    "            losses[k] = loss.item()\n",
    "        out[split]= losses.mean()\n",
    "    model.train()\n",
    "    return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/abdullah/miniforge3/envs/gog/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step:0,train loss4.426, val loss: 4.425\n",
      "step:250,train loss1.807, val loss: 1.850\n",
      "step:500,train loss1.487, val loss: 1.583\n",
      "step:750,train loss1.340, val loss: 1.482\n",
      "1.3297969102859497\n"
     ]
    }
   ],
   "source": [
    "optimizer  = torch.optim.AdamW(model.parameters(), lr = learning_rate)\n",
    "\n",
    "for iter in range(max_iterations):\n",
    "    if iter % eval_iters==0:\n",
    "        losses = estimate_loss()\n",
    "        print(f\"step:{iter},train loss{losses['train']:.3f}, val loss: {losses['val']:.3f}\")\n",
    "\n",
    "    xb, yb = get_batch('train')## batch data sample \n",
    "\n",
    "    ## loss eval\n",
    "    logits, loss = model.forward(xb,yb)\n",
    "    optimizer.zero_grad(set_to_none= True)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "print(loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context = torch.zeros((1,1), dtype=torch.long, device=mps_device)\n",
    "generated_chars = decode(m.generate(context, max_new_tokens=500)[0].tolist())\n",
    "print(generated_chars)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf",
   "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.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}