train/evmxo-training-lstm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0557f840",
   "metadata": {},
   "source": [
    "Simple RNN trained on the dataset tokenized, no token between calls, only at the beginning of a new trace.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 341,
   "id": "20f0b695",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "! [ -e /content ] && pip install -Uqq fastbook\n",
    "import fastbook\n",
    "\n",
    "fastbook.setup_book()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 342,
   "id": "9805fd92",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import json\n",
    "\n",
    "from fastbook import *\n",
    "from fastai.text.all import *\n",
    "\n",
    "# dtype = torch.float\n",
    "# device = torch.device(\"mps\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21cf561a",
   "metadata": {},
   "source": [
    "# Prepare data for model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 343,
   "id": "e6b52567",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ensure proper ordering of the batch\n",
    "def group_chunks(ds, bs):\n",
    "    m = len(ds) // bs\n",
    "    new_ds = L()\n",
    "    for i in range(m): new_ds += L(ds[i + m*j] for j in range(bs))\n",
    "    return new_ds\n",
    "\n",
    "def get_dls(ds_div, bs, sl):\n",
    "    with open('../extract/data/model_input_med.txt', 'r') as f:\n",
    "        text = f.read()\n",
    "\n",
    "    text = '[UNK] ' + text\n",
    "    text = text[:int(len(text)/ds_div)]\n",
    "\n",
    "    tokens = text.strip().split(' ')\n",
    "    vocab = L(*tokens).unique()\n",
    "\n",
    "    print(f'text length: {len(text):,d}')\n",
    "    print(f'tokens length: {len(tokens):,d}')\n",
    "    print(f'vocab length: {len(vocab):,d}')\n",
    "\n",
    "    word2idx = {w:i for i,w in enumerate(vocab)}\n",
    "    nums = L(word2idx[i] for i in tokens)\n",
    "    \n",
    "    seqs = L((tensor(nums[i:i+sl]), tensor(nums[i+sl:i+sl+sl]))\n",
    "             for i in range(0,len(nums)-sl-1,sl))\n",
    "    m = len(seqs)//bs\n",
    "    m,bs,len(seqs)\n",
    "    cut = int(len(seqs) * 0.8)\n",
    "    dls = DataLoaders.from_dsets(group_chunks(seqs[:cut], bs),\n",
    "                                 group_chunks(seqs[cut:], bs),\n",
    "                                 bs=bs, drop_last=True, shuffle=False)\n",
    "    return dls, vocab\n",
    "\n",
    "class LMModel(Module):\n",
    "    def __init__(self, vocab_sz, n_hidden, n_layers, p):\n",
    "        self.i_h = nn.Embedding(vocab_sz, n_hidden)\n",
    "        self.rnn = nn.LSTM(n_hidden, n_hidden, n_layers, batch_first=True)\n",
    "        self.drop = nn.Dropout(p)\n",
    "        self.h_o = nn.Linear(n_hidden, vocab_sz)\n",
    "        self.h_o.weight = self.i_h.weight\n",
    "        self.h = [torch.zeros(n_layers, bs, n_hidden) for _ in range(2)]\n",
    "\n",
    "    def forward(self, x):\n",
    "        raw,h = self.rnn(self.i_h(x), self.h)\n",
    "        out = self.drop(raw)\n",
    "        self.h = [h_.detach() for h_ in h]\n",
    "        return self.h_o(out),raw,out\n",
    "    \n",
    "    def reset(self): \n",
    "        for h in self.h: h.zero_()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35a75c2c",
   "metadata": {},
   "source": [
    "# Data model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 344,
   "id": "67c11702",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "text length: 47,324,779\n",
      "tokens length: 1,541,629\n",
      "vocab length: 17,351\n",
      "ds_div: 8, bs: 8, emb_sz: 128, n_layers: 2, sl: 16, lr: 0.01\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "<style>\n",
       "    /* Turns off some styling */\n",
       "    progress {\n",
       "        /* gets rid of default border in Firefox and Opera. */\n",
       "        border: none;\n",
       "        /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "        background-size: auto;\n",
       "    }\n",
       "    progress:not([value]), progress:not([value])::-webkit-progress-bar {\n",
       "        background: repeating-linear-gradient(45deg, #7e7e7e, #7e7e7e 10px, #5c5c5c 10px, #5c5c5c 20px);\n",
       "    }\n",
       "    .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "        background: #F44336;\n",
       "    }\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>my_accuracy</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>2.182962</td>\n",
       "      <td>3.600276</td>\n",
       "      <td>0.599596</td>\n",
       "      <td>0.599596</td>\n",
       "      <td>03:53</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>2.300854</td>\n",
       "      <td>1.567296</td>\n",
       "      <td>0.809434</td>\n",
       "      <td>0.809434</td>\n",
       "      <td>03:56</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "emb_sz = 128\n",
    "n_layers = 2\n",
    "wd = 0.1\n",
    "n_epoch = 2\n",
    "lr = 1e-2\n",
    "bs = 8\n",
    "ds_div = 8\n",
    "\n",
    "sl = 16\n",
    "\n",
    "inputs = []\n",
    "preds = []\n",
    "targs = []\n",
    "cnt = 0\n",
    "# def my_accuracy(inp, targ, axis=-1):\n",
    "#     \"Compute accuracy with `targ` when `pred` is bs * n_classes\"\n",
    "# #     print('Inp:', inp)\n",
    "#     if len(inputs) < 10:\n",
    "#         inputs.append(inp)\n",
    "#         targs.append(targ)\n",
    "#     pred,targ = flatten_check(inp.argmax(dim=axis), targ)\n",
    "#     if len(preds) < 10:\n",
    "#         preds.append(pred)\n",
    "# #     print('Pred: ', pred)\n",
    "# #     print('Targ: ', targ)\n",
    "#     return (pred == targ).float().mean()\n",
    "\n",
    "dls, vocab = get_dls(ds_div, bs, sl)\n",
    "\n",
    "model = torch.compile(LMModel(len(vocab), emb_sz, n_layers, 0.4))\n",
    "print(f'ds_div: {ds_div}, bs: {bs}, emb_sz: {emb_sz}, n_layers: {n_layers}, sl: {sl}, lr: {lr}')\n",
    "learn = TextLearner(dls, model,\n",
    "                    loss_func=CrossEntropyLossFlat(), metrics=[accuracy])\n",
    "learn.fit_one_cycle(n_epoch, lr, wd=wd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "3ce5906c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "saved model: evmxo_vocab17351_bs8_sl16_emb128\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "evmxo_path = Path('.')\n",
    "model_path = Path('serve')\n",
    "vocab_path = Path('index_to_names')\n",
    "state_dict_path = Path('model-store')\n",
    "\n",
    "vocab_len = len(vocab)\n",
    "model_name = f'evmxo_vocab{vocab_len}_bs{bs}_sl{sl}_emb{emb_sz}'\n",
    "\n",
    "state_dict_filename = f'state_dict_{model_name}.pt'\n",
    "torch.save(learn.model.state_dict(), evmxo_path/state_dict_path/state_dict_filename)\n",
    "\n",
    "vocab_dict = {k: v for k, v in enumerate(vocab)}\n",
    "vocab_filename = f'index_to_name_{model_name}.json'\n",
    "with open(evmxo_path/vocab_path/vocab_filename, 'w') as f:\n",
    "    json.dump(vocab_dict, f, indent=2)\n",
    "\n",
    "model_filename = f'model_{model_name}.py'\n",
    "with open(evmxo_path/model_path/model_filename, 'w') as f:\n",
    "    model_code = f'''\n",
    "import base_model\n",
    "\n",
    "batch_size = {bs}\n",
    "embeddings_size = {emb_sz}\n",
    "vocab_length = {vocab_len}\n",
    "\n",
    "class LMModel_{model_name}(base_model.LMModel):\n",
    "    def __init__(self):\n",
    "        super(LMModel_{model_name}, self).__init__(vocab_length, embeddings_size, 2, batch_size)\n",
    "'''\n",
    "    f.write(model_code)\n",
    "\n",
    "print(f'saved model: {model_name}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "77a0a8d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING - Overwriting ../evmxo/model-store/evmxo_vocab17351_bs8_sl16_emb128.mar ...\r\n"
     ]
    }
   ],
   "source": [
    "! cp {evmxo_path}/{vocab_path}/{vocab_filename} {evmxo_path}/{vocab_path}/index_to_name.json\n",
    "! torch-model-archiver --model-name {model_name} --model-file {evmxo_path}/{model_path}/{model_filename} \\\n",
    "--serialized-file {evmxo_path}/{state_dict_path}/{state_dict_filename} \\\n",
    "--extra-files {evmxo_path}/{model_path}/evmxHandler.py,{evmxo_path}/{model_path}/base_model.py,{evmxo_path}/{vocab_path}/index_to_name.json \\\n",
    "--handler {evmxo_path}/serve/handler.py -v 1.0 \\\n",
    "--export-path {evmxo_path}/model-store -f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6d888d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training to the max on different batch sizes\n",
    "\n",
    "# still slightly increasing after 8 epoch\n",
    "# ds_div: 4, bs: 64, emb_sz: 128, n_layers: 2, sl: 64\n",
    "# epoch\ttrain_loss\tvalid_loss\taccuracy\ttime\n",
    "# 0\t3.281244\t3.325129\t0.578501\t03:14\n",
    "# 1\t1.484431\t1.413652\t0.765938\t03:08\n",
    "# 2\t1.174471\t1.098444\t0.790480\t03:09\n",
    "# 3\t0.992564\t0.849274\t0.821012\t03:10\n",
    "# 4\t0.874084\t0.709975\t0.842021\t03:10\n",
    "# 5\t0.769414\t0.605410\t0.858727\t03:08\n",
    "# 6\t0.700507\t0.542126\t0.868866\t03:08\n",
    "# 7\t0.673132\t0.523573\t0.870518\t03:12\n",
    "\n",
    "# try 24 epochs\n",
    "# ds_div: 4, bs: 8, emb_sz: 128, n_layers: 2, sl: 32, lr: 0.001\n",
    "# epoch\ttrain_loss\tvalid_loss\taccuracy\ttime\n",
    "# 0\t5.406795\t4.948601\t0.410613\t06:39\n",
    "# 1\t3.266705\t3.139425\t0.624549\t06:40\n",
    "# 2\t2.074166\t2.418597\t0.669145\t06:45\n",
    "# 3\t1.576855\t2.224808\t0.696416\t06:41\n",
    "# 4\t1.498893\t2.107253\t0.706125\t06:37\n",
    "# 5\t1.425497\t2.024300\t0.718755\t06:38\n",
    "# 6\t1.391082\t2.117145\t0.707255\t06:37\n",
    "# 7\t1.368256\t1.924885\t0.728522\t06:38\n",
    "# 8\t1.361894\t1.828822\t0.744813\t06:37\n",
    "# 9\t1.350665\t1.753814\t0.752652\t06:31\n",
    "# 10\t1.339556\t1.677444\t0.761733\t06:34\n",
    "# 11\t1.332376\t1.580631\t0.771484\t06:34\n",
    "# 12\t1.314510\t1.491990\t0.782072\t06:34\n",
    "# 13\t1.298998\t1.410088\t0.789027\t06:42\n",
    "# 14\t1.285932\t1.331676\t0.795300\t06:44\n",
    "# 15\t1.292275\t1.289153\t0.800699\t06:40\n",
    "\n",
    "\n",
    "# SAVED\n",
    "# vocab length: 18,260\n",
    "# ds_div: 2, bs: 64, emb_sz: 128, n_layers: 2, sl: 32, lr: 0.001\n",
    "# epoch\ttrain_loss\tvalid_loss\taccuracy\ttime\n",
    "# 0\t4.025664\t3.036482\t0.657743\t06:58\n",
    "# 1\t1.821263\t1.476473\t0.780719\t06:55\n",
    "# 2\t1.225658\t0.898250\t0.845399\t06:58\n",
    "# 3\t0.946340\t0.644613\t0.879003\t06:46\n",
    "# 4\t0.810226\t0.555904\t0.888096\t06:45\n",
    "# 5\t0.740880\t0.522233\t0.891788\t06:41\n",
    "# 6\t0.711271\t0.506037\t0.893532\t06:44\n",
    "# 7\t0.703302\t0.500850\t0.894477\t06:44\n",
    "\n",
    "# SAVED\n",
    "# vocab length: 18,260\n",
    "# ds_div: 2, bs: 8, emb_sz: 128, n_layers: 2, sl: 32, lr: 0.001\n",
    "# epoch\ttrain_loss\tvalid_loss\taccuracy\ttime\n",
    "# 0\t4.392361\t3.268613\t0.601746\t13:26\n",
    "# 1\t2.662708\t2.112759\t0.721673\t13:24\n",
    "# 2\t1.782815\t1.478820\t0.773827\t13:22\n",
    "# 3\t1.457928\t1.234130\t0.800126\t13:19\n",
    "# 4\t1.376427\t1.260093\t0.801919\t13:24\n",
    "# 5\t1.315052\t1.258130\t0.801879\t13:46\n",
    "# 6\t1.280592\t1.217529\t0.805961\t13:41\n",
    "# 7\t1.245183\t1.205008\t0.808077\t13:28\n",
    "# 8\t1.228678\t1.194641\t0.809214\t17:04\n",
    "# 9\t1.214260\t1.165168\t0.812769\t13:27\n",
    "# 10\t1.200192\t1.159776\t0.812537\t13:28\n",
    "# 11\t1.194933\t1.153038\t0.813744\t13:24\n",
    "# 12\t1.174281\t1.140359\t0.815773\t13:22\n",
    "# 13\t1.170180\t1.129120\t0.818285\t13:19\n",
    "# 14\t1.159167\t1.108682\t0.821779\t13:20\n",
    "# 15\t1.145710\t1.092171\t0.823530\t13:20\n",
    "# 16\t1.139795\t1.081515\t0.823574\t13:20\n",
    "# 17\t1.126797\t1.061210\t0.827687\t13:20\n",
    "# 18\t1.118184\t1.034794\t0.831409\t13:16\n",
    "# 19\t1.105856\t1.008426\t0.835146\t13:18\n",
    "# 20\t1.088293\t0.979371\t0.839519\t13:19\n",
    "# 21\t1.085581\t0.950177\t0.843593\t13:20\n",
    "# 22\t1.062206\t0.916075\t0.848545\t13:15\n",
    "# 23\t1.049483\t0.880122\t0.853547\t13:14\n",
    "# 24\t1.030618\t0.838580\t0.859141\t13:15\n",
    "# 25\t1.027435\t0.811198\t0.864295\t13:17\n",
    "# 26\t1.010219\t0.779236\t0.869996\t13:18\n",
    "# 27\t0.999372\t0.758293\t0.872126\t13:18\n",
    "# 28\t0.990476\t0.740457\t0.873269\t13:19\n",
    "# 29\t0.985521\t0.730486\t0.873999\t13:20\n",
    "# 30\t0.980521\t0.724510\t0.874585\t3:46:08\n",
    "# 31\t0.980267\t0.723266\t0.874661\t4:36:07\n",
    "\n",
    "# SAVE xlarge\n",
    "# vocab length: 19,032\n",
    "# ds_div: 1, bs: 64, emb_sz: 128, n_layers: 2, sl: 32, lr: 0.001\n",
    "# epoch\ttrain_loss\tvalid_loss\taccuracy\ttime\n",
    "# 0\t2.977237\t2.647233\t0.679372\t13:35\n",
    "# 1\t1.058554\t0.942487\t0.845939\t13:42\n",
    "# 2\t0.670738\t0.689060\t0.872087\t13:39\n",
    "# 3\t0.573277\t0.571642\t0.883391\t13:41\n",
    "# 4\t0.524220\t0.503763\t0.893860\t13:40\n",
    "# 5\t0.494350\t0.463308\t0.901011\t13:44\n",
    "# 6\t0.477324\t0.442572\t0.904965\t13:48\n",
    "# 7\t0.471718\t0.437447\t0.906063\t13:37"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a451726",
   "metadata": {},
   "source": [
    "# Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "e137bc11",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1387,\n",
       " 2698,\n",
       " 2699,\n",
       " 2700,\n",
       " 2701,\n",
       " 2702,\n",
       " 665,\n",
       " 669,\n",
       " 2703,\n",
       " 451,\n",
       " 452,\n",
       " 2704,\n",
       " 1975,\n",
       " 1976,\n",
       " 2705,\n",
       " 2232]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "word2idx = {v: k for k, v in enumerate(vocab)}\n",
    "\n",
    "def words2idx(l):\n",
    "    return [word2idx[i] for i in l]\n",
    "\n",
    "def idx2words(l):\n",
    "    return [vocab[i] for i in l]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 345,
   "id": "19a3c6f5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.1 ms ± 8.84 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "inp = [\"KS:bcb\", \"KS:bcc\", \"KP:0x102874f8c031d6ffdd1f56044c9694f8da5a94e8544779489cdbe3c6f1aeb\", \"KS:c3b\", \"KS:c3c\", \"KP:0xab17cf9d32a81474da175e525185d4126797203f1935d379f0e5cbbcfe621\", \"KS:e76\", \"KS:e77\", \"KP:0x717fe4ed56f5623deae4023431b785ee24a24f1cc87afbd6841ada19332f5\", \"KS:df5\", \"KS:df6\", \"KP:0x405db1111f2d2da95e58ea488695497e1059df8ca77b713e262cde10b9e09\", \"KS:1dd\", \"KS:1de\", \"KP:0x7bb4d269e2b923067312d8b5314084b3f8c7504c5d71107c3b4d1874dfb08\", \"KS:297\"]\n",
    "inp_idx = torch.Tensor(words2idx(inp)).type(torch.int)\n",
    "\n",
    "t = torch.zeros(7, 16, dtype=int)\n",
    "t = torch.cat((inp_idx.unsqueeze(dim=0), t))\n",
    "\n",
    "learn.model.eval()\n",
    "with torch.inference_mode():\n",
    "#     print(t)\n",
    "    %timeit learn.model(t)\n",
    "#     print(results)\n",
    "\n",
    "# print('inp: ', json.dumps(inp))\n",
    "# pred = idx2words(results[0].argmax(dim=-1))\n",
    "# print('results: ', results[0])\n",
    "# print('pred: ', json.dumps(pred))\n",
    "# # target = idx2words(train_pred[2][idx])\n",
    "# # print('target: ', json.dumps(target))\n",
    "# # print('accuracy: ', len(set(pred).intersection(target)) / len(pred))\n",
    "\n",
    "\n",
    "# # model.eval()\n",
    "# # with torch.inference_mode():\n",
    "# #     results_loaded, _, _ = model(t)\n",
    "# # pred_loaded = idx2words(results.argmax(dim=-1)[0])\n",
    "# # print('pred_loaded: ', json.dumps(pred_loaded))\n",
    "# # print('accuracy_loaded: ', len(set(pred_loaded).intersection(target)) / len(pred_loaded))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}