MLM_DNA_BERT.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/ems/elsc-labs/habib-n/yuval.rom/miniconda3/envs/mutation_correction_env_2/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, BertForMaskedLM\n",
    "from transformers.models.bert.configuration_bert import BertConfig\n",
    "from tqdm import tqdm\n",
    "from data_handling_for_MLM import MutationDetectionDataset, collate_fn\n",
    "from torch.utils.data import DataLoader\n",
    "import Levenshtein\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data\n",
    "Tokenizing the normal and mutated data,\n",
    "Marking what token has been changed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "fasta_m = '/ems/elsc-labs/habib-n/yuval.rom/school/ANLP/final_project/Mutation-Simulator/data/sample_data/data_m.fa'\n",
    "fasta_t = '/ems/elsc-labs/habib-n/yuval.rom/school/ANLP/final_project/Mutation-Simulator/data/sample_data/data.fa'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cpu\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/ems/elsc-labs/habib-n/yuval.rom/miniconda3/envs/mutation_correction_env_2/lib/python3.10/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print('Using device:', device)\n",
    "config = BertConfig.from_pretrained(\"zhihan1996/DNABERT-2-117M\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"zhihan1996/DNABERT-2-117M\", trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1, 3, 2]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.get_vocab()\n",
    "tokenizer('[PAD]').input_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataloaders and Padding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x tensor([   1,    4, 1004,   67,   36,  726,  528, 1104,  319,  746,  296,   28,\n",
      "          75, 1507,   55,  362,  123,  130,   82,  443,  184, 2063, 2169,  161,\n",
      "          83,  180,    4,    4,  588,  126,  545,   66,  374, 1602,  283, 1108,\n",
      "         152,  645,  215,    4,  678, 2045,  556, 1176,  727,   97,  173,  448,\n",
      "        1227,  486,   48,  220,   65,   20,    4,  268,   27,  283,  104, 1184,\n",
      "          73, 3532,  245,   61,  208, 3056,  552,  635,   99,  819,   42,  558,\n",
      "         283,   65,  232,  204,   32,  289,   75, 3715,  151,  987, 1435,  226,\n",
      "          33,  411,  149, 3654,  494,  163, 1321,   53, 2975,  112,  131, 1069,\n",
      "           4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,\n",
      "           4,    4,    4,    4,    4,    4,    4,    4,    4,    4, 1127, 2293,\n",
      "         448, 3462, 3454,  942,  307,   82, 2491,   50, 1431,  116,   28,  347,\n",
      "         220,   95,  366,  637,    4,    4,    4,    4,    4,    4,    4,    4,\n",
      "           4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,\n",
      "           4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,\n",
      "           4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,\n",
      "           4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,\n",
      "           4,    4,    4,    4,    4,    4,    4,    4,    4,  776,   86,  177,\n",
      "          95,   64, 3371,    4,  174,  390,   95,   54, 1023,  232,   88,  228,\n",
      "         133,   77,   62,  205,  308,  115,  102,   37,   41,  379,   20,  330,\n",
      "          26,  249,  460,  599, 2662,   92,    4,    4,   95, 1763,  257,  456,\n",
      "          74, 2384,  241,   79,   74,  204, 2384,  192,   73,   88,  121, 1883,\n",
      "         103,    4,    4,  144,   41, 3901,  209,   78, 1355,   34,  487,  124,\n",
      "         415,  100, 1585,  888,  210,  477, 2804, 1539,  731,  552,   32,  800,\n",
      "         536,    4,  102,   30, 1621,  357, 3087,   30,  329,  186,  209,  837,\n",
      "         428,   62,  828, 2523,  232,   40,  418,  552, 1708,   45,   53,   22,\n",
      "        1922,   60,  452,   95,  452,   99,  993,   53, 2041,  112,   67,   59,\n",
      "         102,   50,  134,  111,  177,   44,  105,   50,  130, 1269,  476,  146,\n",
      "         101,   54,    2])\n",
      "labels_mask tensor([-100,  153, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, 1719,   54, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, 1347, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100,  150, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "          28,  619,  285,   78,  108, 1957,   80,  716, 3833,  177,  100,  118,\n",
      "          48,  941, 3115,  519, 1624,  107,   89,  586,  535,  173, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100,   35,  136,   47,  257,   53,  769,  210, 1795,\n",
      "          19,  402,  232,  586, 1889,  403, 3671,  341, 2793,   74,  220,   53,\n",
      "        3667, 1702, 1490, 1767,  215,  108, 2046,  163, 2034,  567,   44,  112,\n",
      "          24,  621,  134, 2288,   53,   57,   26,  166,  348,  281,  177,   89,\n",
      "          26,   61,  822,  567,   32,  241,  261, 2293,  144,   54, 1403,  177,\n",
      "         123,   36,   71,  382, 1533,   28,  136,   54, 2809, -100, -100, -100,\n",
      "        -100, -100, -100, 1158, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, 2278,  209, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100,   34,  506, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100,   50, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100])\n",
      "y tensor([   1,  153, 1004,   67,   36,  726,  528, 1104,  319,  746,  296,   28,\n",
      "          75, 1507,   55,  362,  123,  130,   82,  443,  184, 2063, 2169,  161,\n",
      "          83,  180, 1719,   54,  588,  126,  545,   66,  374, 1602,  283, 1108,\n",
      "         152,  645,  215, 1347,  678, 2045,  556, 1176,  727,   97,  173,  448,\n",
      "        1227,  486,   48,  220,   65,   20,  150,  268,   27,  283,  104, 1184,\n",
      "          73, 3532,  245,   61,  208, 3056,  552,  635,   99,  819,   42,  558,\n",
      "         283,   65,  232,  204,   32,  289,   75, 3715,  151,  987, 1435,  226,\n",
      "          33,  411,  149, 3654,  494,  163, 1321,   53, 2975,  112,  131, 1069,\n",
      "          28,  619,  285,   78,  108, 1957,   80,  716, 3833,  177,  100,  118,\n",
      "          48,  941, 3115,  519, 1624,  107,   89,  586,  535,  173, 1127, 2293,\n",
      "         448, 3462, 3454,  942,  307,   82, 2491,   50, 1431,  116,   28,  347,\n",
      "         220,   95,  366,  637,   35,  136,   47,  257,   53,  769,  210, 1795,\n",
      "          19,  402,  232,  586, 1889,  403, 3671,  341, 2793,   74,  220,   53,\n",
      "        3667, 1702, 1490, 1767,  215,  108, 2046,  163, 2034,  567,   44,  112,\n",
      "          24,  621,  134, 2288,   53,   57,   26,  166,  348,  281,  177,   89,\n",
      "          26,   61,  822,  567,   32,  241,  261, 2293,  144,   54, 1403,  177,\n",
      "         123,   36,   71,  382, 1533,   28,  136,   54, 2809,  776,   86,  177,\n",
      "          95,   64, 3371, 1158,  174,  390,   95,   54, 1023,  232,   88,  228,\n",
      "         133,   77,   62,  205,  308,  115,  102,   37,   41,  379,   20,  330,\n",
      "          26,  249,  460,  599, 2662,   92, 2278,  209,   95, 1763,  257,  456,\n",
      "          74, 2384,  241,   79,   74,  204, 2384,  192,   73,   88,  121, 1883,\n",
      "         103,   34,  506,  144,   41, 3901,  209,   78, 1355,   34,  487,  124,\n",
      "         415,  100, 1585,  888,  210,  477, 2804, 1539,  731,  552,   32,  800,\n",
      "         536,   50,  102,   30, 1621,  357, 3087,   30,  329,  186,  209,  837,\n",
      "         428,   62,  828, 2523,  232,   40,  418,  552, 1708,   45,   53,   22,\n",
      "        1922,   60,  452,   95,  452,   99,  993,   53, 2041,  112,   67,   59,\n",
      "         102,   50,  134,  111,  177,   44,  105,   50,  130, 1269,  476,  146,\n",
      "         101,   54,    2])\n",
      "----------------\n",
      "x tensor([   1,   30,  126,  545,   66,  374, 1602,  283, 1108,  152,  793,  128,\n",
      "         245,   61,  208, 3056,  552,  635,   99,  819,   42,  558,  283,   65,\n",
      "         232,  204,   32,    4,   75, 3715,  151,  987, 1435,  226,   33,  411,\n",
      "         149,   63,  307,   53, 2695,  102, 1183, 3656,   30,  145,  678,  804,\n",
      "         189,   56,   34,  553,   99,   53, 2975,  112,  131, 1069,   28,  619,\n",
      "         285,   78,  108, 1957,   80,  716, 3833,  177,  100,  118,   48,  941,\n",
      "           4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,\n",
      "           4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,\n",
      "           4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,\n",
      "           4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,\n",
      "           4,    4,    4,    4,    4,    4,    4])\n",
      "labels_mask tensor([-100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100,  289, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        3115,  519, 1624,  107,   89,  586,  535,  173, 1127,   42,  132,  517,\n",
      "         860,  136,   47,  257,   53,  769,  210, 1795,   19,  402,  232,  586,\n",
      "        1889,  403, 3671,  341, 2793,   74,  220,   53, 3667, 1702, 1490, 1767,\n",
      "         215,  108, 2046,  163, 2034,  567,   44,  112,   24,  621,  134, 2288,\n",
      "          53,   57,   26,  166,    2,    3,    3])\n",
      "y tensor([   1,   30,  126,  545,   66,  374, 1602,  283, 1108,  152,  793,  128,\n",
      "         245,   61,  208, 3056,  552,  635,   99,  819,   42,  558,  283,   65,\n",
      "         232,  204,   32,  289,   75, 3715,  151,  987, 1435,  226,   33,  411,\n",
      "         149,   63,  307,   53, 2695,  102, 1183, 3656,   30,  145,  678,  804,\n",
      "         189,   56,   34,  553,   99,   53, 2975,  112,  131, 1069,   28,  619,\n",
      "         285,   78,  108, 1957,   80,  716, 3833,  177,  100,  118,   48,  941,\n",
      "        3115,  519, 1624,  107,   89,  586,  535,  173, 1127,   42,  132,  517,\n",
      "         860,  136,   47,  257,   53,  769,  210, 1795,   19,  402,  232,  586,\n",
      "        1889,  403, 3671,  341, 2793,   74,  220,   53, 3667, 1702, 1490, 1767,\n",
      "         215,  108, 2046,  163, 2034,  567,   44,  112,   24,  621,  134, 2288,\n",
      "          53,   57,   26,  166,    2,    3,    3])\n",
      "----------------\n",
      "x tensor([   1,    5,   23,   78,  479, 3270,   32,    4,    4,    4,  158,   70,\n",
      "         886,  232, 1115,  560,   74,   72,  425, 1432,   41,  148,   83,  708,\n",
      "         139,   28,  240,   74,  151,   15,  304,   50,  734,  148, 1425,   32,\n",
      "         307,   33,  382,   11,   53,  173, 3896,   15,  816,  136,   53,   76,\n",
      "          73,  145,  962,  332, 2250,  229,   56,  999, 3812,   76,   35,  205,\n",
      "         932,   19,  192,  290,   70,  694,   47,  866,  601,  203,  205, 2586,\n",
      "          50,  169, 1765, 3431,  462,   66,  541,    4,    4,  987,  247,  519,\n",
      "         203,   25,    4,   33,  752,    4,    4, 2517,   36, 1036, 3342,  499,\n",
      "         151,   71,  226,   92,  161,   95,   66,  197,  365,   22, 1456,   62,\n",
      "         149,   27,  210,   13,  145,   74,  207,   25,  312, 3273, 1799, 2018,\n",
      "         343,  998,  158,   65,  188, 3422,  371,   55,    4,  450,  188,  102,\n",
      "        2409,  197,   66,   50,  208, 3097,  148,  624,   55,  921,  702,   42,\n",
      "         454,   74,  229,  103,  257,   73,    4,   34,  561,   74,  257,   74,\n",
      "         716,   51,  156,   93,  190,  106,   51,   36,  989, 1705,  954,   36,\n",
      "         372,  139, 3535,  173,  637,   28,    4,   64,  442,   59,  138,   36,\n",
      "         307,  413, 1907, 1959,  473,   79,  210,  906, 1002,   95,  407,  906,\n",
      "         241, 2190,   28, 2969,  140,   95,   78,  140,  164,   13, 1965,   74,\n",
      "          99,   76,    4,   53,  203,   93,  148,  349,  511,  777,  115, 1556,\n",
      "         212, 1232,  809,    4, 2305,    4,   62,   89,  178,   37,  611,  293,\n",
      "         156,   62,   92,  278,   15,  303,   76,   62,    4, 2027,   33, 1538,\n",
      "          62,  145,  349,  496,   33,  595,   66,   61, 1112,   93,  156,  885,\n",
      "         219,    4,    4,  271,  361,   62,   67,  221,  547,  332,  304,  128,\n",
      "        2690,  905,   61,   92, 2996,    4,    4,    4,    4,    4,    4,    4,\n",
      "           4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,\n",
      "           4,    4])\n",
      "labels_mask tensor([-100, -100, -100, -100, -100, -100, -100, 2171,  253,   22, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100,   36,  205, -100, -100, -100,\n",
      "        -100, -100,  437, -100, -100,  339,  368, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100,  383, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100,  742, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100,   95, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, 1589, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100,  512, -100,  103, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100,  159, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100,   83,  757, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, 1036,   86, 1765,   30,  346,   40,  293,\n",
      "        2741,   13, 1260,  349,  460,  840,  367,   93, 2849, 1513,  725,    7,\n",
      "           2,    3])\n",
      "y tensor([   1,    5,   23,   78,  479, 3270,   32, 2171,  253,   22,  158,   70,\n",
      "         886,  232, 1115,  560,   74,   72,  425, 1432,   41,  148,   83,  708,\n",
      "         139,   28,  240,   74,  151,   15,  304,   50,  734,  148, 1425,   32,\n",
      "         307,   33,  382,   11,   53,  173, 3896,   15,  816,  136,   53,   76,\n",
      "          73,  145,  962,  332, 2250,  229,   56,  999, 3812,   76,   35,  205,\n",
      "         932,   19,  192,  290,   70,  694,   47,  866,  601,  203,  205, 2586,\n",
      "          50,  169, 1765, 3431,  462,   66,  541,   36,  205,  987,  247,  519,\n",
      "         203,   25,  437,   33,  752,  339,  368, 2517,   36, 1036, 3342,  499,\n",
      "         151,   71,  226,   92,  161,   95,   66,  197,  365,   22, 1456,   62,\n",
      "         149,   27,  210,   13,  145,   74,  207,   25,  312, 3273, 1799, 2018,\n",
      "         343,  998,  158,   65,  188, 3422,  371,   55,  383,  450,  188,  102,\n",
      "        2409,  197,   66,   50,  208, 3097,  148,  624,   55,  921,  702,   42,\n",
      "         454,   74,  229,  103,  257,   73,  742,   34,  561,   74,  257,   74,\n",
      "         716,   51,  156,   93,  190,  106,   51,   36,  989, 1705,  954,   36,\n",
      "         372,  139, 3535,  173,  637,   28,   95,   64,  442,   59,  138,   36,\n",
      "         307,  413, 1907, 1959,  473,   79,  210,  906, 1002,   95,  407,  906,\n",
      "         241, 2190,   28, 2969,  140,   95,   78,  140,  164,   13, 1965,   74,\n",
      "          99,   76, 1589,   53,  203,   93,  148,  349,  511,  777,  115, 1556,\n",
      "         212, 1232,  809,  512, 2305,  103,   62,   89,  178,   37,  611,  293,\n",
      "         156,   62,   92,  278,   15,  303,   76,   62,  159, 2027,   33, 1538,\n",
      "          62,  145,  349,  496,   33,  595,   66,   61, 1112,   93,  156,  885,\n",
      "         219,   83,  757,  271,  361,   62,   67,  221,  547,  332,  304,  128,\n",
      "        2690,  905,   61,   92, 2996, 1036,   86, 1765,   30,  346,   40,  293,\n",
      "        2741,   13, 1260,  349,  460,  840,  367,   93, 2849, 1513,  725,    7,\n",
      "           2,    3])\n",
      "----------------\n",
      "x tensor([   1,   42,  329,   57,   65,  720,   33,  181,  351, 1587,  876,  102,\n",
      "          65,  651, 1005,    7,    2])\n",
      "labels_mask tensor([-100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100])\n",
      "y tensor([   1,   42,  329,   57,   65,  720,   33,  181,  351, 1587,  876,  102,\n",
      "          65,  651, 1005,    7,    2])\n",
      "----------------\n",
      "x tensor([   1, 2089,   72,   78,   65,  134,  161,   13,  212,   67,    4,    4,\n",
      "           4,    4, 1281,  178,   55, 2306, 1176,  209,   90,   86,   26,  278,\n",
      "          62,   64,  102,   56,  143,   15,  156,   62,   26,  156,  261, 1025,\n",
      "          26,  332, 2227, 1220,  103,  169, 2684,  203,  906,  549,  169,   53,\n",
      "          36,  169,   74,  372,  231,   37,    4,    4,    4,    4,    4,    4,\n",
      "           4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,\n",
      "           4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,    4,\n",
      "           4,    4,    4,    4,    4, 2914,  151,  515,   41, 3690,   97, 1689,\n",
      "        3506,  316,    4,    4,  163,  238,   87,    4,    4,  205,  209,    4,\n",
      "           4,    4,  992,  222,   68,  466,   65,  582,   75, 4017,  181,  235,\n",
      "          90,  292,  454,   61,  256,  166,  145,  103,   62, 3301, 1681,   50,\n",
      "        2047,   13,  200, 3574,  769, 2065,  803,  984,   76, 2152,  762, 3061,\n",
      "        3416,    8,    2])\n",
      "labels_mask tensor([-100, -100, -100, -100, -100, -100, -100, -100, -100, -100,  136,   36,\n",
      "         114, 2617, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, 1219,  158,  457,  177,  270,  116,\n",
      "          36,  171, 1559, 1682, 3739,  112,  241, 2316,   54,   81, 2920,  102,\n",
      "        3780,  132,   27, 2039,  560, 3728, 2249,  276, 1229,  553,   87,  135,\n",
      "          30,  114,  370,   41,  377, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100,   32,  239, -100, -100, -100,   90,  184, -100, -100,   47,\n",
      "         277, 2880, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,\n",
      "        -100, -100, -100])\n",
      "y tensor([   1, 2089,   72,   78,   65,  134,  161,   13,  212,   67,  136,   36,\n",
      "         114, 2617, 1281,  178,   55, 2306, 1176,  209,   90,   86,   26,  278,\n",
      "          62,   64,  102,   56,  143,   15,  156,   62,   26,  156,  261, 1025,\n",
      "          26,  332, 2227, 1220,  103,  169, 2684,  203,  906,  549,  169,   53,\n",
      "          36,  169,   74,  372,  231,   37, 1219,  158,  457,  177,  270,  116,\n",
      "          36,  171, 1559, 1682, 3739,  112,  241, 2316,   54,   81, 2920,  102,\n",
      "        3780,  132,   27, 2039,  560, 3728, 2249,  276, 1229,  553,   87,  135,\n",
      "          30,  114,  370,   41,  377, 2914,  151,  515,   41, 3690,   97, 1689,\n",
      "        3506,  316,   32,  239,  163,  238,   87,   90,  184,  205,  209,   47,\n",
      "         277, 2880,  992,  222,   68,  466,   65,  582,   75, 4017,  181,  235,\n",
      "          90,  292,  454,   61,  256,  166,  145,  103,   62, 3301, 1681,   50,\n",
      "        2047,   13,  200, 3574,  769, 2065,  803,  984,   76, 2152,  762, 3061,\n",
      "        3416,    8,    2])\n",
      "----------------\n"
     ]
    }
   ],
   "source": [
    "mutation_dataset = MutationDetectionDataset(fasta_m, fasta_t, tokenizer, verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader = DataLoader(mutation_dataset, batch_size=2, collate_fn=collate_fn, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model\n",
    "We are using [DNABERT2](https://github.com/MAGICS-LAB/DNABERT_2/tree/main?tab=readme-ov-file#1-introduction)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cpu\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print('Using device:', device)\n",
    "lr = 5e-5\n",
    "weight_decay = 0.01\n",
    "batch_size = 2 # TODO: change to 32\n",
    "num_epochs = 3 # TODO: change to ???\n",
    "\n",
    "\n",
    "model = BertForMaskedLM(config).to(device)\n",
    "loss_func = torch.nn.CrossEntropyLoss() \n",
    "# metric = evaluate.load('seqeval')\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
    "train_loader = DataLoader(mutation_dataset, batch_size=batch_size, collate_fn=collate_fn, shuffle=True)\n",
    "eval_loader = DataLoader(mutation_dataset, batch_size=batch_size, collate_fn=collate_fn, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_median_edit_distance(edit_distances):\n",
    "    return np.median(edit_distances)\n",
    "\n",
    "\n",
    "def compute_mean_edit_distance(edit_distances):\n",
    "    return np.mean(edit_distances)\n",
    "\n",
    "\n",
    "def compute_normalized_mean_edit_distance(edit_distances, label_texts):\n",
    "    # Calculate normalized mean edit distance\n",
    "    normalized_edit_distances = [edit_distance / len(label_text) for edit_distance, label_text in zip(edit_distances, label_texts)]\n",
    "    return np.mean(normalized_edit_distances)\n",
    "\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "    logits, labels = eval_pred\n",
    "    predictions = np.argmax(logits, axis=-1)\n",
    "    \n",
    "    # Convert predictions and labels to lists of token IDs\n",
    "    predictions = predictions.tolist()\n",
    "    labels = labels.tolist()\n",
    "    \n",
    "    # Decode predictions and labels, filtering out invalid labels\n",
    "    pred_texts = []\n",
    "    label_texts = []\n",
    "    for pred, label in zip(predictions, labels):\n",
    "        # Filter out -100 (or any invalid token IDs) from labels\n",
    "        valid_label_ids = [id for id in label if id != -100]\n",
    "        # Decode only valid token IDs\n",
    "        pred_texts.append(tokenizer.decode(pred, skip_special_tokens=True))\n",
    "        label_texts.append(tokenizer.decode(valid_label_ids, skip_special_tokens=True))\n",
    "    \n",
    "    # Calculate Levenshtein distance\n",
    "    edit_distances = [Levenshtein.distance(pred, label) for pred, label in zip(pred_texts, label_texts)]\n",
    "    \n",
    "    \n",
    "    return {\"avg_edit_distance\": compute_mean_edit_distance(edit_distances),\n",
    "            \"median_edit_distance\": compute_median_edit_distance(edit_distances),\n",
    "            \"normalized_avg_edit_distance\": compute_normalized_mean_edit_distance(edit_distances, labels)}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='9' max='9' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [9/9 02:09, Epoch 3/3]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Avg Edit Distance</th>\n",
       "      <th>Median Edit Distance</th>\n",
       "      <th>Normalized Avg Edit Distance</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>No log</td>\n",
       "      <td>1.412605</td>\n",
       "      <td>169.400000</td>\n",
       "      <td>175.000000</td>\n",
       "      <td>0.518043</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>No log</td>\n",
       "      <td>1.346656</td>\n",
       "      <td>217.400000</td>\n",
       "      <td>275.000000</td>\n",
       "      <td>0.664832</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>No log</td>\n",
       "      <td>1.273410</td>\n",
       "      <td>167.200000</td>\n",
       "      <td>157.000000</td>\n",
       "      <td>0.511315</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=9, training_loss=3.339323255750868, metrics={'train_runtime': 130.745, 'train_samples_per_second': 0.115, 'train_steps_per_second': 0.069, 'total_flos': 1869102832128.0, 'train_loss': 3.339323255750868, 'epoch': 3.0})"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import Trainer, TrainingArguments\n",
    "\n",
    "# Define training arguments\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./models/correction/results\",  # output directory\n",
    "    overwrite_output_dir=True,\n",
    "    num_train_epochs=num_epochs,\n",
    "    per_device_train_batch_size=batch_size,\n",
    "    save_steps=10_000,\n",
    "    save_total_limit=2,\n",
    "    logging_dir='./models/correction/logs',\n",
    "    evaluation_strategy=\"epoch\",  # Ensure evaluations occur\n",
    "\n",
    ")\n",
    "\n",
    "# Initialize the Trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    data_collator=collate_fn,\n",
    "    train_dataset=mutation_dataset,\n",
    "    eval_dataset=mutation_dataset,\n",
    "    tokenizer=tokenizer,\n",
    "    optimizers=(optimizer, None),\n",
    "    compute_metrics=compute_metrics\n",
    ")\n",
    "\n",
    "# Train the model\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model,\n",
    "                loss_func,\n",
    "                train_dataloader,\n",
    "                eval_dataloader,\n",
    "                lr,\n",
    "                weight_decay,\n",
    "                batch_size,\n",
    "                num_epochs,\n",
    "                device,\n",
    "                optimizer,\n",
    "                lora=False):\n",
    "    # if not lora:\n",
    "    #     optimizer = AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)\n",
    "    for epoch in range(num_epochs):\n",
    "        # Training\n",
    "        model.train()\n",
    "        train_loss = 0\n",
    "        for batch in tqdm(train_dataloader, desc=f\"Training Epoch {epoch}\"):\n",
    "            torch.cuda.empty_cache()\n",
    "            batch = {k: v.to(device) for k, v in batch.items()}\n",
    "            outputs = model(**batch)\n",
    "            loss = loss_func(outputs.logits, batch['labels'])\n",
    "            train_loss += loss.item()\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        avg_train_loss = train_loss / len(train_dataloader)\n",
    "        print(f\"Average train loss: {avg_train_loss}\")\n",
    "\n",
    "        # Validation\n",
    "        model.eval()\n",
    "        eval_loss = 0\n",
    "        correct = 0\n",
    "        with torch.no_grad():\n",
    "            for eval_batch in eval_dataloader:\n",
    "                eval_batch = {k: v.to(device) for k, v in eval_batch.items()}\n",
    "                outputs = model(**eval_batch)\n",
    "                loss = outputs.loss\n",
    "                correct += (outputs.logits.argmax(dim=1) == eval_batch['labels']).float().sum()\n",
    "                eval_loss += loss.item()\n",
    "\n",
    "            print(f\"Average eval loss: {eval_loss / len(eval_dataloader)}\")\n",
    "            accuracy = correct / (len(eval_dataloader) * batch_size)\n",
    "            print(f\"Eval Accuracy: {accuracy}\")\n",
    "\n",
    "    model.save_pretrained(f\"models/correction/fine_tuned_model_e{num_epochs}_bc{batch_size}_lr{lr}_wd{weight_decay}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training Epoch 0:   0%|          | 0/3 [00:00<?, ?it/s]\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "Expected target size [2, 4096], got [2, 327]",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[13], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtrain_model\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      2\u001b[0m \u001b[43m            \u001b[49m\u001b[43mloss_func\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      3\u001b[0m \u001b[43m            \u001b[49m\u001b[43mtrain_loader\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      4\u001b[0m \u001b[43m            \u001b[49m\u001b[43meval_loader\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      5\u001b[0m \u001b[43m            \u001b[49m\u001b[43mlr\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      6\u001b[0m \u001b[43m            \u001b[49m\u001b[43mweight_decay\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      7\u001b[0m \u001b[43m            \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      8\u001b[0m \u001b[43m            \u001b[49m\u001b[43mnum_epochs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      9\u001b[0m \u001b[43m            \u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     10\u001b[0m \u001b[43m            \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[12], line 22\u001b[0m, in \u001b[0;36mtrain_model\u001b[0;34m(model, loss_func, train_dataloader, eval_dataloader, lr, weight_decay, batch_size, num_epochs, device, optimizer, lora)\u001b[0m\n\u001b[1;32m     20\u001b[0m batch \u001b[38;5;241m=\u001b[39m {k: v\u001b[38;5;241m.\u001b[39mto(device) \u001b[38;5;28;01mfor\u001b[39;00m k, v \u001b[38;5;129;01min\u001b[39;00m batch\u001b[38;5;241m.\u001b[39mitems()}\n\u001b[1;32m     21\u001b[0m outputs \u001b[38;5;241m=\u001b[39m model(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mbatch)\n\u001b[0;32m---> 22\u001b[0m loss \u001b[38;5;241m=\u001b[39m \u001b[43mloss_func\u001b[49m\u001b[43m(\u001b[49m\u001b[43moutputs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlogits\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mlabels\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     23\u001b[0m train_loss \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m loss\u001b[38;5;241m.\u001b[39mitem()\n\u001b[1;32m     24\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mzero_grad()\n",
      "File \u001b[0;32m~/miniconda3/envs/mutation_correction_env_2/lib/python3.10/site-packages/torch/nn/modules/module.py:1532\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1530\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1531\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1532\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/mutation_correction_env_2/lib/python3.10/site-packages/torch/nn/modules/module.py:1541\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1536\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1537\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1538\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1539\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1540\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1541\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1543\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1544\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/mutation_correction_env_2/lib/python3.10/site-packages/torch/nn/modules/loss.py:1185\u001b[0m, in \u001b[0;36mCrossEntropyLoss.forward\u001b[0;34m(self, input, target)\u001b[0m\n\u001b[1;32m   1184\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Tensor, target: Tensor) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tensor:\n\u001b[0;32m-> 1185\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcross_entropy\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtarget\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mweight\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1186\u001b[0m \u001b[43m                           \u001b[49m\u001b[43mignore_index\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mignore_index\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mreduction\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mreduction\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1187\u001b[0m \u001b[43m                           \u001b[49m\u001b[43mlabel_smoothing\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlabel_smoothing\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/mutation_correction_env_2/lib/python3.10/site-packages/torch/nn/functional.py:3086\u001b[0m, in \u001b[0;36mcross_entropy\u001b[0;34m(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)\u001b[0m\n\u001b[1;32m   3084\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m size_average \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mor\u001b[39;00m reduce \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m   3085\u001b[0m     reduction \u001b[38;5;241m=\u001b[39m _Reduction\u001b[38;5;241m.\u001b[39mlegacy_get_string(size_average, reduce)\n\u001b[0;32m-> 3086\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_C\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_nn\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcross_entropy_loss\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtarget\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m_Reduction\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_enum\u001b[49m\u001b[43m(\u001b[49m\u001b[43mreduction\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mignore_index\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlabel_smoothing\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Expected target size [2, 4096], got [2, 327]"
     ]
    }
   ],
   "source": [
    "train_model(model,\n",
    "            loss_func,\n",
    "            train_loader,\n",
    "            eval_loader,\n",
    "            lr,\n",
    "            weight_decay,\n",
    "            batch_size,\n",
    "            num_epochs,\n",
    "            device,\n",
    "            optimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mutation_sim",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}