Finetuning+evaluation_VulBERTa-CNN.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import copy\n",
    "from pathlib import Path\n",
    "import pandas as pd\n",
    "import random\n",
    "import re\n",
    "import os\n",
    "import time\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchtext\n",
    "import torchtext.vocab as vocab\n",
    "import sklearn.metrics\n",
    "from transformers import RobertaModel\n",
    "from transformers import RobertaConfig\n",
    "from sklearn.metrics import confusion_matrix\n",
    "from sklearn.utils.class_weight import compute_class_weight\n",
    "from torch.autograd import Variable\n",
    "from torch import nn, optim\n",
    "from torch.optim import SGD,Adam,RMSprop\n",
    "from torch.utils.data import Dataset, DataLoader, IterableDataset\n",
    "from clang import *\n",
    "\n",
    "\n",
    "seed = 1234\n",
    "os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.cuda.manual_seed(seed)\n",
    "torch.cuda.manual_seed_all(seed)\n",
    "np.random.seed(seed)\n",
    "random.seed(seed)\n",
    "torch.backends.cudnn.enabled = True\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "multigpu = False\n",
    "if device == torch.device('cuda'):\n",
    "\tmultigpu = torch.cuda.device_count() > 1\n",
    "print('Device: ',device)\n",
    "print('MultiGPU: ',multigpu)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Training & vocab parameters\n",
    "DATA_PATH = 'data'\n",
    "VOCAB_SIZE = 50000\n",
    "BATCH_SIZE = 128\n",
    "EMBED_SIZE = VOCAB_SIZE+2\n",
    "EMBED_DIM = 768 #768"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Tokenizer\n",
    "\n",
    "from tokenizers.pre_tokenizers import PreTokenizer\n",
    "from tokenizers.pre_tokenizers import Whitespace\n",
    "from tokenizers import NormalizedString,PreTokenizedString\n",
    "from typing import List \n",
    "\n",
    "class MyTokenizer:\n",
    "    \n",
    "    cidx = cindex.Index.create()\n",
    "        \n",
    "\n",
    "    def clang_split(self, i: int, normalized_string: NormalizedString) -> List[NormalizedString]:\n",
    "        ## Tokkenize using clang\n",
    "        tok = []\n",
    "        tu = self.cidx.parse('tmp.c',\n",
    "                       args=[''],  \n",
    "                       unsaved_files=[('tmp.c', str(normalized_string.original))],  \n",
    "                       options=0)\n",
    "        for t in tu.get_tokens(extent=tu.cursor.extent):\n",
    "            spelling = t.spelling.strip()\n",
    "            \n",
    "            if spelling == '':\n",
    "                continue\n",
    "                \n",
    "            ## Keyword no need\n",
    "\n",
    "            ## Punctuations no need\n",
    "\n",
    "            ## Literal all to BPE\n",
    "            \n",
    "            #spelling = spelling.replace(' ', '')\n",
    "            tok.append(NormalizedString(spelling))\n",
    "\n",
    "        return(tok)\n",
    "    \n",
    "    def pre_tokenize(self, pretok: PreTokenizedString):\n",
    "        pretok.split(self.clang_split)\n",
    "        \n",
    "## Custom tokenizer\n",
    "\n",
    "from tokenizers import Tokenizer\n",
    "from tokenizers import normalizers,decoders\n",
    "from tokenizers.normalizers import StripAccents, unicode_normalizer_from_str, Replace\n",
    "from tokenizers.processors import TemplateProcessing\n",
    "from tokenizers import processors,pre_tokenizers\n",
    "from tokenizers.models import BPE\n",
    "\n",
    "## Load pre-trained tokenizers\n",
    "vocab, merges = BPE.read_file(vocab=\"./tokenizer/drapgh-vocab.json\", merges=\"./tokenizer/drapgh-merges.txt\")\n",
    "my_tokenizer = Tokenizer(BPE(vocab, merges, unk_token=\"<unk>\"))\n",
    "\n",
    "my_tokenizer.normalizer = normalizers.Sequence([StripAccents(), Replace(\" \", \"Ä\")])\n",
    "my_tokenizer.pre_tokenizer = PreTokenizer.custom(MyTokenizer())\n",
    "my_tokenizer.post_processor = processors.ByteLevel(trim_offsets=False)\n",
    "my_tokenizer.post_processor = TemplateProcessing(\n",
    "    single=\"<s> $A </s>\",\n",
    "    special_tokens=[\n",
    "    (\"<s>\",0),\n",
    "    (\"<pad>\",1),\n",
    "    (\"</s>\",2),\n",
    "    (\"<unk>\",3),\n",
    "    (\"<mask>\",4)\n",
    "    ]\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PREPARE DATA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TEST_ONLY = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mydataset = 'd2a'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_tokenizer.enable_truncation(max_length=1024)\n",
    "my_tokenizer.enable_padding(direction='right', pad_id=1, pad_type_id=0, pad_token='<pad>', length=None, pad_to_multiple_of=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cleaner(code):\n",
    "    ## Remove code comments\n",
    "    pat = re.compile(r'(/\\*([^*]|(\\*+[^*/]))*\\*+/)|(//.*)')\n",
    "    code = re.sub(pat,'',code)\n",
    "    code = re.sub('\\n','',code)\n",
    "    code = re.sub('\\t','',code)\n",
    "    return(code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_encodings(encodings):\n",
    "    input_ids=[]\n",
    "    attention_mask=[]\n",
    "    for enc in encodings:\n",
    "        input_ids.append(enc.ids)\n",
    "        attention_mask.append(enc.attention_mask)\n",
    "    return {'input_ids':input_ids, 'attention_mask':attention_mask}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def replace_colname(x):\n",
    "    try:\n",
    "        x = x.rename(columns={'functionSource': \"func\"})\n",
    "    except:\n",
    "        None\n",
    "\n",
    "    try:\n",
    "        x = x.rename(columns={'code': \"func\"})\n",
    "    except:\n",
    "        None\n",
    "\n",
    "    try:\n",
    "        x = x.rename(columns={'label': \"target\"})\n",
    "    except:\n",
    "        None\n",
    "    return(x)\n",
    "\n",
    "if mydataset =='devign':\n",
    "    if TEST_ONLY:\n",
    "        \n",
    "        test_index=set()\n",
    "        with open('data/finetune/devign/test.txt') as f:\n",
    "            for line in f:\n",
    "                line=line.strip()\n",
    "                test_index.add(int(line))\n",
    "        mydata = pd.read_json('data/finetune/devign/Devign.json')\n",
    "        m3=mydata.iloc[list(test_index)]\n",
    "        mydata = None\n",
    "        del(mydata)\n",
    "        \n",
    "    else:\n",
    "        train_index=set()\n",
    "        valid_index=set()\n",
    "        test_index=set()\n",
    "\n",
    "        with open('data/finetune/devign/train.txt') as f:\n",
    "            for line in f:\n",
    "                line=line.strip()\n",
    "                train_index.add(int(line))\n",
    "\n",
    "        with open('data/finetune/devign/valid.txt') as f:\n",
    "            for line in f:\n",
    "                line=line.strip()\n",
    "                valid_index.add(int(line))\n",
    "\n",
    "        with open('data/finetune/devign/test.txt') as f:\n",
    "            for line in f:\n",
    "                line=line.strip()\n",
    "                test_index.add(int(line))\n",
    "\n",
    "        mydata = pd.read_json('data/finetune/devign/Devign.json')\n",
    "        m1=mydata.iloc[list(train_index)]\n",
    "        m2=mydata.iloc[list(valid_index)]\n",
    "        m3=mydata.iloc[list(test_index)]\n",
    "\n",
    "        mydata = None\n",
    "        del(mydata)\n",
    "    \n",
    "\n",
    "elif mydataset =='d2a':\n",
    "    task = 'function'\n",
    "    \n",
    "    if TEST_ONLY:\n",
    "        m3 = pd.read_csv('data/finetune/%s/%s/d2a_lbv1_%s_dev.csv'%(mydataset,task,task))\n",
    "        m3 = replace_colname(m3)\n",
    "    else:\n",
    "        m1 = pd.read_csv('data/finetune/%s/%s/d2a_lbv1_%s_train.csv'%(mydataset,task,task))\n",
    "        m2 = pd.read_csv('data/finetune/%s/%s/d2a_lbv1_%s_dev.csv'%(mydataset,task,task))\n",
    "        m3 = pd.read_csv('data/finetune/%s/%s/d2a_lbv1_%s_test.csv'%(mydataset,task,task))\n",
    "       \n",
    "        m1 = replace_colname(m1)\n",
    "        m2 = replace_colname(m2)\n",
    "        m3 = replace_colname(m3)\n",
    "        \n",
    "        \n",
    "else:\n",
    "    \n",
    "    def replace_colname(x):\n",
    "        try:\n",
    "            x = x.rename(columns={'functionSource': \"func\"})\n",
    "        except:\n",
    "            None\n",
    "            \n",
    "        try:\n",
    "            x = x.rename(columns={'code': \"func\"})\n",
    "        except:\n",
    "            None\n",
    "\n",
    "        try:\n",
    "            x = x.rename(columns={'label': \"target\"})\n",
    "        except:\n",
    "            None\n",
    "        return(x)\n",
    "    \n",
    "    \n",
    "    if TEST_ONLY:\n",
    "        m3 = pd.read_pickle('data/finetune/%s/%s_test.pkl'%(mydataset,mydataset))\n",
    "        m3 = replace_colname(m3)\n",
    "        \n",
    "    else:\n",
    "        m1 = pd.read_pickle('data/finetune/%s/%s_train.pkl'%(mydataset,mydataset))\n",
    "        m2 = pd.read_pickle('data/finetune/%s/%s_val.pkl'%(mydataset,mydataset))\n",
    "        m3 = pd.read_pickle('data/finetune/%s/%s_test.pkl'%(mydataset,mydataset))\n",
    "\n",
    "        m1 = replace_colname(m1)\n",
    "        m2 = replace_colname(m2)\n",
    "        m3 = replace_colname(m3)\n",
    "\n",
    "if TEST_ONLY:\n",
    "    m3.func = m3.func.apply(cleaner)\n",
    "    test_encodings = my_tokenizer.encode_batch(m3.func)\n",
    "    try:\n",
    "        test_encodings = [{'func':enc.ids,'target':lab} for enc,lab in zip(test_encodings,m3.target.tolist())]\n",
    "    except:\n",
    "        test_encodings = [{'func':enc.ids,'target':lab} for enc,lab in zip(test_encodings,(m3['combine']*1).tolist())]\n",
    "\n",
    "else:\n",
    "    \n",
    "    m1.func = m1.func.apply(cleaner)\n",
    "    train_encodings = my_tokenizer.encode_batch(m1.func)\n",
    "    try:\n",
    "        train_encodings = [{'func':enc.ids,'target':lab} for enc,lab in zip(train_encodings,m1.target.tolist())]\n",
    "    except:\n",
    "        train_encodings = [{'func':enc.ids,'target':lab} for enc,lab in zip(train_encodings,(m1['combine']*1).tolist())]\n",
    "\n",
    "\n",
    "    m2.func = m2.func.apply(cleaner)\n",
    "    val_encodings = my_tokenizer.encode_batch(m2.func)\n",
    "    try:\n",
    "        val_encodings = [{'func':enc.ids,'target':lab} for enc,lab in zip(val_encodings,m2.target.tolist())]\n",
    "    except:\n",
    "        val_encodings = [{'func':enc.ids,'target':lab} for enc,lab in zip(val_encodings,(m2['combine']*1).tolist())]\n",
    "\n",
    "        \n",
    "    m3.func = m3.func.apply(cleaner)\n",
    "    test_encodings = my_tokenizer.encode_batch(m3.func)\n",
    "    try:\n",
    "        test_encodings = [{'func':enc.ids,'target':lab} for enc,lab in zip(test_encodings,m3.target.tolist())]\n",
    "    except:\n",
    "        test_encodings = [{'func':enc.ids,'target':lab} for enc,lab in zip(test_encodings,(m3['combine']*1).tolist())]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CODES = torchtext.data.Field(batch_first=True, fix_length=1024,use_vocab=False)\n",
    "LABEL = torchtext.data.LabelField(dtype=torch.long, is_target=True,use_vocab=False)\n",
    "fields = {'func': ('codes', CODES), 'target': ('label', LABEL)}\n",
    "\n",
    "class TabularDataset_From_List(torchtext.data.Dataset):\n",
    "    def __init__(self, input_list, format, fields, skip_header=False, **kwargs):\n",
    "        make_example = {\n",
    "            'json': torchtext.data.Example.fromJSON, 'dict': torchtext.data.Example.fromdict}[format.lower()]\n",
    "\n",
    "        examples = [make_example(item, fields) for item in input_list]\n",
    "\n",
    "        if make_example in (torchtext.data.Example.fromdict, torchtext.data.Example.fromJSON):\n",
    "            fields, field_dict = [], fields\n",
    "            for field in field_dict.values():\n",
    "                if isinstance(field, list):\n",
    "                    fields.extend(field)\n",
    "                else:\n",
    "                    fields.append(field)\n",
    "\n",
    "        super(TabularDataset_From_List, self).__init__(examples, fields, **kwargs)\n",
    "\n",
    "    @classmethod\n",
    "    def splits(cls, path=None, root='.data', train=None, validation=None,\n",
    "               test=None, **kwargs):\n",
    "        if path is None:\n",
    "            path = cls.download(root)\n",
    "        train_data = None if train is None else cls(\n",
    "            train, **kwargs)\n",
    "        val_data = None if validation is None else cls(\n",
    "            validation, **kwargs)\n",
    "        test_data = None if test is None else cls(\n",
    "            test, **kwargs)\n",
    "        return tuple(d for d in (train_data, val_data, test_data)\n",
    "                     if d is not None)\n",
    "\n",
    "\n",
    "## Import the 100K data as TabularDataset\n",
    "\n",
    "if TEST_ONLY:\n",
    "    test_data = TabularDataset_From_List(test_encodings,'dict',fields = fields)\n",
    "else:\n",
    "    train_data = TabularDataset_From_List(train_encodings,'dict',fields = fields)\n",
    "    val_data = TabularDataset_From_List(val_encodings,'dict',fields = fields)\n",
    "    test_data = TabularDataset_From_List(test_encodings,'dict',fields = fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### IF ITERABLE DATASETTEST_ONLY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyDataset(IterableDataset):\n",
    "    \n",
    "    def __init__(self,filename,rcount):\n",
    "     \n",
    "        self.filename=filename\n",
    "        self.len_labels=rcount\n",
    "        super().__init__()\n",
    "                    \n",
    "    def process(self,filename):\n",
    "        import pickle \n",
    "        with open(filename, \"rb\") as f:\n",
    "            while True:\n",
    "                try:\n",
    "                    item = pickle.load(f)\n",
    "                    yield {'input_ids': torch.tensor(item['input_ids']), 'attention_mask':torch.tensor(item['attention_mask']), 'labels':torch.tensor(item['labels'])}\n",
    "                except EOFError:\n",
    "                    break\n",
    "                    \n",
    "    def __len__(self):\n",
    "        return self.len_labels\n",
    "\n",
    "    def __iter__(self):\n",
    "        dataset=self.process(self.filename)          \n",
    "        return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_rcount = len(pd.read_pickle('data/draper/draper_train.pkl'))\n",
    "train_dataset = MyDataset('data/draper/draper_stream_train.pkl', train_rcount)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_rcount = len(pd.read_pickle('data/draper/draper_val.pkl'))\n",
    "val_dataset = MyDataset('data/draper/draper_stream_val.pkl', val_rcount)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_rcount = len(pd.read_pickle('data/draper/draper_test.pkl'))\n",
    "test_dataset = MyDataset('data/draper/draper_stream_test.pkl', test_rcount)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### END ITERABLE DATASET"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_VOCAB_SIZE = VOCAB_SIZE\n",
    "\n",
    "# place into iterators\n",
    "\n",
    "if TEST_ONLY:\n",
    "    test_iterator = torchtext.data.BucketIterator(\n",
    "        test_data, \n",
    "        batch_size = 1,\n",
    "        sort = False,\n",
    "        shuffle = False)\n",
    "    \n",
    "else:\n",
    "    train_iterator, valid_iterator, test_iterator = torchtext.data.BucketIterator.splits(\n",
    "        (train_data, val_data, test_data), \n",
    "        batch_size = BATCH_SIZE,\n",
    "        sort = False,\n",
    "        shuffle = False)\n",
    "\n",
    "UNK_IDX = 3\n",
    "PAD_IDX = 1\n",
    "\n",
    "# test_iterator = torchtext.data.BucketIterator(\n",
    "#     test_data, \n",
    "#     batch_size = BATCH_SIZE,\n",
    "#     sort = False,\n",
    "#     shuffle = False)\n",
    "\n",
    "#from torch.utils.data import DataLoader\n",
    "\n",
    "# train_dataloader = DataLoader(train_dataset, shuffle=False, batch_size=BATCH_SIZE)\n",
    "# val_dataloader = DataLoader(val_dataset, shuffle=False, batch_size=BATCH_SIZE)\n",
    "# test_dataloader = DataLoader(test_dataset, shuffle=False, batch_size=BATCH_SIZE)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define VulBERTa-CNN model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class myCNN(nn.Module):\n",
    "    def __init__(self, EMBED_SIZE, EMBED_DIM):\n",
    "        super(myCNN,self).__init__()\n",
    "        \n",
    "        pretrained_weights = RobertaModel.from_pretrained('./models/VulBERTa/').embeddings.word_embeddings.weight\n",
    "\n",
    "        self.embed = nn.Embedding.from_pretrained(pretrained_weights,\n",
    "                                                  freeze=True,\n",
    "                                                  padding_idx=1)\n",
    "\n",
    "        self.conv1 = nn.Conv1d(in_channels=EMBED_DIM, out_channels=200, kernel_size=3)\n",
    "        self.conv2 = nn.Conv1d(in_channels=EMBED_DIM, out_channels=200, kernel_size=4)\n",
    "        self.conv3 = nn.Conv1d(in_channels=EMBED_DIM, out_channels=200, kernel_size=5)\n",
    "\n",
    "        self.dropout = nn.Dropout(0.5)\n",
    "\n",
    "        self.fc1 = nn.Linear(200*3,256) #500\n",
    "        self.fc2 = nn.Linear(256,128)\n",
    "        self.fc3 = nn.Linear(128,2)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.embed(x)\n",
    "        x = x.permute(0,2,1)\n",
    "\n",
    "        x1 = F.relu(self.conv1(x))\n",
    "        x2 = F.relu(self.conv2(x))\n",
    "        x3 = F.relu(self.conv3(x))\n",
    "        \n",
    "        x1 = F.max_pool1d(x1, x1.shape[2])\n",
    "        x2 = F.max_pool1d(x2, x2.shape[2])\n",
    "        x3 = F.max_pool1d(x3, x3.shape[2])\n",
    "        \n",
    "        x = torch.cat([x1,x2,x3],dim=1)\n",
    "        \n",
    "        # flatten the tensor\n",
    "        x = x.flatten(1)\n",
    "        \n",
    "        # apply mean over the last dimension\n",
    "        #x = torch.mean(x, -1)\n",
    "\n",
    "        x = self.dropout(x)\n",
    "\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return(x)\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = myCNN(EMBED_SIZE,EMBED_DIM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#model.embed.weight.data[UNK_IDX] = torch.zeros(EMBED_DIM)\n",
    "model.embed.weight.data[PAD_IDX] = torch.zeros(EMBED_DIM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if multigpu:\n",
    "    model = torch.nn.DataParallel(model)\n",
    "model.to(device)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Num of trainable param: ',sum(p.numel() for p in model.parameters() if p.requires_grad))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare loss function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import sklearn\n",
    "\n",
    "try:\n",
    "    cw = sklearn.utils.class_weight.compute_class_weight(class_weight='balanced',classes=[0,1],y=m1.label.tolist())\n",
    "except:\n",
    "    cw = sklearn.utils.class_weight.compute_class_weight(class_weight='balanced',classes=[0,1],y=m1.target.tolist())\n",
    "    \n",
    "c_weights = torch.FloatTensor([cw[0], cw[1]])\n",
    "criterion = nn.CrossEntropyLoss(weight=c_weights)\n",
    "criterion = criterion.to(device)\n",
    "\n",
    "optimizer = Adam(model.parameters(), lr=0.0005)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def softmax_accuracy(probs,all_labels):\n",
    "    def getClass(x):\n",
    "        return(x.index(max(x)))\n",
    "    \n",
    "    all_labels = all_labels.tolist()\n",
    "    probs = pd.Series(probs.tolist())\n",
    "    all_predicted = probs.apply(getClass)\n",
    "    all_predicted.reset_index(drop=True, inplace=True)\n",
    "    vc = pd.value_counts(all_predicted == all_labels)\n",
    "    try:\n",
    "        acc = vc[1]/len(all_labels)\n",
    "    except:\n",
    "        if(vc.index[0]==False):\n",
    "            acc = 0\n",
    "        else:\n",
    "            acc = 1\n",
    "    return(acc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    model_foldername = 'VB-CNN_%s'%(mydataset)\n",
    "except FileExistsError:\n",
    "    print('Folder exists')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Training started.....')\n",
    "\n",
    "EPOCHS=20\n",
    "BEST_VAL = 9999.9\n",
    "BEST_MODEL = None\n",
    "BEST_EPOCH = None\n",
    "\n",
    "for e in range(EPOCHS):\n",
    "    running_acc = 0\n",
    "    running_loss = 0\n",
    "    timer = time.time()\n",
    "    model.train()\n",
    "\n",
    "    for batch in train_iterator:\n",
    "        batch.codes, batch.label = batch.codes.to(device), batch.label.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(batch.codes)\n",
    "        loss = criterion(output, batch.label)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        acc = softmax_accuracy(output,batch.label)\n",
    "        running_acc += acc\n",
    "        running_loss += loss.item()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        running_acc_val = 0\n",
    "        running_loss_val = 0\n",
    "        for batch in valid_iterator:\n",
    "            batch.codes, batch.label = batch.codes.to(device), batch.label.to(device)\n",
    "            output_val = model(batch.codes)\n",
    "            loss_val = criterion(output_val,batch.label)\n",
    "            acc_val = softmax_accuracy(output_val,batch.label)\n",
    "            running_acc_val += acc_val\n",
    "            running_loss_val += loss_val.item()\n",
    "\n",
    "    print_out = \"Epoch %d - Training acc: %.4f -Training loss: %.4f - Val acc: %.4f - Val loss: %.4f - Time: %.4fs \\n\" % (e+1,\n",
    "    running_acc/len(train_iterator),\n",
    "    running_loss/len(train_iterator),\n",
    "    running_acc_val/len(valid_iterator),\n",
    "    running_loss_val/len(valid_iterator),\n",
    "    (time.time()-timer))\n",
    "    \n",
    "    \n",
    "    selected_model = False\n",
    "    \n",
    "    if selected_model:\n",
    "        \n",
    "        myfile = open(\"res.txt\", \"a\")\n",
    "\n",
    "        if (running_loss_val/len(valid_iterator)) < BEST_VAL:\n",
    "            print('Val_loss decreased!')\n",
    "            print(print_out, end='')\n",
    "            myfile.write('Val_loss decreased!')\n",
    "            myfile.write(print_out)\n",
    "\n",
    "            BEST_VAL = (running_loss_val/len(valid_iterator))\n",
    "            BEST_MODEL = copy.deepcopy(model)\n",
    "            BEST_EPOCH = e+1\n",
    "            model_name = 'models/%s/model_ep_%d.tar' % (model_foldername,e+1)\n",
    "            torch.save({\n",
    "                'epoch': e+1,\n",
    "                'model_state_dict': model.state_dict(),\n",
    "                'optimizer_state_dict': optimizer.state_dict(),\n",
    "                'loss': loss}, model_name)\n",
    "\n",
    "        else:\n",
    "            print(print_out, end='')\n",
    "            myfile.write(print_out)\n",
    "\n",
    "        myfile.close()\n",
    "        \n",
    "    else:\n",
    "        print(print_out, end='')\n",
    "        model_name = 'models/%s/model_ep_%d.tar' % (model_foldername,e+1)\n",
    "        torch.save({\n",
    "            'epoch': e+1,\n",
    "            'model_state_dict': model.state_dict(),\n",
    "            'optimizer_state_dict': optimizer.state_dict(),\n",
    "            'loss': loss}, model_name)\n",
    "\n",
    "        \n",
    "\n",
    "print('Training completed!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define evaluation function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_testing(all_pred, all_labels):\n",
    "    def getClass(x):\n",
    "        return(x.index(max(x)))\n",
    "\n",
    "    probs = pd.Series(all_pred)\n",
    "    all_predicted = probs.apply(getClass)\n",
    "    all_predicted.reset_index(drop=True, inplace=True)\n",
    "    vc = pd.value_counts(all_predicted == all_labels)\n",
    "\n",
    "    probs2=[]\n",
    "    for x in probs:\n",
    "        probs2.append(x[1])\n",
    "\n",
    "    confusion = sklearn.metrics.confusion_matrix(y_true=all_labels, y_pred=all_predicted)\n",
    "    print('Confusion matrix: \\n',confusion)\n",
    "\n",
    "    try:\n",
    "        tn, fp, fn, tp = confusion.ravel()\n",
    "        print('\\nTP:',tp)\n",
    "        print('FP:',fp)\n",
    "        print('TN:',tn)\n",
    "        print('FN:',fn)\n",
    "\n",
    "        ## Performance measure\n",
    "        print('\\nAccuracy: '+ str(sklearn.metrics.accuracy_score(y_true=all_labels, y_pred=all_predicted)))\n",
    "        print('Precision: '+ str(sklearn.metrics.precision_score(y_true=all_labels, y_pred=all_predicted)))\n",
    "        print('F-measure: '+ str(sklearn.metrics.f1_score(y_true=all_labels, y_pred=all_predicted)))\n",
    "        print('Recall: '+ str(sklearn.metrics.recall_score(y_true=all_labels, y_pred=all_predicted)))\n",
    "        print('Precision-Recall AUC: '+ str(sklearn.metrics.average_precision_score(y_true=all_labels, y_score=probs2)))\n",
    "        print('AUC: '+ str(sklearn.metrics.roc_auc_score(y_true=all_labels, y_score=probs2)))\n",
    "        print('MCC: '+ str(sklearn.metrics.matthews_corrcoef(y_true=all_labels, y_pred=all_predicted)))\n",
    "    except:\n",
    "        None\n",
    "        print('This is multiclass prediction')\n",
    "    return(all_predicted)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate on the testing set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Testing started.......')\n",
    "## Testing\n",
    "checkpoint = torch.load('models/VB-CNN_draper/model_ep_15.tar', map_location='cuda')\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    running_acc_test = 0\n",
    "    running_loss_test = 0\n",
    "    all_pred=[]\n",
    "    all_labels=[]\n",
    "    for batch in test_iterator:\n",
    "        batch.codes, batch.label = batch.codes.to(device), batch.label.to(device)\n",
    "        output_test = model(batch.codes).squeeze(1)\n",
    "        loss_test = criterion(output_test,batch.label)\n",
    "        acc_test = softmax_accuracy(output_test,batch.label)\n",
    "        running_acc_test += acc_test\n",
    "        running_loss_test += loss_test.item()\n",
    "        all_pred += output_test.tolist()\n",
    "        all_labels += batch.label.tolist()\n",
    "\n",
    "ap=evaluate_testing(all_pred, all_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Use below only on MVD dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tn=['non-vulnerable','CWE-404','CWE-476','CWE-119','CWE-706','CWE-670','CWE-673','CWE-119, CWE-666, CWE-573','CWE-573','CWE-668','CWE-400, CWE-665, CWE-020','CWE-662','CWE-400','CWE-665','CWE-020','CWE-074','CWE-362','CWE-191','CWE-190','CWE-610','CWE-704','CWE-170','CWE-676','CWE-187','CWE-138','CWE-369','CWE-662, CWE-573','CWE-834','CWE-400, CWE-665','CWE-400, CWE-404','CWE-221','CWE-754','CWE-311','CWE-404, CWE-668','CWE-506','CWE-758','CWE-666','CWE-467','CWE-327','CWE-666, CWE-573','CWE-469']\n",
    "report = sklearn.metrics.classification_report(y_true=all_labels, y_pred=ap, digits=6,labels=np.arange(0,41),target_names=tn)\n",
    "print(report)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "confusion = sklearn.metrics.confusion_matrix(y_true=[1 if x == 0 else 0 for x in all_labels], y_pred=[1 if x == 0 else 0 for x in ap])\n",
    "tn, fp, fn, tp = confusion.ravel()\n",
    "print('\\nTP:',tp)\n",
    "print('FP:',fp)\n",
    "print('TN:',tn)\n",
    "print('FN:',fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_fpr = []\n",
    "w_all_fpr = []\n",
    "aug_y_true_sum = 0\n",
    "for counter in range(41):\n",
    "    aug_y_true = [1 if x == counter else 0 for x in all_labels]\n",
    "    aug_y_pred = [1 if x == counter else 0 for x in ap]\n",
    "    confusion = sklearn.metrics.confusion_matrix(y_true=aug_y_true, y_pred=aug_y_pred)\n",
    "    tn, fp, fn, tp = confusion.ravel()\n",
    "    all_fpr.append(fp/(fp+tn))  ## FPR\n",
    "    w_all_fpr.append((fp/(fp+tn))*aug_y_true.count(1))  ## w_FPR\n",
    "    aug_y_true_sum += aug_y_true.count(1)\n",
    "\n",
    "print('FPR: ', sum(all_fpr)/41.0*100.0)\n",
    "print('Weighted FPR: ', sum(w_all_fpr)/aug_y_true_sum*100.0)"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}