paren_2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import json\n",
    "from typing import Union\n",
    "from jaxtyping import Int\n",
    "from torch import Tensor\n",
    "import random\n",
    "import numpy as np\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import accuracy_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "VOCAB = {\"[pad]\": 0, \"[start]\": 1, \"[end]\": 2, \"(\": 3, \")\": 4}\n",
    "HIDDEN_SIZE = 56\n",
    "HEAD_SIZE = 28\n",
    "NUM_LAYERS = 3\n",
    "NUM_HEADS = 2\n",
    "MAX_LEN = 100\n",
    "BATCH_SIZE = 32\n",
    "EPOCHS = 10\n",
    "LEARNING_RATE = 0.001\n",
    "device = \"mps\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class SimpleTokenizer:\n",
    "    START_TOKEN = 0\n",
    "    PAD_TOKEN = 1\n",
    "    END_TOKEN = 2\n",
    "    base_d = {\"[start]\": START_TOKEN, \"[pad]\": PAD_TOKEN, \"[end]\": END_TOKEN}\n",
    "\n",
    "    def __init__(self, alphabet: str):\n",
    "        self.alphabet = alphabet\n",
    "        # the 3 is because there are 3 special tokens (defined just above)\n",
    "        self.t_to_i = {**{c: i + 3 for i, c in enumerate(alphabet)}, **self.base_d}\n",
    "        self.i_to_t = {i: c for c, i in self.t_to_i.items()}\n",
    "\n",
    "    def tokenize(self, strs: list[str], max_len=None) -> Int[Tensor, \"batch seq\"]:\n",
    "        def c_to_int(c: str) -> int:\n",
    "            if c in self.t_to_i:\n",
    "                return self.t_to_i[c]\n",
    "            else:\n",
    "                raise ValueError(c)\n",
    "\n",
    "        if isinstance(strs, str):\n",
    "            strs = [strs]\n",
    "\n",
    "        if max_len is None:\n",
    "            max_len = max((max(len(s) for s in strs), 1))\n",
    "\n",
    "        ints = [\n",
    "            [self.START_TOKEN]\n",
    "            + [c_to_int(c) for c in s]\n",
    "            + [self.END_TOKEN]\n",
    "            + [self.PAD_TOKEN] * (max_len - len(s))\n",
    "            for s in strs\n",
    "        ]\n",
    "        return torch.tensor(ints)\n",
    "\n",
    "    def decode(self, tokens) -> list[str]:\n",
    "        assert tokens.ndim >= 2, \"Need to have a batch dimension\"\n",
    "\n",
    "        def int_to_c(c: int) -> str:\n",
    "            if c < len(self.i_to_t):\n",
    "                return self.i_to_t[c]\n",
    "            else:\n",
    "                raise ValueError(c)\n",
    "\n",
    "        return [\n",
    "            \"\".join(\n",
    "                int_to_c(i.item()) for i in seq[1:] if i != self.PAD_TOKEN and i != self.END_TOKEN\n",
    "            )\n",
    "            for seq in tokens\n",
    "        ]\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        return f\"SimpleTokenizer({self.alphabet!r})\"\n",
    "\n",
    "\n",
    "class BracketsDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, data_tuples, tokenizer):\n",
    "        self.tokenizer = SimpleTokenizer(\"()\")\n",
    "        self.strs = [x[0] for x in data_tuples]\n",
    "        self.isbal = torch.tensor([x[1] for x in data_tuples])\n",
    "        self.toks = self.tokenizer.tokenize(self.strs)\n",
    "        self.open_proportion = torch.tensor([s.count(\"(\") / len(s) for s in self.strs])\n",
    "        self.starts_open = torch.tensor([s[0] == \"(\" for s in self.strs]).bool()\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.strs)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.strs[idx], self.isbal[idx], self.toks[idx]\n",
    "\n",
    "    def to(self, device):\n",
    "        self.isbal = self.isbal.to(device)\n",
    "        self.toks = self.toks.to(device)\n",
    "        self.open_proportion = self.open_proportion.to(device)\n",
    "        self.starts_open = self.starts_open.to(device)\n",
    "        return self\n",
    "\n",
    "def load_data():\n",
    "    with open(\"/Users/utkarsh/Documents/neural-toc/naacl_work/ARENA_3.0/chapter1_transformer_interp/exercises/part51_balanced_bracket_classifier/brackets_data.json\") as f:\n",
    "        data_tuples = json.load(f)\n",
    "    data_tuples = data_tuples\n",
    "    random.shuffle(data_tuples)\n",
    "\n",
    "    train_size = int(0.8 * len(data_tuples))\n",
    "    val_size = int(0.1 * len(data_tuples))\n",
    "    test_size = len(data_tuples) - train_size - val_size\n",
    "\n",
    "    train_data = data_tuples[:train_size]\n",
    "    val_data = data_tuples[train_size:train_size+val_size]\n",
    "    test_data = data_tuples[train_size+val_size:]\n",
    "\n",
    "    tokenizer = SimpleTokenizer(\"()\")\n",
    "    train_dataset = BracketsDataset(train_data, tokenizer).to(device)\n",
    "    val_dataset = BracketsDataset(val_data, tokenizer).to(device)\n",
    "    test_dataset = BracketsDataset(test_data, tokenizer).to(device)\n",
    "\n",
    "    return train_dataset, val_dataset, test_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, hidden_size, head_size, num_heads):\n",
    "        super().__init__()\n",
    "        self.num_heads = num_heads\n",
    "        self.head_size = head_size\n",
    "        self.W_q = nn.Linear(hidden_size, num_heads * head_size)\n",
    "        self.W_k = nn.Linear(hidden_size, num_heads * head_size)\n",
    "        self.W_v = nn.Linear(hidden_size, num_heads * head_size)\n",
    "        self.W_o = nn.Linear(num_heads * head_size, hidden_size)\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        batch_size, seq_len, hidden_size = x.size()\n",
    "        Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.head_size).transpose(1, 2)\n",
    "        K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.head_size).transpose(1, 2)\n",
    "        V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.head_size).transpose(1, 2)\n",
    "\n",
    "        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_size ** 0.5)\n",
    "        if mask is not None:\n",
    "            scores = scores.masked_fill(mask == 0, float('-inf'))\n",
    "        attention_weights = torch.softmax(scores, dim=-1)\n",
    "        context = torch.matmul(attention_weights, V).transpose(1, 2).contiguous()\n",
    "        context = context.view(batch_size, seq_len, -1)\n",
    "        return self.W_o(context)\n",
    "\n",
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, hidden_size, head_size, num_heads):\n",
    "        super().__init__()\n",
    "        self.attention = MultiHeadAttention(hidden_size, head_size, num_heads)\n",
    "        self.layernorm1 = nn.LayerNorm(hidden_size)\n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Linear(hidden_size, 4 * hidden_size),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(4 * hidden_size, hidden_size)\n",
    "        )\n",
    "        self.layernorm2 = nn.LayerNorm(hidden_size)\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        attn_output = self.attention(x, mask)\n",
    "        x = self.layernorm1(x + attn_output)\n",
    "        mlp_output = self.mlp(x)\n",
    "        return self.layernorm2(x + mlp_output)\n",
    "\n",
    "# class BalancedParenthesesModel(nn.Module):\n",
    "#     def __init__(self, vocab_size, hidden_size, max_len, num_layers, num_heads):\n",
    "#         super().__init__()\n",
    "#         self.embedding = nn.Embedding(vocab_size, hidden_size)\n",
    "#         self.positional_encodings = nn.Parameter(torch.zeros(1, max_len, hidden_size))\n",
    "#         self.layers = nn.ModuleList([\n",
    "#             TransformerBlock(hidden_size, HEAD_SIZE, num_heads)\n",
    "#             for _ in range(num_layers)\n",
    "#         ])\n",
    "#         self.layernorm_final = nn.LayerNorm(hidden_size)\n",
    "#         self.unembedding = nn.Linear(hidden_size, 2)\n",
    "\n",
    "#     def forward(self, x, mask=None):\n",
    "#         seq_len = x.size(1)\n",
    "#         x = self.embedding(x) + self.positional_encodings[:, :seq_len, :]\n",
    "#         for layer in self.layers:\n",
    "#             x = layer(x, mask)\n",
    "#         x = self.layernorm_final(x)\n",
    "#         logits = self.unembedding(x[:, 0, :])\n",
    "#         return logits\n",
    "\n",
    "class BalancedParenthesesModel(nn.Module):\n",
    "    def __init__(self, vocab_size, hidden_size, max_len, num_layers, num_heads):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(vocab_size, hidden_size)\n",
    "        self.positional_encodings = nn.Parameter(torch.zeros(1, max_len, hidden_size))\n",
    "        self.layers = nn.ModuleList([\n",
    "            TransformerBlock(hidden_size, HEAD_SIZE, num_heads)\n",
    "            for _ in range(num_layers)\n",
    "        ])\n",
    "        self.layernorm_final = nn.LayerNorm(hidden_size)\n",
    "        self.unembedding = nn.Linear(hidden_size, 2)\n",
    "        \n",
    "        # Store intermediate states during forward pass\n",
    "        self.token_states = None\n",
    "\n",
    "    def forward(self, x, mask=None, return_states=False):\n",
    "        seq_len = x.size(1)\n",
    "        \n",
    "        # Embedding with positional encodings\n",
    "        x = self.embedding(x) + self.positional_encodings[:, :seq_len, :]\n",
    "        \n",
    "        # Pass through transformer layers\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, mask)\n",
    "        \n",
    "        # Final layer normalization\n",
    "        x = self.layernorm_final(x)\n",
    "        \n",
    "        # Compute logits\n",
    "        logits = self.unembedding(x[:, 0, :])\n",
    "        \n",
    "        # Store token states if requested\n",
    "        if return_states:\n",
    "            return logits, x.detach()\n",
    "        \n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "def train_model(model, train_dataset, val_dataset, num_epochs, batch_size, lr, device):\n",
    "    \"\"\"\n",
    "    Train the Transformer model on the brackets dataset.\n",
    "    \n",
    "    Args:\n",
    "        model (TransformerModel): The Transformer model to train.\n",
    "        train_dataset (BracketsDataset): The training dataset.\n",
    "        val_dataset (BracketsDataset): The validation dataset.\n",
    "        num_epochs (int): The number of training epochs.\n",
    "        batch_size (int): The batch size.\n",
    "        lr (float): The learning rate.\n",
    "        device (torch.device): The device to use for training (CPU or GPU).\n",
    "    \"\"\"\n",
    "    # Set up data loaders\n",
    "    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    # Set up the optimizer and loss function\n",
    "    optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "    # Move the model to the specified device\n",
    "    model.to(device)\n",
    "\n",
    "    # Train the model\n",
    "    best_val_loss = float('inf')\n",
    "    for epoch in range(num_epochs):\n",
    "        # Training loop\n",
    "        model.train()\n",
    "        total_train_loss = 0\n",
    "        for batch, (_, labels, tokens) in enumerate(train_loader):\n",
    "            optimizer.zero_grad()\n",
    "            output = model(tokens)\n",
    "            loss = criterion(output, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            total_train_loss += loss.item()\n",
    "\n",
    "        # Validation loop\n",
    "        model.eval()\n",
    "        total_val_loss = 0\n",
    "        with torch.no_grad():\n",
    "            for _, (_, labels, tokens) in enumerate(val_loader):\n",
    "                mask = (tokens != VOCAB[\"[pad]\"]).unsqueeze(1).unsqueeze(2).to(device)\n",
    "                output = model(tokens, mask)\n",
    "                loss = criterion(output, labels)\n",
    "                total_val_loss += loss.item()\n",
    "\n",
    "        # Print training and validation results\n",
    "        avg_train_loss = total_train_loss / len(train_loader)\n",
    "        avg_val_loss = total_val_loss / len(val_loader)\n",
    "        print(f\"Epoch {epoch+1}/{num_epochs}, Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}\")\n",
    "\n",
    "    print(\"Training complete.\")\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10, Train Loss: 0.2610, Val Loss: 0.0455\n",
      "Epoch 2/10, Train Loss: 0.1204, Val Loss: 0.1085\n",
      "Epoch 3/10, Train Loss: 0.1858, Val Loss: 0.1491\n",
      "Epoch 4/10, Train Loss: 0.1092, Val Loss: 0.0393\n",
      "Epoch 5/10, Train Loss: 0.0778, Val Loss: 0.0352\n",
      "Epoch 6/10, Train Loss: 0.0477, Val Loss: 0.0534\n",
      "Epoch 7/10, Train Loss: 0.0809, Val Loss: 0.1294\n",
      "Epoch 8/10, Train Loss: 0.0410, Val Loss: 0.0220\n",
      "Epoch 9/10, Train Loss: 0.0884, Val Loss: 0.3257\n",
      "Epoch 10/10, Train Loss: 0.0677, Val Loss: 0.0252\n",
      "Training complete.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "BalancedParenthesesModel(\n",
       "  (embedding): Embedding(5, 56)\n",
       "  (layers): ModuleList(\n",
       "    (0-2): 3 x TransformerBlock(\n",
       "      (attention): MultiHeadAttention(\n",
       "        (W_q): Linear(in_features=56, out_features=56, bias=True)\n",
       "        (W_k): Linear(in_features=56, out_features=56, bias=True)\n",
       "        (W_v): Linear(in_features=56, out_features=56, bias=True)\n",
       "        (W_o): Linear(in_features=56, out_features=56, bias=True)\n",
       "      )\n",
       "      (layernorm1): LayerNorm((56,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): Sequential(\n",
       "        (0): Linear(in_features=56, out_features=224, bias=True)\n",
       "        (1): ReLU()\n",
       "        (2): Linear(in_features=224, out_features=56, bias=True)\n",
       "      )\n",
       "      (layernorm2): LayerNorm((56,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "  )\n",
       "  (layernorm_final): LayerNorm((56,), eps=1e-05, elementwise_affine=True)\n",
       "  (unembedding): Linear(in_features=56, out_features=2, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data, val_data, test_data = load_data()\n",
    "train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)\n",
    "val_loader = DataLoader(val_data, batch_size=BATCH_SIZE)\n",
    "\n",
    "model = BalancedParenthesesModel(len(VOCAB), HIDDEN_SIZE, MAX_LEN, NUM_LAYERS, NUM_HEADS)\n",
    "train_model(model, train_data, val_data, EPOCHS, BATCH_SIZE, LEARNING_RATE, device=\"mps\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.7610,  0.0120]], device='mps:0')\n",
      "1\n",
      "torch.Size([1, 18, 56])\n"
     ]
    }
   ],
   "source": [
    "tokeniser = SimpleTokenizer(\"()\")\n",
    "device = \"mps\"\n",
    "with torch.no_grad():\n",
    "    output, internal = model(tokeniser.tokenize(\"(()()()()()()())\").to(device), return_states=True)\n",
    "    print(output)\n",
    "    print(torch.argmax(output, dim=1).item())\n",
    "    print(internal.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "()())())))))((((()()))()(())))))()(()((()()((((()((()(()(()))(()())))(()())))))))()))()(()(()()\n"
     ]
    }
   ],
   "source": [
    "import random\n",
    "\n",
    "def generate_unbalanced_parenthesis_string(length):\n",
    "    parenthesis = ['(', ')']\n",
    "    result = ''.join(random.choices(parenthesis, k=length))\n",
    "    return result\n",
    "\n",
    "unbalanced_string = generate_unbalanced_parenthesis_string(95)\n",
    "print(unbalanced_string)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), \"model_mask_paren_4_dec_maxlen100_probable.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model saved to model_mask_paren_4_dec_maxlen_100_probable.pkl\n"
     ]
    }
   ],
   "source": [
    "import pickle\n",
    "save_path = 'model_mask_paren_4_dec_maxlen_100_probable.pkl'\n",
    "\n",
    "# Save the model\n",
    "with open(save_path, 'wb') as file:\n",
    "    pickle.dump(model, file)\n",
    "\n",
    "print(f\"Model saved to {save_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loader = DataLoader(test_data, batch_size=BATCH_SIZE)\n",
    "correct = 0\n",
    "total = 0\n",
    "incorrect = []\n",
    "with torch.no_grad():\n",
    "    for batch, (_, labels, tokens) in enumerate(test_loader):\n",
    "        output = model(tokens)\n",
    "        _, predicted = torch.max(output.data, 1)\n",
    "        total += labels.size(0)\n",
    "        correct += (predicted == labels).sum().item()\n",
    "        if (predicted != labels).any():\n",
    "            incorrect.append((tokeniser.decode(tokens), labels, predicted))\n",
    "\n",
    "test_accuracy = 100 * correct / total\n",
    "print(f\"Test Accuracy: {test_accuracy:.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_parenthesis_dataset(num_samples, max_length):\n",
    "    import random\n",
    "    data = []\n",
    "    for _ in range(num_samples):\n",
    "        length = random.randint(2, max_length)\n",
    "        string = \"\"\n",
    "        depth = 0\n",
    "        stack_depths = []\n",
    "        for _ in range(length):\n",
    "            if random.random() < 0.5 and depth > 0:  # Close parenthesis\n",
    "                string += ')'\n",
    "                depth -= 1\n",
    "            else:  # Open parenthesis\n",
    "                string += '('\n",
    "                depth += 1\n",
    "            stack_depths.append(depth)\n",
    "        data.append((string, stack_depths))\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "def generate_anbn_dataset(max_n):\n",
    "    dataset = []\n",
    "    for n in range(max_n + 1):\n",
    "        valid_string = 'a' * n + 'b' * n\n",
    "        dataset.append([valid_string, True])\n",
    "        \n",
    "        # Generate some invalid strings\n",
    "        if n > 0:\n",
    "            invalid_string = 'a' * n + 'b' * (n + 1)\n",
    "            dataset.append([invalid_string, False])\n",
    "            invalid_string = 'a' * (n + 1) + 'b' * n\n",
    "            dataset.append([invalid_string, False])\n",
    "    \n",
    "    return dataset\n",
    "\n",
    "anbn_dataset = generate_anbn_dataset(45)\n",
    "\n",
    "# Write dataset to a JSON file\n",
    "with open('anbn_data.json', 'w') as file:\n",
    "    json.dump(anbn_dataset, file, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}