haiku.word.transformer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Haiku Generator \n",
    "This neural net uses transformer to generate Haikus word by word."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(143137,\n",
       " ['\\nHaiku:\\n', 'Memorial', ' ', 'Day', ' ', '--', '\\n', 'a', ' ', 'shadow'],\n",
       " 4071931)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import csv\n",
    "import itertools\n",
    "import re\n",
    "\n",
    "with open('data/haikus.csv') as data:\n",
    "    csv_data = [row for row in csv.DictReader(data)]\n",
    "\n",
    "def row_to_lines(row):\n",
    "    return [\"\\nHaiku:\\n\"] + list(itertools.chain.from_iterable([re.split(r'(\\s+)', line) + ['\\n'] for line in [row[str(col)] for col in range(3)]]))\n",
    "\n",
    "tokens = list(itertools.chain.from_iterable([row_to_lines(row) for row in csv_data]))\n",
    "\n",
    "len(csv_data), tokens[:10], len(tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "71749"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "context_size = 10\n",
    "vocab = set(tokens)\n",
    "vocab_size = len(vocab)\n",
    "word_to_ix = {w:i for i, w in enumerate(vocab)}\n",
    "ix_to_word = {i:w for i, w in enumerate(vocab)}\n",
    "encode = lambda line: [word_to_ix[w] for w in list(line)]\n",
    "decode = lambda ixs: ''.join([ix_to_word[ix] for ix in ixs])\n",
    "vocab_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class Head(nn.Module):\n",
    "    def __init__(self, embedding_dim, head_size, dropout=0.3, masked=True):\n",
    "        super().__init__()\n",
    "        self.masked = masked\n",
    "        self.key = nn.Linear(embedding_dim, head_size, bias=False)\n",
    "        self.query = nn.Linear(embedding_dim, head_size, bias=False)\n",
    "        self.value = nn.Linear(embedding_dim, head_size, bias=False)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.register_buffer('tril', torch.tril(torch.ones(context_size, context_size)))\n",
    "\n",
    "    def forward(self, x):\n",
    "        B,T,C = x.shape\n",
    "        k = self.key(x)\n",
    "        q = self.query(x)\n",
    "        wei = q @ k.transpose(-2, -1)\n",
    "        if self.masked:\n",
    "            wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))\n",
    "        wei = F.softmax(wei, dim=-1)\n",
    "        wei = self.dropout(wei)\n",
    "        v = self.value(x)\n",
    "\n",
    "        return wei @ v\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, embedding_dim, n_heads, dropout=0.3):\n",
    "        super().__init__()\n",
    "        head_size = embedding_dim // n_heads\n",
    "        self.heads = nn.ModuleList([Head(embedding_dim, head_size) for _ in range(n_heads)])\n",
    "        self.proj = nn.Linear(embedding_dim, embedding_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = torch.cat([h(x) for h in self.heads], dim=-1)\n",
    "        return self.dropout(self.proj(x))\n",
    "\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, embedding_dim, dropout=0.3):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(embedding_dim, 4 * embedding_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(4 * embedding_dim, embedding_dim),\n",
    "            nn.Dropout(dropout)\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "class Block(nn.Module):\n",
    "    def __init__(self, embedding_dim, n_heads, dropout=0.3):\n",
    "        super().__init__()\n",
    "        self.sa = MultiHeadAttention(embedding_dim, n_heads, dropout)\n",
    "        self.feed_fwd = FeedForward(embedding_dim, dropout)\n",
    "        self.ln1 = nn.Linear(embedding_dim, embedding_dim)\n",
    "        self.ln2 = nn.Linear(embedding_dim, embedding_dim)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = x + self.sa(self.ln1(x))\n",
    "        x = x + self.feed_fwd(self.ln1(x))\n",
    "        return x\n",
    "    \n",
    "class TransformerModel(nn.Module):\n",
    "    def __init__(self, vocab_size, embedding_dim, classes, n_heads=4, dropout_prob=0.25, device='mps', masked=True):\n",
    "        super().__init__()\n",
    "        self.device = device\n",
    "        self.tok_emb = nn.Embedding(vocab_size, embedding_dim)\n",
    "        self.pos_emb = nn.Embedding(context_size, embedding_dim)\n",
    "        self.blocks = nn.Sequential(\n",
    "            Block(embedding_dim, n_heads, dropout_prob),\n",
    "            Block(embedding_dim, n_heads, dropout_prob),\n",
    "            Block(embedding_dim, n_heads, dropout_prob),\n",
    "            Block(embedding_dim, n_heads, dropout_prob),\n",
    "            Block(embedding_dim, n_heads, dropout_prob),\n",
    "            Block(embedding_dim, n_heads, dropout_prob),\n",
    "            nn.LayerNorm(embedding_dim)\n",
    "        )\n",
    "\n",
    "        self.fc = nn.Linear(embedding_dim, classes)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        _, T = x.size()\n",
    "        tok_emb = self.tok_emb(x)\n",
    "        pos_emb = self.pos_emb(torch.arange(T, device=self.device))\n",
    "        x = tok_emb + pos_emb\n",
    "        x = self.blocks(x)\n",
    "        out = self.fc(x)\n",
    "        return out\n",
    "    \n",
    "    def generate(self, idx, max_new_tokens):\n",
    "        for _ in range(max_new_tokens):\n",
    "            idx_cond = idx[:, -context_size:]\n",
    "            logits = self(idx_cond)\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(8, 8, 10001, 10001)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "device = 'mps'\n",
    "batch_size = 128\n",
    "train_size = context_size * 1000 + 1\n",
    "# control the size of words for training, it's too large for my m2 to process lol\n",
    "train_tokens = tokens[:train_size]\n",
    "val_tokens = tokens[train_size:2*train_size]\n",
    "\n",
    "def get_dataloader(bow):\n",
    "    n = len(bow) - 1\n",
    "    ixs = [word_to_ix[word] for word in bow]\n",
    "    x = torch.tensor(ixs[:n]).reshape((-1, context_size))\n",
    "    y = torch.tensor(ixs[1:n+1]).reshape((-1, context_size))\n",
    "    x = x.to(device)\n",
    "    y = y.to(device)\n",
    "    dataset = TensorDataset(x, y)\n",
    "    return DataLoader(dataset, batch_size)\n",
    "\n",
    "train_dataloader = get_dataloader(train_tokens)\n",
    "val_dataloader = get_dataloader(val_tokens)\n",
    "\n",
    "len(train_dataloader), len(val_dataloader),  len(train_tokens), len(val_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "device ='mps'\n",
    "embedding_dim = 64\n",
    "classes = vocab_size\n",
    "\n",
    "model = TransformerModel(vocab_size, embedding_dim, classes, n_heads=4)\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 100\n",
    "lr = 1e-3\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=lr)\n",
    "criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(10/100): train loss: 3.9485, val loss: 3.8031 (22.66s)\n",
      "(20/100): train loss: 3.3173, val loss: 3.2526 (21.37s)\n",
      "(30/100): train loss: 3.0124, val loss: 2.9165 (26.65s)\n",
      "(40/100): train loss: 2.6971, val loss: 2.5723 (34.34s)\n",
      "(50/100): train loss: 2.4342, val loss: 2.2658 (26.61s)\n",
      "(60/100): train loss: 2.1254, val loss: 1.9427 (26.75s)\n",
      "(70/100): train loss: 1.8812, val loss: 1.6497 (20.57s)\n",
      "(80/100): train loss: 1.6384, val loss: 1.3749 (19.04s)\n",
      "(90/100): train loss: 1.4393, val loss: 1.1159 (20.94s)\n",
      "(100/100): train loss: 1.2225, val loss: 0.8941 (17.08s)\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "for k in range(epochs):\n",
    "    model.train()\n",
    "    train_loss = 0.0\n",
    "    start_time = time.time()\n",
    "    for xb, yb in train_dataloader:\n",
    "        preds = model(xb)\n",
    "        B, T, C = preds.shape\n",
    "        loss = criterion(preds.view(B*T, C), yb.view(B*T))\n",
    "\n",
    "        model.zero_grad(set_to_none=True)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        train_loss += loss.item()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        val_loss = 0.0\n",
    "        for xb, yb in train_dataloader:\n",
    "            preds = model(xb)\n",
    "            B, T, C = preds.shape\n",
    "            loss = criterion(preds.view(B*T, C), yb.view(B*T))\n",
    "            val_loss += loss.item()\n",
    "        \n",
    "    if (k+1) % max(1, int(0.1*epochs)) == 0:\n",
    "        end_time = time.time()\n",
    "        torch.save(model, f'model.word.haiku.pt')\n",
    "        print(f\"({k+1}/{epochs}): train loss: {train_loss/len(train_dataloader):.4f}, val loss: {val_loss/len(val_dataloader):.4f} ({end_time - start_time:.2f}s)\")\n",
    "        start_time = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " trains rejection Summer\n",
      "too\n",
      "in\n"
     ]
    }
   ],
   "source": [
    "context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
    "print(decode(model.generate(context, max_new_tokens=10)[0].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate(model, idx, max_new_tokens):\n",
    "    def inner(idx):\n",
    "        idx_cond = idx[:, -context_size:]\n",
    "        logits = model(idx_cond)\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",
    "    for _ in range(max_new_tokens):\n",
    "        idx = inner(idx)\n",
    "        yield idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rainy day him\n",
      "your exit nightfall the shade\n",
      "now off wisdom\n",
      "\n",
      "Haiku:\n",
      "summer's rising\n",
      "my moon\n",
      "my press\n",
      "\n",
      "Haiku:\n",
      "february steamed\n",
      "just pepper\n"
     ]
    }
   ],
   "source": [
    "prompt = encode(\"rainy day\")\n",
    "context = torch.tensor(prompt, device=device).view(1, -1)\n",
    "for i, c in enumerate(generate(model, context, max_new_tokens=50)):\n",
    "    for c in decode([c[0][i].tolist()]):\n",
    "        time.sleep(.05)\n",
    "        print(c, end='', flush=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.11.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}