train_test_phi_fine-tuned.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare the best F1-score on the first 40 epochs and extract the best model\n",
    "\n",
    "Embedded: \n",
    "\"\"\"\n",
    "### CVE description:\n",
    "CVE description\n",
    "\n",
    "### Predicted CWE:\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train examples:  84025\n",
      "Number of classes :  64\n",
      "Test examples:  15508\n",
      "Number of classes :  64\n"
     ]
    }
   ],
   "source": [
    "from collections import Counter\n",
    "import pickle\n",
    "with open('train_final.pickle', 'rb') as f:\n",
    "    train = pickle.load(f)\n",
    "\n",
    "print(\"Train examples: \", len(train))\n",
    "cve_counter = Counter(entry['cwe']for entry in train)\n",
    "print(\"Number of classes : \", len(cve_counter))\n",
    "\n",
    "with open('test_phi-2_decisivo_0.pickle', 'rb') as f:\n",
    "    test = pickle.load(f)\n",
    "\n",
    "print(\"Test examples: \", len(test))\n",
    "cve_counter = Counter(entry['cwe']for entry in test)\n",
    "print(\"Number of classes : \", len(cve_counter))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n",
      "Epoch 1/40\n",
      "263/263 [==============================] - 1s 3ms/step loss\n",
      "Epoch 1 - F1 Score: 0.6872\n",
      "Saved best model\n",
      "[0.6872336935313559]\n",
      "2364/2364 [==============================] - 21s 8ms/step - loss: 1.1965 - accuracy: 0.6779 - val_loss: 1.0902 - val_accuracy: 0.7002\n",
      "Epoch 2/40\n",
      "263/263 [==============================] - 1s 3ms/step loss:\n",
      "Epoch 2 - F1 Score: 0.7084\n",
      "Saved best model\n",
      "[0.6872336935313559, 0.7083562236747423]\n",
      "2364/2364 [==============================] - 19s 8ms/step - loss: 1.0232 - accuracy: 0.7124 - val_loss: 1.0039 - val_accuracy: 0.7152\n",
      "Epoch 3/40\n",
      "263/263 [==============================] - 1s 3ms/step lo\n",
      "Epoch 3 - F1 Score: 0.7138\n",
      "Saved best model\n",
      "[0.6872336935313559, 0.7083562236747423, 0.7138461665905502]\n",
      "2364/2364 [==============================] - 20s 9ms/step - loss: 0.9825 - accuracy: 0.7201 - val_loss: 0.9911 - val_accuracy: 0.7189\n",
      "Epoch 4/40\n",
      "263/263 [==============================] - 1s 3ms/step loss\n",
      "2364/2364 [==============================] - 27s 11ms/step - loss: 0.9508 - accuracy: 0.7280 - val_loss: 0.9704 - val_accuracy: 0.7216\n",
      "Epoch 5/40\n",
      "263/263 [==============================] - 1s 2ms/step loss: 0.9\n",
      "Epoch 5 - F1 Score: 0.7232\n",
      "Saved best model\n",
      "[0.6872336935313559, 0.7083562236747423, 0.7138461665905502, 0.713544042278393, 0.7232258550856164]\n",
      "2364/2364 [==============================] - 23s 10ms/step - loss: 0.9286 - accuracy: 0.7322 - val_loss: 0.9462 - val_accuracy: 0.7270\n",
      "Epoch 6/40\n",
      "263/263 [==============================] - 1s 2ms/step loss: 0\n",
      "2364/2364 [==============================] - 19s 8ms/step - loss: 0.9093 - accuracy: 0.7358 - val_loss: 0.9834 - val_accuracy: 0.7181\n",
      "Epoch 7/40\n",
      "263/263 [==============================] - 1s 3ms/step\n",
      "Epoch 7 - F1 Score: 0.7233\n",
      "Saved best model\n",
      "[0.6872336935313559, 0.7083562236747423, 0.7138461665905502, 0.713544042278393, 0.7232258550856164, 0.715572613256233, 0.7232557763603135]\n",
      "2364/2364 [==============================] - 26s 11ms/step - loss: 0.8948 - accuracy: 0.7393 - val_loss: 0.9617 - val_accuracy: 0.7282\n",
      "Epoch 8/40\n",
      "263/263 [==============================] - 1s 3ms/step loss\n",
      "2364/2364 [==============================] - 27s 11ms/step - loss: 0.8810 - accuracy: 0.7426 - val_loss: 0.9641 - val_accuracy: 0.7241\n",
      "Epoch 9/40\n",
      "263/263 [==============================] - 1s 3ms/step\n",
      "Epoch 9 - F1 Score: 0.7239\n",
      "Saved best model\n",
      "[0.6872336935313559, 0.7083562236747423, 0.7138461665905502, 0.713544042278393, 0.7232258550856164, 0.715572613256233, 0.7232557763603135, 0.7159720183330646, 0.7238847368762382]\n",
      "2364/2364 [==============================] - 22s 9ms/step - loss: 0.8672 - accuracy: 0.7459 - val_loss: 0.9772 - val_accuracy: 0.7264\n",
      "Epoch 10/40\n",
      "263/263 [==============================] - 1s 3ms/step los\n",
      "Epoch 10 - F1 Score: 0.7288\n",
      "Saved best model\n",
      "[0.6872336935313559, 0.7083562236747423, 0.7138461665905502, 0.713544042278393, 0.7232258550856164, 0.715572613256233, 0.7232557763603135, 0.7159720183330646, 0.7238847368762382, 0.7287909994297945]\n",
      "2364/2364 [==============================] - 24s 10ms/step - loss: 0.8575 - accuracy: 0.7475 - val_loss: 0.9501 - val_accuracy: 0.7313\n",
      "Epoch 11/40\n",
      "263/263 [==============================] - 1s 3ms/step los\n",
      "2364/2364 [==============================] - 20s 9ms/step - loss: 0.8463 - accuracy: 0.7502 - val_loss: 0.9563 - val_accuracy: 0.7307\n",
      "Epoch 12/40\n",
      "263/263 [==============================] - 1s 3ms/step los\n",
      "Epoch 12 - F1 Score: 0.7312\n",
      "Saved best model\n",
      "[0.6872336935313559, 0.7083562236747423, 0.7138461665905502, 0.713544042278393, 0.7232258550856164, 0.715572613256233, 0.7232557763603135, 0.7159720183330646, 0.7238847368762382, 0.7287909994297945, 0.7264359811387763, 0.7311699698422199]\n",
      "2364/2364 [==============================] - 20s 9ms/step - loss: 0.8380 - accuracy: 0.7509 - val_loss: 0.9577 - val_accuracy: 0.7364\n",
      "Epoch 13/40\n",
      "263/263 [==============================] - 1s 3ms/step\n",
      "2364/2364 [==============================] - 19s 8ms/step - loss: 0.8269 - accuracy: 0.7548 - val_loss: 0.9457 - val_accuracy: 0.7339\n",
      "Epoch 14/40\n",
      "263/263 [==============================] - 1s 3ms/step \n",
      "2364/2364 [==============================] - 20s 9ms/step - loss: 0.8188 - accuracy: 0.7548 - val_loss: 1.0076 - val_accuracy: 0.7208\n",
      "Epoch 15/40\n",
      "263/263 [==============================] - 1s 2ms/step loss: 0.814\n",
      "2364/2364 [==============================] - 22s 9ms/step - loss: 0.8142 - accuracy: 0.7576 - val_loss: 0.9707 - val_accuracy: 0.7305\n",
      "Epoch 16/40\n",
      "263/263 [==============================] - 1s 3ms/step\n",
      "2364/2364 [==============================] - 22s 9ms/step - loss: 0.8059 - accuracy: 0.7589 - val_loss: 0.9787 - val_accuracy: 0.7257\n",
      "Epoch 17/40\n",
      "263/263 [==============================] - 1s 4ms/step\n",
      "2364/2364 [==============================] - 23s 10ms/step - loss: 0.7987 - accuracy: 0.7606 - val_loss: 0.9667 - val_accuracy: 0.7259\n",
      "Epoch 18/40\n",
      "263/263 [==============================] - 1s 3ms/step loss:\n",
      "2364/2364 [==============================] - 26s 11ms/step - loss: 0.7888 - accuracy: 0.7627 - val_loss: 0.9651 - val_accuracy: 0.7305\n",
      "Epoch 19/40\n",
      "263/263 [==============================] - 1s 3ms/step l\n",
      "2364/2364 [==============================] - 20s 9ms/step - loss: 0.7859 - accuracy: 0.7612 - val_loss: 0.9831 - val_accuracy: 0.7376\n",
      "Epoch 20/40\n",
      "263/263 [==============================] - 1s 4ms/step\n",
      "Epoch 20 - F1 Score: 0.7336\n",
      "Saved best model\n",
      "[0.6872336935313559, 0.7083562236747423, 0.7138461665905502, 0.713544042278393, 0.7232258550856164, 0.715572613256233, 0.7232557763603135, 0.7159720183330646, 0.7238847368762382, 0.7287909994297945, 0.7264359811387763, 0.7311699698422199, 0.7309045684602405, 0.7178320173237428, 0.7272710225172093, 0.7187070797830019, 0.7252718359222966, 0.7298785843659817, 0.7292715898473155, 0.7336021111787876]\n",
      "2364/2364 [==============================] - 24s 10ms/step - loss: 0.7808 - accuracy: 0.7656 - val_loss: 0.9764 - val_accuracy: 0.7338\n",
      "Epoch 21/40\n",
      "263/263 [==============================] - 1s 3ms/step\n",
      "2364/2364 [==============================] - 23s 10ms/step - loss: 0.7742 - accuracy: 0.7663 - val_loss: 0.9991 - val_accuracy: 0.7288\n",
      "Epoch 22/40\n",
      "263/263 [==============================] - 1s 3ms/step\n",
      "2364/2364 [==============================] - 25s 10ms/step - loss: 0.7683 - accuracy: 0.7660 - val_loss: 0.9957 - val_accuracy: 0.7303\n",
      "Epoch 23/40\n",
      "263/263 [==============================] - 1s 3ms/step \n",
      "2364/2364 [==============================] - 26s 11ms/step - loss: 0.7633 - accuracy: 0.7686 - val_loss: 0.9605 - val_accuracy: 0.7380\n",
      "Epoch 24/40\n",
      "263/263 [==============================] - 1s 2ms/step loss:\n",
      "2364/2364 [==============================] - 20s 9ms/step - loss: 0.7560 - accuracy: 0.7704 - val_loss: 0.9809 - val_accuracy: 0.7322\n",
      "Epoch 25/40\n",
      "263/263 [==============================] - 1s 3ms/step\n",
      "Epoch 25 - F1 Score: 0.7357\n",
      "Saved best model\n",
      "[0.6872336935313559, 0.7083562236747423, 0.7138461665905502, 0.713544042278393, 0.7232258550856164, 0.715572613256233, 0.7232557763603135, 0.7159720183330646, 0.7238847368762382, 0.7287909994297945, 0.7264359811387763, 0.7311699698422199, 0.7309045684602405, 0.7178320173237428, 0.7272710225172093, 0.7187070797830019, 0.7252718359222966, 0.7298785843659817, 0.7292715898473155, 0.7336021111787876, 0.7217708880175926, 0.7255293232238774, 0.7331332765021082, 0.7278698842947577, 0.7356619274937776]\n",
      "2364/2364 [==============================] - 20s 9ms/step - loss: 0.7523 - accuracy: 0.7715 - val_loss: 0.9578 - val_accuracy: 0.7395\n",
      "Epoch 26/40\n",
      "263/263 [==============================] - 1s 2ms/step loss:\n",
      "2364/2364 [==============================] - 21s 9ms/step - loss: 0.7519 - accuracy: 0.7710 - val_loss: 0.9831 - val_accuracy: 0.7371\n",
      "Epoch 27/40\n",
      "263/263 [==============================] - 1s 3ms/step los\n",
      "2364/2364 [==============================] - 20s 8ms/step - loss: 0.7414 - accuracy: 0.7731 - val_loss: 1.0084 - val_accuracy: 0.7360\n",
      "Epoch 28/40\n",
      "263/263 [==============================] - 1s 2ms/step loss: 0\n",
      "2364/2364 [==============================] - 12s 5ms/step - loss: 0.7389 - accuracy: 0.7728 - val_loss: 0.9836 - val_accuracy: 0.7385\n",
      "Epoch 29/40\n",
      "263/263 [==============================] - 0s 1ms/step loss: 0.7323 - \n",
      "2364/2364 [==============================] - 16s 7ms/step - loss: 0.7325 - accuracy: 0.7750 - val_loss: 0.9936 - val_accuracy: 0.7315\n",
      "Epoch 30/40\n",
      "263/263 [==============================] - 0s 1ms/step loss: 0.7319 - \n",
      "2364/2364 [==============================] - 16s 7ms/step - loss: 0.7318 - accuracy: 0.7746 - val_loss: 1.0117 - val_accuracy: 0.7319\n",
      "Epoch 31/40\n",
      "263/263 [==============================] - 0s 2ms/step loss: 0.7252 \n",
      "2364/2364 [==============================] - 12s 5ms/step - loss: 0.7247 - accuracy: 0.7780 - val_loss: 1.0204 - val_accuracy: 0.7333\n",
      "Epoch 32/40\n",
      "263/263 [==============================] - 1s 2ms/step loss: 0.721\n",
      "Epoch 32 - F1 Score: 0.7364\n",
      "Saved best model\n",
      "[0.6872336935313559, 0.7083562236747423, 0.7138461665905502, 0.713544042278393, 0.7232258550856164, 0.715572613256233, 0.7232557763603135, 0.7159720183330646, 0.7238847368762382, 0.7287909994297945, 0.7264359811387763, 0.7311699698422199, 0.7309045684602405, 0.7178320173237428, 0.7272710225172093, 0.7187070797830019, 0.7252718359222966, 0.7298785843659817, 0.7292715898473155, 0.7336021111787876, 0.7217708880175926, 0.7255293232238774, 0.7331332765021082, 0.7278698842947577, 0.7356619274937776, 0.7320807256772742, 0.729076020529005, 0.7334947728376845, 0.726209862999425, 0.729895876147236, 0.7305730770196492, 0.7364267306722198]\n",
      "2364/2364 [==============================] - 14s 6ms/step - loss: 0.7215 - accuracy: 0.7776 - val_loss: 1.0074 - val_accuracy: 0.7376\n",
      "Epoch 33/40\n",
      "263/263 [==============================] - 1s 3ms/step\n",
      "Epoch 33 - F1 Score: 0.7383\n",
      "Saved best model\n",
      "[0.6872336935313559, 0.7083562236747423, 0.7138461665905502, 0.713544042278393, 0.7232258550856164, 0.715572613256233, 0.7232557763603135, 0.7159720183330646, 0.7238847368762382, 0.7287909994297945, 0.7264359811387763, 0.7311699698422199, 0.7309045684602405, 0.7178320173237428, 0.7272710225172093, 0.7187070797830019, 0.7252718359222966, 0.7298785843659817, 0.7292715898473155, 0.7336021111787876, 0.7217708880175926, 0.7255293232238774, 0.7331332765021082, 0.7278698842947577, 0.7356619274937776, 0.7320807256772742, 0.729076020529005, 0.7334947728376845, 0.726209862999425, 0.729895876147236, 0.7305730770196492, 0.7364267306722198, 0.738315855353733]\n",
      "2364/2364 [==============================] - 17s 7ms/step - loss: 0.7154 - accuracy: 0.7800 - val_loss: 1.0174 - val_accuracy: 0.7407\n",
      "Epoch 34/40\n",
      "263/263 [==============================] - 0s 2ms/step loss: 0.7116 - \n",
      "2364/2364 [==============================] - 15s 6ms/step - loss: 0.7120 - accuracy: 0.7800 - val_loss: 1.0206 - val_accuracy: 0.7365\n",
      "Epoch 35/40\n",
      "263/263 [==============================] - 0s 2ms/step loss: 0.7084 \n",
      "2364/2364 [==============================] - 13s 6ms/step - loss: 0.7083 - accuracy: 0.7807 - val_loss: 1.0499 - val_accuracy: 0.7310\n",
      "Epoch 36/40\n",
      "263/263 [==============================] - 0s 2ms/step loss: 0.7071 \n",
      "2364/2364 [==============================] - 13s 5ms/step - loss: 0.7070 - accuracy: 0.7794 - val_loss: 1.0372 - val_accuracy: 0.7305\n",
      "Epoch 37/40\n",
      "263/263 [==============================] - 1s 2ms/step loss:\n",
      "Epoch 37 - F1 Score: 0.7387\n",
      "Saved best model\n",
      "[0.6872336935313559, 0.7083562236747423, 0.7138461665905502, 0.713544042278393, 0.7232258550856164, 0.715572613256233, 0.7232557763603135, 0.7159720183330646, 0.7238847368762382, 0.7287909994297945, 0.7264359811387763, 0.7311699698422199, 0.7309045684602405, 0.7178320173237428, 0.7272710225172093, 0.7187070797830019, 0.7252718359222966, 0.7298785843659817, 0.7292715898473155, 0.7336021111787876, 0.7217708880175926, 0.7255293232238774, 0.7331332765021082, 0.7278698842947577, 0.7356619274937776, 0.7320807256772742, 0.729076020529005, 0.7334947728376845, 0.726209862999425, 0.729895876147236, 0.7305730770196492, 0.7364267306722198, 0.738315855353733, 0.7333107723004711, 0.7258754005730349, 0.7273315101546731, 0.7387283652068511]\n",
      "2364/2364 [==============================] - 19s 8ms/step - loss: 0.7036 - accuracy: 0.7818 - val_loss: 1.0431 - val_accuracy: 0.7404\n",
      "Epoch 38/40\n",
      "263/263 [==============================] - 1s 3ms/step los\n",
      "2364/2364 [==============================] - 21s 9ms/step - loss: 0.6970 - accuracy: 0.7837 - val_loss: 1.0776 - val_accuracy: 0.7310\n",
      "Epoch 39/40\n",
      "263/263 [==============================] - 1s 3ms/step \n",
      "2364/2364 [==============================] - 20s 9ms/step - loss: 0.6974 - accuracy: 0.7838 - val_loss: 1.0621 - val_accuracy: 0.7307\n",
      "Epoch 40/40\n",
      "263/263 [==============================] - 1s 3ms/step l\n",
      "2364/2364 [==============================] - 20s 8ms/step - loss: 0.6932 - accuracy: 0.7843 - val_loss: 1.0834 - val_accuracy: 0.7245\n",
      "485/485 [==============================] - 1s 3ms/step\n",
      "Classification Report:\n",
      "               precision    recall  f1-score   support\n",
      "\n",
      "         119     0.8850    0.4131    0.5633      1174\n",
      "         120     0.4358    0.8522    0.5767       203\n",
      "         125     0.8409    0.8689    0.8547       572\n",
      "         134     0.8182    0.9474    0.8780        19\n",
      "         189     0.5323    0.8992    0.6688       119\n",
      "         190     0.8367    0.7847    0.8099       209\n",
      "          20     0.3199    0.2563    0.2846       870\n",
      "         200     0.6399    0.5301    0.5799       647\n",
      "         203     0.6279    0.7500    0.6835        36\n",
      "          22     0.8672    0.9198    0.8927       561\n",
      "         254     0.1053    0.0526    0.0702        38\n",
      "         255     0.3542    0.4928    0.4121        69\n",
      "         264     0.7258    0.4982    0.5908       542\n",
      "         269     0.3810    0.5333    0.4444       120\n",
      "         276     0.4151    0.2973    0.3465        74\n",
      "         284     0.2970    0.2326    0.2609       129\n",
      "         287     0.4241    0.7072    0.5302       304\n",
      "         295     0.8358    0.6667    0.7417        84\n",
      "         306     0.4352    0.5000    0.4653        94\n",
      "         310     0.7483    0.8327    0.7882       257\n",
      "         312     0.3889    0.1373    0.2029        51\n",
      "         319     0.5278    0.6909    0.5984        55\n",
      "         326     0.3429    0.3529    0.3478        34\n",
      "         327     0.5000    0.3684    0.4242        38\n",
      "         345     0.2500    0.3846    0.3030        26\n",
      "         347     0.5152    0.6296    0.5667        27\n",
      "         352     0.7951    0.9399    0.8614       549\n",
      "         362     0.6730    0.7698    0.7181       139\n",
      "         399     0.4196    0.6691    0.5158       269\n",
      "         400     0.3320    0.5755    0.4211       139\n",
      "         401     0.5094    0.4821    0.4954        56\n",
      "         415     0.8864    0.8478    0.8667        46\n",
      "         416     0.8100    0.9076    0.8560       357\n",
      "         426     0.4595    0.7727    0.5763        44\n",
      "         427     0.5128    0.4255    0.4651        47\n",
      "         434     0.6902    0.8341    0.7554       211\n",
      "         476     0.6688    0.8974    0.7664       234\n",
      "         502     0.7826    0.8108    0.7965       111\n",
      "         522     0.4907    0.5699    0.5274        93\n",
      "         532     0.3925    0.8400    0.5350        50\n",
      "          59     0.6822    0.7857    0.7303       112\n",
      "         601     0.7009    0.9036    0.7895        83\n",
      "         611     0.8889    0.9697    0.9275        99\n",
      "         617     0.6400    0.7442    0.6882        43\n",
      "         639     0.4643    0.4194    0.4407        31\n",
      "         668     0.1224    0.2400    0.1622        50\n",
      "         732     0.3577    0.4314    0.3911       102\n",
      "          74     0.1098    0.5488    0.1829        82\n",
      "         755     0.4130    0.5135    0.4578        37\n",
      "          77     0.3868    0.2697    0.3178       152\n",
      "         770     0.3611    0.5270    0.4286        74\n",
      "         772     0.4333    0.3333    0.3768        39\n",
      "          78     0.5850    0.8307    0.6865       319\n",
      "         787     0.7096    0.7740    0.7404       960\n",
      "          79     0.9970    0.7821    0.8766      2552\n",
      "         798     0.7051    0.8527    0.7719       129\n",
      "         835     0.6735    0.7333    0.7021        45\n",
      "         843     0.7674    0.8250    0.7952        40\n",
      "         862     0.5606    0.7200    0.6304       225\n",
      "         863     0.4384    0.2602    0.3265       123\n",
      "          89     0.9961    0.7132    0.8313      1081\n",
      "         908     0.5366    0.7857    0.6377        28\n",
      "         918     0.9655    0.8485    0.9032        99\n",
      "          94     0.5851    0.6405    0.6115       306\n",
      "\n",
      "    accuracy                         0.6749     15508\n",
      "   macro avg     0.5711    0.6311    0.5851     15508\n",
      "weighted avg     0.7237    0.6749    0.6812     15508\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['label_encoder_train.joblib']"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pickle\n",
    "import numpy as np\n",
    "from sklearn.metrics import classification_report, f1_score\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "from tensorflow.keras.models import Sequential\n",
    "from tensorflow.keras.layers import Dense\n",
    "import joblib\n",
    "from keras.callbacks import Callback\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "\n",
    "class F1ScoreCallback(Callback):\n",
    "    def __init__(self, X_val, y_val):\n",
    "        super(F1ScoreCallback, self).__init__()\n",
    "        self.X_val = X_val\n",
    "        self.y_val = y_val\n",
    "        self.best_f1 = 0.0\n",
    "        self.best_model = None\n",
    "        self.f1_scores = []\n",
    "\n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        y_val_pred = np.argmax(self.model.predict(self.X_val), axis=1)\n",
    "        f1 = f1_score(self.y_val, y_val_pred, average='weighted')\n",
    "        self.f1_scores.append(f1)\n",
    "        \n",
    "\n",
    "        if f1 > self.best_f1:\n",
    "            self.best_f1 = f1\n",
    "            self.best_model = self.model\n",
    "            print(f\"Epoch {epoch + 1} - F1 Score: {f1:.4f}\")\n",
    "            print(\"Saved best model\")\n",
    "            print(self.f1_scores)\n",
    "\n",
    "with open('train_final.pickle', 'rb') as f1:\n",
    "    balanced = pickle.load(f1)\n",
    "\n",
    "print(len(balanced[0]))\n",
    "with open('test_phi-2_decisivo_0.pickle', 'rb') as f2:\n",
    "    unbalanced = pickle.load(f2)\n",
    "\n",
    "train = np.array([item['cve_description_phi_mean'] for item in balanced if item['cwe'] != 'None'])\n",
    "test = np.array([item['cwe'] for item in balanced if item['cwe'] != 'None'])\n",
    "np.random.seed(42)\n",
    "X_train, X_val, y_train, y_val = train_test_split(train,test,test_size=0.1,random_state=42)\n",
    "\n",
    "X_test = np.array([item['cve_description_phi_mean'] for item in unbalanced if item['cwe'] != 'None'])\n",
    "y_test = np.array([item['cwe'] for item in unbalanced if item['cwe'] != 'None'])\n",
    "\n",
    "label_encoder_train = LabelEncoder()\n",
    "y_train_encoded = label_encoder_train.fit_transform(y_train)\n",
    "label_encoder_test = LabelEncoder()\n",
    "y_test_encoded = label_encoder_test.fit_transform(y_test)\n",
    "\n",
    "\n",
    "input_dim = X_train.shape[1]\n",
    "output_dim = len(np.unique(y_train))\n",
    "\n",
    "model = Sequential()\n",
    "model.add(Dense(128, input_dim=input_dim, activation='relu'))\n",
    "model.add(Dense(64, activation='relu'))\n",
    "model.add(Dense(output_dim, activation='softmax'))\n",
    "\n",
    "\n",
    "model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])\n",
    "\n",
    "f1_callback = F1ScoreCallback(X_val, label_encoder_train.transform(y_val))\n",
    "\n",
    "history = model.fit(X_train, y_train_encoded, epochs=40, batch_size=32, validation_data=(X_val, label_encoder_train.transform(y_val)), verbose=1, callbacks=[f1_callback])\n",
    "\n",
    "best_model = f1_callback.best_model\n",
    "\n",
    "\n",
    "# Save the best model\n",
    "joblib.dump(best_model, 'CWE_classes.joblib')\n",
    "\n",
    "# Make predictions on the test set\n",
    "y_pred_probs = best_model.predict(X_test)\n",
    "y_pred = np.argmax(y_pred_probs, axis=1)\n",
    "\n",
    "y_pred_original = label_encoder_train.inverse_transform(y_pred)\n",
    "\n",
    "print(\"Classification Report:\\n\", classification_report(y_test, y_pred_original, digits=4))\n",
    "joblib.dump(label_encoder_train, 'label_encoder_train.joblib')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "485/485 [==============================] - 5s 9ms/step\n",
      "Classification Report:\n",
      "               precision    recall  f1-score   support\n",
      "\n",
      "         119     0.8850    0.4131    0.5633      1174\n",
      "         120     0.4358    0.8522    0.5767       203\n",
      "         125     0.8409    0.8689    0.8547       572\n",
      "         134     0.8182    0.9474    0.8780        19\n",
      "         189     0.5323    0.8992    0.6688       119\n",
      "         190     0.8367    0.7847    0.8099       209\n",
      "          20     0.3199    0.2563    0.2846       870\n",
      "         200     0.6399    0.5301    0.5799       647\n",
      "         203     0.6279    0.7500    0.6835        36\n",
      "          22     0.8672    0.9198    0.8927       561\n",
      "         254     0.1053    0.0526    0.0702        38\n",
      "         255     0.3542    0.4928    0.4121        69\n",
      "         264     0.7258    0.4982    0.5908       542\n",
      "         269     0.3810    0.5333    0.4444       120\n",
      "         276     0.4151    0.2973    0.3465        74\n",
      "         284     0.2970    0.2326    0.2609       129\n",
      "         287     0.4241    0.7072    0.5302       304\n",
      "         295     0.8358    0.6667    0.7417        84\n",
      "         306     0.4352    0.5000    0.4653        94\n",
      "         310     0.7483    0.8327    0.7882       257\n",
      "         312     0.3889    0.1373    0.2029        51\n",
      "         319     0.5278    0.6909    0.5984        55\n",
      "         326     0.3429    0.3529    0.3478        34\n",
      "         327     0.5000    0.3684    0.4242        38\n",
      "         345     0.2500    0.3846    0.3030        26\n",
      "         347     0.5152    0.6296    0.5667        27\n",
      "         352     0.7951    0.9399    0.8614       549\n",
      "         362     0.6730    0.7698    0.7181       139\n",
      "         399     0.4196    0.6691    0.5158       269\n",
      "         400     0.3320    0.5755    0.4211       139\n",
      "         401     0.5094    0.4821    0.4954        56\n",
      "         415     0.8864    0.8478    0.8667        46\n",
      "         416     0.8100    0.9076    0.8560       357\n",
      "         426     0.4595    0.7727    0.5763        44\n",
      "         427     0.5128    0.4255    0.4651        47\n",
      "         434     0.6902    0.8341    0.7554       211\n",
      "         476     0.6688    0.8974    0.7664       234\n",
      "         502     0.7826    0.8108    0.7965       111\n",
      "         522     0.4907    0.5699    0.5274        93\n",
      "         532     0.3925    0.8400    0.5350        50\n",
      "          59     0.6822    0.7857    0.7303       112\n",
      "         601     0.7009    0.9036    0.7895        83\n",
      "         611     0.8889    0.9697    0.9275        99\n",
      "         617     0.6400    0.7442    0.6882        43\n",
      "         639     0.4643    0.4194    0.4407        31\n",
      "         668     0.1224    0.2400    0.1622        50\n",
      "         732     0.3577    0.4314    0.3911       102\n",
      "          74     0.1098    0.5488    0.1829        82\n",
      "         755     0.4130    0.5135    0.4578        37\n",
      "          77     0.3868    0.2697    0.3178       152\n",
      "         770     0.3611    0.5270    0.4286        74\n",
      "         772     0.4333    0.3333    0.3768        39\n",
      "          78     0.5850    0.8307    0.6865       319\n",
      "         787     0.7096    0.7740    0.7404       960\n",
      "          79     0.9970    0.7821    0.8766      2552\n",
      "         798     0.7051    0.8527    0.7719       129\n",
      "         835     0.6735    0.7333    0.7021        45\n",
      "         843     0.7674    0.8250    0.7952        40\n",
      "         862     0.5606    0.7200    0.6304       225\n",
      "         863     0.4384    0.2602    0.3265       123\n",
      "          89     0.9961    0.7132    0.8313      1081\n",
      "         908     0.5366    0.7857    0.6377        28\n",
      "         918     0.9655    0.8485    0.9032        99\n",
      "          94     0.5851    0.6405    0.6115       306\n",
      "\n",
      "    accuracy                         0.6749     15508\n",
      "   macro avg     0.5711    0.6311    0.5851     15508\n",
      "weighted avg     0.7237    0.6749    0.6812     15508\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import pickle\n",
    "import numpy as np\n",
    "from sklearn.metrics import classification_report\n",
    "import joblib\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "# Save the best model\n",
    "with open('test_phi-2_decisivo_0.pickle', 'rb') as f2:\n",
    "    unbalanced = pickle.load(f2)\n",
    "\n",
    "X_test = np.array([item['cve_description_phi_mean'] for item in unbalanced if item['cwe'] != 'None'])\n",
    "y_test = np.array([item['cwe'] for item in unbalanced if item['cwe'] != 'None'])\n",
    "\n",
    "best_model=joblib.load('CWE_classes.joblib')\n",
    "label_encoder_train=joblib.load('label_encoder_train.joblib')\n",
    "\n",
    "# Make predictions on the test set\n",
    "y_pred_probs = best_model.predict(X_test)\n",
    "y_pred = np.argmax(y_pred_probs, axis=1)\n",
    "\n",
    "y_pred_original = label_encoder_train.inverse_transform(y_pred)\n",
    "\n",
    "print(\"Classification Report:\\n\", classification_report(y_test, y_pred_original, digits=4))"
   ]
  }
 ],
 "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}