033_simplified_k3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.init as init\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from tqdm import tqdm\n",
    "from time import time\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.colors as mcolors\n",
    "\n",
    "colors = ['white', '#eae2b7', '#fcbf49', '#f77f00', '#d62828', '#003049']\n",
    "boundaries = [-1, 0.001, 0.10, 0.25, 0.5, 0.75, 1]\n",
    "cmap = mcolors.ListedColormap(colors)\n",
    "norm = mcolors.BoundaryNorm(boundaries, cmap.N, clip=True)\n",
    "rose = '#ff99ac'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_fname = 'data/llm/events_train_k3.txt'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Organizing data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(data_fname, sep=' ', header=None, dtype=str)\n",
    "df = df.fillna('<NaN>')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = []\n",
    "for i in range(33, 40):\n",
    "    _df = df.copy()\n",
    "    _df['target'] = _df[i]\n",
    "    for j in range(i, 40):\n",
    "        _df[j] = '<NaN>'\n",
    "\n",
    "    train_df.append(_df)\n",
    "train_df = pd.concat(train_df, ignore_index=True)\n",
    "train_df = train_df[train_df['target'] != '<NaN>']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tokenizing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer_map = json.load(open('models/llm/tokenizer_map.json', 'r'))\n",
    "detokenizer_map = {v: k for k, v in tokenizer_map.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(0, 40):\n",
    "    train_df[i] = train_df[i].map(tokenizer_map)\n",
    "train_df['target'] = train_df['target'].map(tokenizer_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = train_df.drop('target', axis=1).astype(float).values\n",
    "Y_train = pd.get_dummies(train_df['target']).astype(float).values\n",
    "\n",
    "X_train_tensor = torch.tensor(X_train, dtype=torch.float32)\n",
    "Y_train_tensor = torch.tensor(Y_train, dtype=torch.float32)\n",
    "\n",
    "train_dataset = TensorDataset(X_train_tensor, Y_train_tensor)\n",
    "\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=1024, shuffle=True)\n",
    "\n",
    "input_size = X_train.shape[1]\n",
    "output_size = len(tokenizer_map) - 1 # because of <NaN> token not existing in the output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "MODEL_NAME = 'llm_v1_tokens_v2_k3'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def flatten(l):\n",
    "    return [item for sublist in l for item in sublist]\n",
    "\n",
    "class MultiLayerBinaryClassifier(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, output_size, activation='relu'):\n",
    "        super(MultiLayerBinaryClassifier, self).__init__()\n",
    "\n",
    "        activation_dict = {\n",
    "            'relu': nn.ReLU,\n",
    "            'sigmoid': nn.Sigmoid,\n",
    "            'tanh': nn.Tanh,\n",
    "            'leaky_relu': nn.LeakyReLU,\n",
    "        }\n",
    "        layers = [\n",
    "            nn.Linear(input_size, hidden_size[0]),\n",
    "            activation_dict[activation]()\n",
    "        ] + flatten([\n",
    "            [nn.Linear(hidden_size[i], hidden_size[i+1]),\n",
    "            activation_dict[activation]()] for i in range(len(hidden_size) - 1)\n",
    "        ]) + [\n",
    "            nn.Linear(hidden_size[-1], output_size),\n",
    "            nn.Sigmoid()\n",
    "        ]\n",
    "\n",
    "        self.model = nn.Sequential(*layers)\n",
    "        \n",
    "        # Initialize the linear layers\n",
    "        self.init_weights()\n",
    "\n",
    "    def init_weights(self):\n",
    "        for m in self.model.modules():\n",
    "            if isinstance(m, nn.Linear):\n",
    "                init.xavier_uniform_(m.weight)\n",
    "                init.zeros_(m.bias)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "    \n",
    "def cyclic_cosine_annealing_lr(lr, T_max, eta_min=0, last_epoch=-1):\n",
    "    if last_epoch == 0:\n",
    "        return lr\n",
    "\n",
    "    if last_epoch % (2 * T_max) < T_max:\n",
    "        return (\n",
    "            eta_min\n",
    "            + (lr - eta_min)\n",
    "            * (1 + torch.cos(torch.tensor(3.1415 * last_epoch / T_max)))\n",
    "            / 2\n",
    "        )\n",
    "    else:\n",
    "        return (\n",
    "            eta_min\n",
    "            + (lr - eta_min)\n",
    "            * (1 + torch.cos(torch.tensor(3.1415 * (last_epoch - T_max) / T_max)))\n",
    "            / 2\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, dataloader, criterion, optimizer, device):\n",
    "    model.train()\n",
    "    running_loss = 0.0\n",
    "\n",
    "    for inputs, labels in dataloader:\n",
    "        inputs, labels = inputs.to(device), labels.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "\n",
    "        running_loss += loss.item() * inputs.size(0)\n",
    "\n",
    "    epoch_loss = running_loss / len(dataloader.dataset)\n",
    "    return epoch_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model, dataloader, criterion, device):\n",
    "    model.eval()\n",
    "    running_loss = 0.0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for inputs, labels in dataloader:\n",
    "            inputs, labels = inputs.to(device), labels.to(device)\n",
    "\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, labels)\n",
    "\n",
    "            running_loss += loss.item() * inputs.size(0)\n",
    "\n",
    "    epoch_loss = running_loss / len(dataloader.dataset)\n",
    "    return epoch_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MultiLayerBinaryClassifier(input_size, [512, 512, 512], output_size).to(DEVICE)\n",
    "learning_rate_init = 0.001\n",
    "num_epochs = 50\n",
    "best_val_loss = 1000\n",
    "\n",
    "criterion = nn.BCELoss()\n",
    "for epoch in range(num_epochs):\n",
    "    t0 = time()\n",
    "    lr_update = cyclic_cosine_annealing_lr(learning_rate_init, num_epochs, 0, epoch)\n",
    "    optimizer = optim.Adam(model.parameters(), lr=lr_update)\n",
    "    train_loss = train(model, train_dataloader, criterion, optimizer, DEVICE)\n",
    "    test_loss = evaluate(model, train_dataloader, criterion, DEVICE)\n",
    "    print(f'Epoch: {epoch+1}/{num_epochs}. Training loss: {train_loss:.4f}. Test loss: {test_loss:.4f}. Time: {time() - t0:.2f}s')\n",
    "\n",
    "    if test_loss < best_val_loss:\n",
    "        best_val_loss = test_loss\n",
    "        torch.save(model, f'models/llm/full_{MODEL_NAME}.pth')"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}