lstm_model.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "51ca972d-31a2-4d21-8f14-ebd996c8787a",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "51ca972d-31a2-4d21-8f14-ebd996c8787a"
      },
      "outputs": [],
      "source": [
        "# Import packages\n",
        "import pandas as pd\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "from tensorflow.keras.layers import LSTM, Activation, Dropout, Dense, Input, Conv1D, MaxPooling1D, GlobalMaxPooling1D\n",
        "from tensorflow.keras.models import Model\n",
        "import string\n",
        "import re\n",
        "from tensorflow.keras.preprocessing.text import Tokenizer\n",
        "from sklearn.preprocessing import LabelBinarizer\n",
        "from tensorflow.keras.preprocessing.sequence import pad_sequences\n",
        "import tensorflow.keras\n",
        "from sklearn.model_selection import train_test_split\n",
        "import pandas as pd\n",
        "from gensim.models import KeyedVectors\n",
        "import matplotlib as py\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Preprocessing\n"
      ],
      "metadata": {
        "id": "h8IDeKWhgAOW"
      },
      "id": "h8IDeKWhgAOW"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8db91bd6-10d0-46ce-bb3d-6fe9754d4ea0",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "8db91bd6-10d0-46ce-bb3d-6fe9754d4ea0"
      },
      "outputs": [],
      "source": [
        "data = pd.read_csv(\"final.csv\")\n",
        "data.dropna(inplace = True)\n",
        "data.reset_index(inplace = True,drop = True)\n",
        "reloaded_word_vectors = KeyedVectors.load('vectors_v2.kv')\n",
        "maxLen = 14638\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "811dd3c6-5006-423d-b34d-1df81a02f17c",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "811dd3c6-5006-423d-b34d-1df81a02f17c"
      },
      "outputs": [],
      "source": [
        "## We only grab 132 of the all data\n",
        "for x in data.index:\n",
        "      if data.loc[x, \"index\"] > 132:\n",
        "              data.drop(x, inplace = True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "15aba856-cd19-46c5-ab38-5dc231bfab78",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "15aba856-cd19-46c5-ab38-5dc231bfab78"
      },
      "outputs": [],
      "source": [
        "embed_vector_len = reloaded_word_vectors[\"SHL\"].shape[0]\n",
        "vocab_len = len(reloaded_word_vectors.index_to_key)\n",
        "emb_matrix = np.zeros((vocab_len, embed_vector_len))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d589a151-3b56-4bd2-af61-8d7702537204",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "d589a151-3b56-4bd2-af61-8d7702537204"
      },
      "outputs": [],
      "source": [
        "for index, word in enumerate(reloaded_word_vectors.index_to_key):\n",
        "  embedding_vector = reloaded_word_vectors[word]\n",
        "  if embedding_vector is not None:\n",
        "    emb_matrix[index, :] = embedding_vector"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e861e180-5153-4cf0-bc14-8f9057c541d7",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "e861e180-5153-4cf0-bc14-8f9057c541d7"
      },
      "outputs": [],
      "source": [
        "embedding_layer = tf.keras.layers.Embedding(input_dim=vocab_len, output_dim=embed_vector_len, input_length=maxLen, weights = [emb_matrix], trainable=False)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f92cb8b3-c24e-4503-8a87-18140488ab86",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "f92cb8b3-c24e-4503-8a87-18140488ab86"
      },
      "outputs": [],
      "source": [
        "## Opcodes with hex values are processed here\n",
        "def adjust_hex(row):\n",
        "    array = row.contract_code.split()\n",
        "    result_string = \"\"\n",
        "    for code in array:\n",
        "        if(\"0x\" in code):\n",
        "            continue\n",
        "        elif(\"PUSH\" in code):\n",
        "            result_string  = result_string + \"PUSH\" + \" \"\n",
        "        elif(\"SWAP\" in code):\n",
        "            result_string  = result_string + \"SWAP\" + \" \"\n",
        "        elif(\"DUP\" in code):\n",
        "            result_string  = result_string + \"DUP\" + \" \"\n",
        "        elif(\"LOG\" in code):\n",
        "            result_string  = result_string + \"LOG\" + \" \"\n",
        "        else:\n",
        "            result_string  = result_string + code + \" \"\n",
        "\n",
        "    return result_string.strip()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5b2e9775-5c04-4358-bf5d-962275d1c096",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "5b2e9775-5c04-4358-bf5d-962275d1c096"
      },
      "outputs": [],
      "source": [
        "data[\"opcode_adjusted\"] = data.apply(lambda row:adjust_hex(row),axis = 1)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "04a4c02a-9a0e-4ddb-ae58-eb0439fc91d6",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "04a4c02a-9a0e-4ddb-ae58-eb0439fc91d6"
      },
      "outputs": [],
      "source": [
        "reviews = data[\"opcode_adjusted\"]\n",
        "\n",
        "reviews_list = []\n",
        "for i in range(len(reviews)):\n",
        "  reviews_list.append(reviews[i])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c31672b6-b715-47ad-8f80-8b4cad72c756",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "c31672b6-b715-47ad-8f80-8b4cad72c756"
      },
      "outputs": [],
      "source": [
        "sentiment = data['vulnerability']\n",
        "unique_labels = np.unique(sentiment)\n",
        "label_numbers = np.searchsorted(unique_labels, sentiment)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training"
      ],
      "metadata": {
        "id": "g0kFVf7fgTJf"
      },
      "id": "g0kFVf7fgTJf"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "25783b7d-f1d1-457b-bcf3-df5a1e9486f1",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "25783b7d-f1d1-457b-bcf3-df5a1e9486f1"
      },
      "outputs": [],
      "source": [
        "X_train, X_test,Y_train, Y_test = train_test_split(reviews_list, label_numbers, test_size=0.2, random_state = 45)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d4a2c1e1-c93f-4a10-b597-5c05e17755e8",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "d4a2c1e1-c93f-4a10-b597-5c05e17755e8"
      },
      "outputs": [],
      "source": [
        "tokenizer = tf.keras.preprocessing.text.Tokenizer()\n",
        "tokenizer.fit_on_texts(X_train)\n",
        "words_to_index = tokenizer.word_index\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ad2e1595-bac2-4d89-b6e9-800dafc06c53",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "ad2e1595-bac2-4d89-b6e9-800dafc06c53"
      },
      "outputs": [],
      "source": [
        "def find_length(row):\n",
        "    return len(row[\"opcode_adjusted\"].split())\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "48dafd3a-2186-46c1-9a87-8758e075da83",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "48dafd3a-2186-46c1-9a87-8758e075da83"
      },
      "outputs": [],
      "source": [
        "def vuln_model(input_shape):\n",
        "\n",
        "  X_indices = Input(input_shape)\n",
        "\n",
        "  embeddings = embedding_layer(X_indices)\n",
        "\n",
        "  X = LSTM(128, return_sequences=True, trainable=True )(embeddings)\n",
        "\n",
        "  X = Dropout(0.6)(X)\n",
        "\n",
        "  X = LSTM(128)(X)\n",
        "\n",
        "  X = Dense(1, activation='sigmoid')(X)\n",
        "\n",
        "  model = Model(inputs=X_indices, outputs=X)\n",
        "\n",
        "\n",
        "\n",
        "  return model\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fd5ba8f5-645a-444e-9f5c-5b9e870bcdc5",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "fd5ba8f5-645a-444e-9f5c-5b9e870bcdc5"
      },
      "outputs": [],
      "source": [
        "from keras.models import Sequential\n",
        "from keras.layers import *\n",
        "from tensorflow.keras.preprocessing.sequence import pad_sequences\n",
        "import gensim\n",
        "import sys\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import string\n",
        "import logging\n",
        "import random\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9319431f-ef26-49ee-9588-e3978840b532",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "9319431f-ef26-49ee-9588-e3978840b532"
      },
      "outputs": [],
      "source": [
        "embed_vector_len = reloaded_word_vectors[\"SHL\"].shape[0]\n",
        "vocab_len = len(reloaded_word_vectors.index_to_key)\n",
        "emb_matrix = np.zeros((vocab_len, embed_vector_len))\n",
        "n_categories = len(unique_labels)\n",
        "n_categories"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "90272da8-24f6-4bad-8cc5-543972b17f84",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "90272da8-24f6-4bad-8cc5-543972b17f84"
      },
      "outputs": [],
      "source": [
        "w2v_weights = reloaded_word_vectors.vectors\n",
        "X_train, X_test,Y_train, Y_test = train_test_split(reviews_list, label_numbers, test_size=0.2, random_state = 45)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e37f8938-e11c-49be-a562-65bad672fe4e",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "e37f8938-e11c-49be-a562-65bad672fe4e"
      },
      "outputs": [],
      "source": [
        "# Convert from opcode string to embedding vector\n",
        "a = reloaded_word_vectors.key_to_index\n",
        "\n",
        "train_indices = []\n",
        "for i in X_train:\n",
        "    array = i.split()\n",
        "    result_array=[]\n",
        "    for i in range(len(array)):\n",
        "        if(array[i] == 'EXTCODEHASH'):\n",
        "            continue\n",
        "        result_array.append(a[array[i]])\n",
        "    train_indices.append(result_array)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "98e92a7e-c4cf-49be-b80e-d79dd7b38b4f",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "98e92a7e-c4cf-49be-b80e-d79dd7b38b4f"
      },
      "outputs": [],
      "source": [
        "model = vuln_model((maxLen,))\n",
        "model.summary()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "83ac5dff-b607-4342-b7cc-32ff2df30871",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "83ac5dff-b607-4342-b7cc-32ff2df30871"
      },
      "outputs": [],
      "source": [
        "model = Sequential()\n",
        "model.add(Embedding(input_dim=vocab_len,\n",
        "                    output_dim=embed_vector_len,\n",
        "                    weights=[w2v_weights],\n",
        "                    input_length=maxLen,\n",
        "                    mask_zero=True,\n",
        "                    trainable=False))\n",
        "\n",
        "model.add(Bidirectional(LSTM(100)))\n",
        "model.add(Dense(n_categories, activation='softmax'))\n",
        "\n",
        "model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])\n",
        "\n",
        "X_train_indices = pad_sequences(train_indices, maxlen=maxLen, padding='post')\n",
        "x_train = X_train_indices[23:,:]\n",
        "x_val = X_train_indices[:23,:]\n",
        "y_train = Y_train[23:]\n",
        "y_val = Y_train[:23]\n",
        "history = model.fit(x_train, y_train, batch_size=16, epochs=10,validation_data=(x_val, y_val))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training Plots"
      ],
      "metadata": {
        "id": "wisjlMsagg0L"
      },
      "id": "wisjlMsagg0L"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "12720d7e-f80e-401c-abab-e0b2e50bb0a1",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "12720d7e-f80e-401c-abab-e0b2e50bb0a1"
      },
      "outputs": [],
      "source": [
        "plt.plot(history.history['accuracy'])\n",
        "plt.plot(history.history['val_accuracy'])\n",
        "plt.title('model accuracy')\n",
        "plt.ylabel('accuracy')\n",
        "plt.xlabel('epoch')\n",
        "plt.legend(['train', 'val'], loc='upper left')\n",
        "plt.show()\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1a48b810",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "1a48b810"
      },
      "outputs": [],
      "source": [
        "plt.plot(history.history['loss'])\n",
        "plt.plot(history.history['val_loss'])\n",
        "plt.title('model loss')\n",
        "plt.ylabel('loss')\n",
        "plt.xlabel('epoch')\n",
        "plt.legend(['train', 'val'], loc='upper left')\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Test"
      ],
      "metadata": {
        "id": "lTajqAVngmPW"
      },
      "id": "lTajqAVngmPW"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "195f2cd9",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "195f2cd9"
      },
      "outputs": [],
      "source": [
        "tokenizer = tf.keras.preprocessing.text.Tokenizer()\n",
        "tokenizer.fit_on_texts(X_test)\n",
        "words_to_index = tokenizer.word_index\n",
        "\n",
        "a = reloaded_word_vectors.key_to_index\n",
        "\n",
        "test_indices = []\n",
        "for i in X_test:\n",
        "    array = i.split()\n",
        "    result_array=[]\n",
        "    for i in range(len(array)):\n",
        "        if(array[i] == 'EXTCODEHASH'):\n",
        "            continue\n",
        "        result_array.append(a[array[i]])\n",
        "    test_indices.append(result_array)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "22ebbf16",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "22ebbf16"
      },
      "outputs": [],
      "source": [
        "X_test_indices = pad_sequences(test_indices, maxlen=maxLen,padding='post')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a2c27b57-7426-4efb-882e-606f2920c67c",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "a2c27b57-7426-4efb-882e-606f2920c67c"
      },
      "outputs": [],
      "source": [
        "model.evaluate(X_test_indices,  Y_test, verbose=2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e5c5b122-dba2-49be-9db1-bad7e9ce25d4",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "e5c5b122-dba2-49be-9db1-bad7e9ce25d4"
      },
      "outputs": [],
      "source": [
        "X_test_indices = pad_sequences(test_indices, maxlen=maxLen, padding='post')\n",
        "\n",
        "model.evaluate(X_test_indices, Y_test)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e2e5ee49",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "e2e5ee49"
      },
      "outputs": [],
      "source": [
        "from sklearn.model_selection import KFold\n",
        "import pickle\n",
        "import joblib"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#KFold Cross Validation"
      ],
      "metadata": {
        "id": "UcKdeYGIgrH7"
      },
      "id": "UcKdeYGIgrH7"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6bda96dd",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "6bda96dd"
      },
      "outputs": [],
      "source": [
        "from tensorflow.keras.models import Sequential\n",
        "from tensorflow.keras.layers import Embedding, Bidirectional, LSTM, Dense\n",
        "\n",
        "def create_model():\n",
        "    model = Sequential()\n",
        "    model.add(Embedding(input_dim=vocab_len,\n",
        "                        output_dim=embed_vector_len,\n",
        "                        weights=[w2v_weights],\n",
        "                        input_length=maxLen,\n",
        "                        mask_zero=True,\n",
        "                        trainable=False))\n",
        "    model.add(Bidirectional(LSTM(100)))\n",
        "    model.add(Dense(32, activation='relu'))\n",
        "    model.add(Dense(n_categories, activation='sigmoid'))\n",
        "    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])\n",
        "    return model\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "049953da",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "049953da"
      },
      "outputs": [],
      "source": [
        "kf = KFold(n_splits=5, shuffle=True, random_state=42)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cefdc021",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "cefdc021"
      },
      "outputs": [],
      "source": [
        "import joblib\n",
        "import matplotlib.pyplot as plt\n",
        "from sklearn.model_selection import StratifiedKFold\n",
        "from keras.models import Sequential\n",
        "from keras.layers import Embedding, Bidirectional, LSTM, Dense\n",
        "\n",
        "w2v_weights = reloaded_word_vectors.vectors\n",
        "X_train, X_test,Y_train, Y_test = train_test_split(reviews_list, label_numbers, test_size=0.2, random_state = 45)\n",
        "\n",
        "fold_histories = []\n",
        "\n",
        "k_fold = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)\n",
        "\n",
        "for fold, (train_indices, val_indices) in enumerate(k_fold.split(X_train_indices, Y_train), 1):\n",
        "    print(f\"Fold {fold}:\")\n",
        "\n",
        "    model = Sequential()\n",
        "    model.add(Embedding(input_dim=vocab_len,\n",
        "                        output_dim=embed_vector_len,\n",
        "                        weights=[w2v_weights],\n",
        "                        input_length=maxLen,\n",
        "                        mask_zero=True,\n",
        "                        trainable=False))\n",
        "    model.add(Bidirectional(LSTM(100)))\n",
        "    model.add(Dense(n_categories, activation='softmax'))\n",
        "    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])\n",
        "\n",
        "    X_train_fold = X_train_indices[train_indices]\n",
        "    Y_train_fold = Y_train[train_indices]\n",
        "    X_val_fold = X_train_indices[val_indices]\n",
        "    Y_val_fold = Y_train[val_indices]\n",
        "\n",
        "    history = model.fit(X_train_fold, Y_train_fold, batch_size=16, epochs=10, validation_data=(X_val_fold, Y_val_fold))\n",
        "\n",
        "    fold_histories.append(history)\n",
        "\n",
        "    model_filename = f'model_fold_{fold}.pkl'\n",
        "    joblib.dump(model, model_filename)\n",
        "\n",
        "    loaded_model = joblib.load(model_filename)\n",
        "\n",
        "    loss, accuracy = loaded_model.evaluate(X_val_fold, Y_val_fold)\n",
        "    print(f\"Validation Loss: {loss:.4f}, Validation Accuracy: {accuracy:.4f}\")\n",
        "\n",
        "for fold, history in enumerate(fold_histories, 1):\n",
        "    plt.figure(figsize=(12, 4))\n",
        "    plt.subplot(1, 2, 1)\n",
        "    plt.plot(history.history['loss'], label='Training Loss')\n",
        "    plt.plot(history.history['val_loss'], label='Validation Loss')\n",
        "    plt.xlabel('Epochs')\n",
        "    plt.ylabel('Loss')\n",
        "    plt.legend()\n",
        "    plt.title(f'Fold {fold} - Loss')\n",
        "\n",
        "    plt.subplot(1, 2, 2)\n",
        "    plt.plot(history.history['accuracy'], label='Training Accuracy')\n",
        "    plt.plot(history.history['val_accuracy'], label='Validation Accuracy')\n",
        "    plt.xlabel('Epochs')\n",
        "    plt.ylabel('Accuracy')\n",
        "    plt.legend()\n",
        "    plt.title(f'Fold {fold} - Accuracy')\n",
        "\n",
        "    plt.tight_layout()\n",
        "    plt.show()\n"
      ]
    }
  ],
  "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.9.12"
    },
    "vscode": {
      "interpreter": {
        "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
      }
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}