demos/llama_dna_sequential_finetune_QLoRA.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7c364732",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2bf01f08",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import transformers\n",
    "import pandas as pd\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments\n",
    "from transformers import BitsAndBytesConfig\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from dataclasses import dataclass, field\n",
    "from typing import Optional, Dict, Sequence, Tuple, List\n",
    "from peft import LoraConfig, get_peft_model\n",
    "from sklearn.metrics import precision_score, recall_score, f1_score, matthews_corrcoef\n",
    "from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score\n",
    "import numpy as np\n",
    "from scipy.special import softmax\n",
    "import sklearn.metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "80e247a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class TrainingArguments(transformers.TrainingArguments):\n",
    "    cache_dir: Optional[str] = field(default=None)\n",
    "    run_name: str = field(default=\"run\")\n",
    "    optim: str = field(default=\"adamw_torch\")\n",
    "    model_max_length: int = field(default=256, metadata={\"help\": \"Maximum sequence length.\"})\n",
    "    gradient_accumulation_steps: int = field(default=1)\n",
    "    per_device_train_batch_size: int = field(default=8)\n",
    "    per_device_eval_batch_size: int = field(default=4)\n",
    "    num_train_epochs: int = field(default=5)\n",
    "    fp16: bool = field(default=False)\n",
    "    #logging_steps: int = field(default=1000)\n",
    "    save_steps: int = field(default=200)\n",
    "    eval_steps: int = field(default=200)\n",
    "    evaluation_strategy: str = field(default=\"steps\")\n",
    "    load_best_model_at_end: bool = field(default=True)     # load the best model when finished training (default metric is loss)\n",
    "    metric_for_best_model: str = field(default=\"eval_loss\") # the metric to use to compare models\n",
    "    greater_is_better: bool = field(default=False)           # whether the `metric_for_best_model` should be maximized or not\n",
    "    logging_strategy: str = field(default=\"steps\")  # Log every \"steps\"\n",
    "    logging_steps: int = field(default=200)  # Log every 100 steps\n",
    "    warmup_ratio: int = field(default=0.1)\n",
    "    weight_decay: float = field(default=1e-2)\n",
    "    learning_rate: float = field(default=1e-5)\n",
    "    lr_scheduler_type: str = field(default='linear')\n",
    "    save_total_limit: int = field(default=5)\n",
    "    load_best_model_at_end: bool = field(default=True)\n",
    "    output_dir: str = field(default=\"/common/zhangz2lab/zhanh/Jupyter_Scripts/output_0828/llama_results\")\n",
    "    find_unused_parameters: bool = field(default=False)\n",
    "    checkpointing: bool = field(default=False)\n",
    "    dataloader_pin_memory: bool = field(default=False)\n",
    "    eval_and_save_results: bool = field(default=True)\n",
    "    save_model: bool = field(default=False)\n",
    "    seed: int = field(default=42)\n",
    "    logging_first_step: bool = field(default=True)\n",
    "    early_stopping_patience: int = field(default = 5)  # number of evaluations without improvement to wait\n",
    "    early_stopping_threshold: float = field(default = 1e-3)  # threshold for an improvement\n",
    "training_args = TrainingArguments() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e05a33b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleDataset(Dataset):\n",
    "    def __init__(self, texts, labels, tokenizer, max_length=512):\n",
    "        self.texts = texts\n",
    "        self.labels = labels\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_length = max_length\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.texts)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        text = self.texts[idx]\n",
    "        label = self.labels[idx]\n",
    "        \n",
    "        encoding = self.tokenizer(text, return_tensors=\"pt\", max_length=self.max_length, padding=\"max_length\", truncation=True)\n",
    "        return {\n",
    "            \"input_ids\": encoding[\"input_ids\"].squeeze(0),\n",
    "            \"attention_mask\": encoding[\"attention_mask\"].squeeze(0),\n",
    "            \"labels\": torch.tensor(label, dtype=torch.long)\n",
    "        }\n",
    "\n",
    "class DNADataset(Dataset):\n",
    "    def __init__(self, dataframe, tokenizer, max_length=training_args.model_max_length):\n",
    "        self.texts = dataframe[\"sequence\"].tolist()  # CHANGED\n",
    "        self.labels = dataframe[\"label\"].tolist()  # CHANGED\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_length = max_length\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.texts)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        text = self.texts[idx]\n",
    "        label = self.labels[idx]\n",
    "        \n",
    "        encoding = self.tokenizer(text, return_tensors=\"pt\", max_length=self.max_length, padding=\"max_length\", truncation=True)\n",
    "        return {\n",
    "            \"input_ids\": encoding[\"input_ids\"].squeeze(0),\n",
    "            \"attention_mask\": encoding[\"attention_mask\"].squeeze(0),\n",
    "            \"labels\": torch.tensor(label, dtype=torch.long)\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "285883eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def custom_data_collator(data):\n",
    "    input_ids = torch.stack([item['input_ids'] for item in data])\n",
    "    attention_mask = torch.stack([item['attention_mask'] for item in data])\n",
    "    labels = torch.stack([item['labels'] for item in data])\n",
    "    return {\n",
    "        'input_ids': input_ids,  # CHANGED key name\n",
    "        'attention_mask': attention_mask,  # CHANGED key name\n",
    "        'labels': labels\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ef016e71",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class ModelArguments:\n",
    "    #model_name_or_path: Optional[str] = field(default=\"guardrail/llama-2-7b-guanaco-instruct-sharded\")\n",
    "    model_name_or_path: Optional[str] = field(default=\"togethercomputer/LLaMA-2-7B-32K\")\n",
    "    #model_name_or_path: Optional[str] = field(default=\"facebook/bart-base\")\n",
    "    use_lora: bool = field(default=True, metadata={\"help\": \"whether to use LoRA\"})\n",
    "    lora_r: int = field(default=8, metadata={\"help\": \"hidden dimension for LoRA\"})\n",
    "    lora_alpha: int = field(default=32, metadata={\"help\": \"alpha for LoRA\"})\n",
    "    lora_dropout: float = field(default=0.05, metadata={\"help\": \"dropout rate for LoRA\"})\n",
    "    #lora_target_modules: str = field(default=\"k_proj,q_proj,v_proj,fc1,fc2,output_proj\", metadata={\"help\": \"where to perform LoRA\"})\n",
    "    lora_target_modules: str = field(default=\"q_proj,v_proj\", metadata={\"help\": \"where to perform LoRA\"})\n",
    "    #lora_target_modules: str = field(default=\"query,key,value\", metadata={\"help\": \"where to perform LoRA\"})\n",
    "    use_4bit: bool = field(default=True, metadata={\"help\": \"whether to use 4-bit quantization\"})\n",
    "    use_nested_quant: bool = field(default=False, metadata={\"help\": \"Activate nested quantization for 4-bit base models\"})\n",
    "    bnb_4bit_compute_dtype: str = field(default=\"bfloat16\", metadata={\"help\": \"Compute dtype for 4-bit base models\"})\n",
    "    bnb_4bit_quant_type: str = field(default=\"nf4\", metadata={\"help\": \"Quantization type (fp4 or nf4)\"})\n",
    "model_args = ModelArguments()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "36715be7",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class DataArguments:\n",
    "    data_path: str = field(default=\"/common/zhangz2lab/zhanh/GUE/EMP/H3\", metadata={\"help\": \"Path to the training data.\"})\n",
    "data_args = DataArguments()        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93d0fa14",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "01a10002",
   "metadata": {},
   "outputs": [],
   "source": [
    "#model_name_or_path = \"guardrail/llama-2-7b-guanaco-instruct-sharded\"\n",
    "#model_name_or_path = \"facebook/bart-base\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path,\n",
    "        #cache_dir=training_args.cache_dir,\n",
    "        model_max_length=training_args.model_max_length,\n",
    "        padding_side=\"right\",\n",
    "        use_fast=True,\n",
    "        trust_remote_code=True)  # CHANGED (though this was originally okay)\n",
    "texts = [\"I love this!\", \"This is bad.\", \"Could be better.\", \"Excellent!\"]\n",
    "labels = [1, 0, 0, 1]\n",
    "\n",
    "train_dataset = SimpleDataset(texts, labels, tokenizer)\n",
    "test_texts = [\"This is amazing!\", \"I don't like it.\"]\n",
    "test_labels = [1, 0]\n",
    "test_dataset = SimpleDataset(test_texts, test_labels, tokenizer)\n",
    "\n",
    "train_df = pd.read_csv(os.path.join(data_args.data_path, \"train.csv\"))  # ADDED\n",
    "test_df = pd.read_csv(os.path.join(data_args.data_path, \"test.csv\"))  # ADDED\n",
    "val_df = pd.read_csv(os.path.join(data_args.data_path, \"dev.csv\"))  # ADDED\n",
    "\n",
    "train_dataset = DNADataset(train_df, tokenizer)  # CHANGED\n",
    "test_dataset = DNADataset(test_df, tokenizer)  # CHANGED\n",
    "val_dataset = DNADataset(val_df, tokenizer)  # CHANGED"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "504ea385",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ACAATAATAATAATAATAATAATAATAATAACAATAACAATAGTGGTAATAGTAGTAATAATAATAACAATAATAACAATAATAAAAATAATAATGACTTCGGCATTAAGATTGATAACAATTCACCGTCTTATGAAGGGTTTCCCCAGTTACAAATACCGCTTTCACAAGACAATTTGAACATAGAAGATAAAGAGGAGATGTCACCTAATATTGAAATTAAAAACGAACAAAATATGACTGACTCAAACGATATTCTTGGAGTATTCGATCAGTTAGATGCTCAGCTATTTGGGAAATACCTACCTTTAAATTACCCCTCTGAATGAAAACGTTATCTTTGATTTATATTCTATAATATCGTGGCTACAGCACTTGCTGAACATAAGCTTAAAACGTTTATGTGTGTATTTATATATGATAATAATAATAATAATAATAATAATAATAATAATAATAATAATAATAATAATAATAATGGTGATAAACTGCAATAACAA\n",
      "500\n",
      "[319, 5454, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 2477, 29909, 8254, 2477, 29909, 8254, 23799, 26788, 6040, 8254, 29954, 16881, 6040, 8254, 8254, 8254, 2477, 29909, 8254, 8254, 2477, 29909, 8254, 1299, 23184, 8254, 8254, 1299, 12739, 1783, 9472, 29954, 8766, 1299, 6040, 10051, 1299, 29911, 29954, 8254, 2477, 29909, 1299, 9472, 2477, 11135, 29911, 1783, 29911, 1299, 29954, 6344, 26788, 29954, 19988, 29911, 4174, 4174, 10051, 29911, 8687, 6344, 1299, 2477, 11135, 1783, 29911, 9472, 2477, 6344, 29954, 2477, 29909, 1299, 19988, 12739, 2477, 8254, 29954, 6344, 29954, 1299, 6344, 10051, 10051, 29954, 10051, 1299, 29954, 9472, 2477, 1783, 29909, 1299, 1299, 29911, 29954, 6344, 1299, 6040, 6344, 29909, 2477, 12739, 2477, 6344, 29909, 1299, 1299, 12739, 1783, 12739, 1783, 22701, 2477, 29954, 1299, 1299, 29911, 1783, 29911, 26788, 10051, 29911, 1299, 9472, 29954, 1299, 5454, 23799, 16881, 1299, 29954, 1783, 5454, 29954, 1783, 1299, 19988, 26788, 29954, 6344, 1299, 2477, 1783, 2477, 1783, 29911, 6040, 29909, 1299, 8687, 4174, 1783, 1783, 12739, 1299, 29954, 6344, 29909, 2477, 29954, 19988, 1299, 1783, 19988, 29954, 1299, 19988, 1299, 1299, 29911, 1783, 8254, 1299, 1299, 11135, 29911, 26788, 1783, 2477, 10051, 5454, 1783, 29911, 29954, 1783, 12739, 2477, 1299, 6344, 29954, 1783, 6040, 6344, 2477, 29954, 19988, 29911, 1299, 23799, 23799, 23799, 1299, 19988, 1299, 1299, 1299, 29954, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 8254, 1299, 29954, 23799, 29954, 1299, 6344, 17923, 29954, 5454, 8254, 2477, 6344]\n",
      "251\n"
     ]
    }
   ],
   "source": [
    "first_sequence = val_df.iloc[0][\"sequence\"]\n",
    "print(first_sequence)\n",
    "print(len(first_sequence))\n",
    "# Tokenize the sequence without padding\n",
    "tokens = tokenizer.encode(first_sequence, add_special_tokens=True, padding=False, return_tensors=\"pt\")\n",
    "\n",
    "# Print out the token IDs\n",
    "print(tokens[0].tolist())\n",
    "print(len(tokens[0].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0144eb2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.bfloat16\n"
     ]
    }
   ],
   "source": [
    "# Load tokenizer and model with QLoRA configuration\n",
    "compute_dtype = getattr(torch, model_args.bnb_4bit_compute_dtype)\n",
    "print(compute_dtype)\n",
    "#device_map = {\"\": 0}\n",
    "#device_map = {\"\": \"cuda:\" + str(int(os.environ.get(\"LOCAL_RANK\") or 0))}\n",
    "device_map = \"auto\"\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=model_args.use_4bit,\n",
    "    bnb_4bit_quant_type=model_args.bnb_4bit_quant_type,\n",
    "    bnb_4bit_compute_dtype=compute_dtype,\n",
    "    bnb_4bit_use_double_quant=model_args.use_nested_quant,\n",
    ")\n",
    "\n",
    "if compute_dtype == torch.float16 and model_args.use_4bit:\n",
    "    major, _ = torch.cuda.get_device_capability()\n",
    "    if major >= 8:\n",
    "        print(\"=\" * 80)\n",
    "        print(\"Your GPU supports bfloat16, you can accelerate training with the argument --bf16\")\n",
    "        print(\"=\" * 80)\n",
    "\n",
    "#n_gpus = torch.cuda.device_count()\n",
    "#max_memory = {i: '80GB' for i in range(n_gpus)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d2a99271",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4779108fd5db432681eb3957347e960c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of LlamaForSequenceClassification were not initialized from the model checkpoint at togethercomputer/LLaMA-2-7B-32K and are newly initialized: ['score.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 4,210,688 || all params: 3,373,551,616 || trainable%: 0.12481469025194841\n"
     ]
    }
   ],
   "source": [
    "model = AutoModelForSequenceClassification.from_pretrained(model_args.model_name_or_path, num_labels=2, device_map=device_map,\n",
    "    quantization_config=bnb_config, torch_dtype=torch.float32)  # CHANGED\n",
    "\n",
    "#model = AutoModelForSequenceClassification.from_pretrained(model_args.model_name_or_path, num_labels=2, device_map=device_map,\n",
    "#    quantization_config=bnb_config) \n",
    "#model = AutoModelForSequenceClassification.from_pretrained(model_args.model_name_or_path, num_labels=2, device_map='sequential',\n",
    "#                                                           max_memory=max_memory, quantization_config=bnb_config, torch_dtype=torch.float32)\n",
    "if model_args.use_lora:\n",
    "        lora_config = LoraConfig(\n",
    "            r = model_args.lora_r,\n",
    "            lora_alpha = model_args.lora_alpha,\n",
    "            target_modules = list(model_args.lora_target_modules.split(\",\")),\n",
    "            lora_dropout = model_args.lora_dropout,\n",
    "            bias=\"none\",\n",
    "            task_type=\"SEQ_CLS\",\n",
    "            inference_mode=False,\n",
    "            #peft_type=\"ADALORA\",\n",
    "        )\n",
    "        #print(list(model_args.lora_target_modules.split(\",\")))\n",
    "        model = get_peft_model(model, lora_config)\n",
    "        #model = AdaLoraModel(model, lora_config, \"default\")\n",
    "        model.print_trainable_parameters()\n",
    "\n",
    "# training_args = TrainingArguments(\n",
    "#     output_dir=\"./results\",\n",
    "#     num_train_epochs=3,\n",
    "#     per_device_train_batch_size=2,\n",
    "#     logging_dir=\"./logs\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "099efa64",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(pred):\n",
    "    logits = pred.predictions[0]\n",
    "    labels = pred.label_ids\n",
    "    preds = logits.argmax(-1)\n",
    "    accuracy = (preds == labels).mean()\n",
    "    return {\"accuracy\": accuracy}\n",
    "\n",
    "def compute_metrics(pred):\n",
    "    logits = pred.predictions[0]\n",
    "    labels = pred.label_ids\n",
    "    \n",
    "    print(\"Logits Shape:\", logits.shape)\n",
    "    print(\"Some Logits:\", logits[:5])  # Print the first 5 samples\n",
    "    print(\"Labels Shape:\", labels.shape)\n",
    "    print(\"Some Labels:\", labels[:5]) \n",
    "    \n",
    "    preds = logits.argmax(-1)\n",
    "    \n",
    "    probs = np.exp(logits) / np.sum(np.exp(logits), axis=-1, keepdims=True)\n",
    "    print(\"Some Probs:\", probs[:5]) \n",
    "    \n",
    "    # If it's binary classification, take the probability of the second class (usually the positive class)\n",
    "    # This assumes that the output shape is [batch_size, 2]\n",
    "    preds_prob = probs[:, 0]\n",
    "\n",
    "    # Apply thresholding to get binary predictions\n",
    "    threshold = 0.5\n",
    "    preds = (preds_prob > threshold).astype(int)\n",
    "\n",
    "    # Now, compute your metrics\n",
    "    accuracy = accuracy_score(labels, preds)\n",
    "    precision = precision_score(labels, preds)\n",
    "    recall = recall_score(labels, preds)\n",
    "    f1 = f1_score(labels, preds)\n",
    "    mcc = matthews_corrcoef(labels, preds)\n",
    "\n",
    "    return {\n",
    "        'accuracy': accuracy,\n",
    "        'precision': precision,\n",
    "        'recall': recall,\n",
    "        'f1': f1,\n",
    "        'mcc': mcc\n",
    "    }\n",
    "\n",
    "def calculate_metric_with_sklearn(pred, threshold=0.5):\n",
    "    logits = pred.predictions\n",
    "    labels = pred.label_ids\n",
    "    # Compute probabilities from logits\n",
    "    probabilities = softmax(logits, axis=-1)\n",
    "    #print(\"Logits:\", logits[:5]) # Printing the first 5 logits\n",
    "    #print(\"Labels:\", labels[:5]) # Printing the first 5 labels\n",
    "    #print(\"Probabilities:\", probabilities[:5]) # Printing the first 5 probabilities\n",
    "\n",
    "    # Assuming the second column is the positive class\n",
    "    valid_scores = probabilities[:, 1] \n",
    "\n",
    "    # Convert probabilities to binary predictions based on threshold\n",
    "    predictions = (valid_scores > threshold).astype(int)\n",
    "\n",
    "    #valid_mask = labels != -100  # Exclude padding tokens (assuming -100 is the padding token ID)\n",
    "    #valid_predictions = predictions[valid_mask]\n",
    "    #valid_labels = labels[valid_mask]\n",
    "    #valid_scores = valid_scores[valid_mask] \n",
    "\n",
    "    # Debugging prints\n",
    "    #print(\"Logits:\", logits[:5]) # Printing the first 5 logits\n",
    "    #print(\"Labels:\", valid_labels[:5]) # Printing the first 5 labels\n",
    "    #print(\"Probabilities:\", probabilities[:5]) # Printing the first 5 probabilities\n",
    "\n",
    "#     return {\n",
    "#         \"accuracy\": sklearn.metrics.accuracy_score(valid_labels, valid_predictions),\n",
    "#         \"f1\": sklearn.metrics.f1_score(valid_labels, valid_predictions, average=\"macro\", zero_division=0),\n",
    "#         \"matthews_correlation\": sklearn.metrics.matthews_corrcoef(valid_labels, valid_predictions),\n",
    "#         \"precision\": sklearn.metrics.precision_score(valid_labels, valid_predictions, average=\"macro\", zero_division=0),\n",
    "#         \"recall\": sklearn.metrics.recall_score(valid_labels, valid_predictions, average=\"macro\", zero_division=0),\n",
    "#         \"pr_auc\": sklearn.metrics.average_precision_score(valid_labels, valid_scores),\n",
    "#         \"roc_auc\": sklearn.metrics.roc_auc_score(valid_labels, valid_scores),\n",
    "#         \"brier_score\": sklearn.metrics.brier_score_loss(valid_labels, valid_scores)\n",
    "#     }\n",
    "    return {\n",
    "        \"accuracy\": sklearn.metrics.accuracy_score(labels, predictions),\n",
    "        \"f1\": sklearn.metrics.f1_score(labels, predictions, average=\"macro\", zero_division=0),\n",
    "        \"matthews_correlation\": sklearn.metrics.matthews_corrcoef(labels, predictions),\n",
    "        \"precision\": sklearn.metrics.precision_score(labels, predictions, average=\"macro\", zero_division=0),\n",
    "        \"recall\": sklearn.metrics.recall_score(labels, predictions, average=\"macro\", zero_division=0),\n",
    "        \"pr_auc\": sklearn.metrics.average_precision_score(labels, valid_scores),\n",
    "        \"roc_auc\": sklearn.metrics.roc_auc_score(labels, valid_scores),\n",
    "        \"brier_score\": sklearn.metrics.brier_score_loss(labels, valid_scores)\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "dbf21600",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=test_dataset,\n",
    "    compute_metrics=calculate_metric_with_sklearn,\n",
    "    data_collator=custom_data_collator\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1520ef96",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='7485' max='7485' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [7485/7485 2:06:58, Epoch 5/5]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Accuracy</th>\n",
       "      <th>F1</th>\n",
       "      <th>Matthews Correlation</th>\n",
       "      <th>Precision</th>\n",
       "      <th>Recall</th>\n",
       "      <th>Pr Auc</th>\n",
       "      <th>Roc Auc</th>\n",
       "      <th>Brier Score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>0.811600</td>\n",
       "      <td>0.750422</td>\n",
       "      <td>0.523046</td>\n",
       "      <td>0.522535</td>\n",
       "      <td>0.045675</td>\n",
       "      <td>0.522873</td>\n",
       "      <td>0.522801</td>\n",
       "      <td>0.506038</td>\n",
       "      <td>0.521217</td>\n",
       "      <td>0.273853</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>400</td>\n",
       "      <td>0.725300</td>\n",
       "      <td>0.716216</td>\n",
       "      <td>0.545758</td>\n",
       "      <td>0.522111</td>\n",
       "      <td>0.097496</td>\n",
       "      <td>0.554172</td>\n",
       "      <td>0.543867</td>\n",
       "      <td>0.526037</td>\n",
       "      <td>0.546631</td>\n",
       "      <td>0.259958</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>600</td>\n",
       "      <td>0.684900</td>\n",
       "      <td>0.612821</td>\n",
       "      <td>0.680695</td>\n",
       "      <td>0.679613</td>\n",
       "      <td>0.362572</td>\n",
       "      <td>0.682337</td>\n",
       "      <td>0.680241</td>\n",
       "      <td>0.746138</td>\n",
       "      <td>0.752861</td>\n",
       "      <td>0.211362</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>800</td>\n",
       "      <td>0.530900</td>\n",
       "      <td>0.481323</td>\n",
       "      <td>0.779559</td>\n",
       "      <td>0.779507</td>\n",
       "      <td>0.559917</td>\n",
       "      <td>0.780165</td>\n",
       "      <td>0.779751</td>\n",
       "      <td>0.831990</td>\n",
       "      <td>0.854471</td>\n",
       "      <td>0.156907</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1000</td>\n",
       "      <td>0.465300</td>\n",
       "      <td>0.465807</td>\n",
       "      <td>0.799599</td>\n",
       "      <td>0.799579</td>\n",
       "      <td>0.599161</td>\n",
       "      <td>0.799589</td>\n",
       "      <td>0.799572</td>\n",
       "      <td>0.835012</td>\n",
       "      <td>0.862805</td>\n",
       "      <td>0.149234</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1200</td>\n",
       "      <td>0.461000</td>\n",
       "      <td>0.473990</td>\n",
       "      <td>0.800935</td>\n",
       "      <td>0.800928</td>\n",
       "      <td>0.601867</td>\n",
       "      <td>0.800924</td>\n",
       "      <td>0.800943</td>\n",
       "      <td>0.843604</td>\n",
       "      <td>0.868483</td>\n",
       "      <td>0.148309</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1400</td>\n",
       "      <td>0.461400</td>\n",
       "      <td>0.505203</td>\n",
       "      <td>0.788911</td>\n",
       "      <td>0.785709</td>\n",
       "      <td>0.592505</td>\n",
       "      <td>0.804837</td>\n",
       "      <td>0.787909</td>\n",
       "      <td>0.853396</td>\n",
       "      <td>0.868269</td>\n",
       "      <td>0.159500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1600</td>\n",
       "      <td>0.433800</td>\n",
       "      <td>0.474986</td>\n",
       "      <td>0.804275</td>\n",
       "      <td>0.803715</td>\n",
       "      <td>0.610672</td>\n",
       "      <td>0.806807</td>\n",
       "      <td>0.803872</td>\n",
       "      <td>0.856440</td>\n",
       "      <td>0.873165</td>\n",
       "      <td>0.146410</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1800</td>\n",
       "      <td>0.417600</td>\n",
       "      <td>0.433525</td>\n",
       "      <td>0.813627</td>\n",
       "      <td>0.813347</td>\n",
       "      <td>0.628169</td>\n",
       "      <td>0.814819</td>\n",
       "      <td>0.813352</td>\n",
       "      <td>0.862431</td>\n",
       "      <td>0.879690</td>\n",
       "      <td>0.137622</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2000</td>\n",
       "      <td>0.420200</td>\n",
       "      <td>0.441847</td>\n",
       "      <td>0.799599</td>\n",
       "      <td>0.799490</td>\n",
       "      <td>0.600626</td>\n",
       "      <td>0.800765</td>\n",
       "      <td>0.799863</td>\n",
       "      <td>0.863263</td>\n",
       "      <td>0.883756</td>\n",
       "      <td>0.140104</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2200</td>\n",
       "      <td>0.417300</td>\n",
       "      <td>0.423074</td>\n",
       "      <td>0.820975</td>\n",
       "      <td>0.820799</td>\n",
       "      <td>0.642446</td>\n",
       "      <td>0.821682</td>\n",
       "      <td>0.820765</td>\n",
       "      <td>0.870474</td>\n",
       "      <td>0.890345</td>\n",
       "      <td>0.132237</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2400</td>\n",
       "      <td>0.435500</td>\n",
       "      <td>0.416375</td>\n",
       "      <td>0.828991</td>\n",
       "      <td>0.828777</td>\n",
       "      <td>0.658740</td>\n",
       "      <td>0.829995</td>\n",
       "      <td>0.828746</td>\n",
       "      <td>0.872130</td>\n",
       "      <td>0.891650</td>\n",
       "      <td>0.130419</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2600</td>\n",
       "      <td>0.407600</td>\n",
       "      <td>0.442869</td>\n",
       "      <td>0.825651</td>\n",
       "      <td>0.825269</td>\n",
       "      <td>0.652944</td>\n",
       "      <td>0.827641</td>\n",
       "      <td>0.825307</td>\n",
       "      <td>0.872809</td>\n",
       "      <td>0.893576</td>\n",
       "      <td>0.134192</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2800</td>\n",
       "      <td>0.417000</td>\n",
       "      <td>0.407801</td>\n",
       "      <td>0.830995</td>\n",
       "      <td>0.830775</td>\n",
       "      <td>0.662800</td>\n",
       "      <td>0.832057</td>\n",
       "      <td>0.830745</td>\n",
       "      <td>0.867881</td>\n",
       "      <td>0.895084</td>\n",
       "      <td>0.125504</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3000</td>\n",
       "      <td>0.421000</td>\n",
       "      <td>0.431013</td>\n",
       "      <td>0.825651</td>\n",
       "      <td>0.825620</td>\n",
       "      <td>0.651276</td>\n",
       "      <td>0.825679</td>\n",
       "      <td>0.825598</td>\n",
       "      <td>0.874301</td>\n",
       "      <td>0.893181</td>\n",
       "      <td>0.131636</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3200</td>\n",
       "      <td>0.385700</td>\n",
       "      <td>0.405181</td>\n",
       "      <td>0.839679</td>\n",
       "      <td>0.839493</td>\n",
       "      <td>0.680071</td>\n",
       "      <td>0.840625</td>\n",
       "      <td>0.839447</td>\n",
       "      <td>0.876688</td>\n",
       "      <td>0.900104</td>\n",
       "      <td>0.123780</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3400</td>\n",
       "      <td>0.357900</td>\n",
       "      <td>0.471537</td>\n",
       "      <td>0.825651</td>\n",
       "      <td>0.825352</td>\n",
       "      <td>0.654870</td>\n",
       "      <td>0.828802</td>\n",
       "      <td>0.826073</td>\n",
       "      <td>0.869229</td>\n",
       "      <td>0.897264</td>\n",
       "      <td>0.136142</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3600</td>\n",
       "      <td>0.411900</td>\n",
       "      <td>0.439947</td>\n",
       "      <td>0.826319</td>\n",
       "      <td>0.824895</td>\n",
       "      <td>0.660986</td>\n",
       "      <td>0.835461</td>\n",
       "      <td>0.825598</td>\n",
       "      <td>0.876207</td>\n",
       "      <td>0.899581</td>\n",
       "      <td>0.133783</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3800</td>\n",
       "      <td>0.406900</td>\n",
       "      <td>0.423484</td>\n",
       "      <td>0.830327</td>\n",
       "      <td>0.829212</td>\n",
       "      <td>0.667204</td>\n",
       "      <td>0.837562</td>\n",
       "      <td>0.829688</td>\n",
       "      <td>0.877427</td>\n",
       "      <td>0.900862</td>\n",
       "      <td>0.128360</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4000</td>\n",
       "      <td>0.397100</td>\n",
       "      <td>0.439754</td>\n",
       "      <td>0.824983</td>\n",
       "      <td>0.824469</td>\n",
       "      <td>0.655501</td>\n",
       "      <td>0.830001</td>\n",
       "      <td>0.825515</td>\n",
       "      <td>0.880119</td>\n",
       "      <td>0.903533</td>\n",
       "      <td>0.132632</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4200</td>\n",
       "      <td>0.405400</td>\n",
       "      <td>0.408283</td>\n",
       "      <td>0.839679</td>\n",
       "      <td>0.839318</td>\n",
       "      <td>0.681170</td>\n",
       "      <td>0.841844</td>\n",
       "      <td>0.839331</td>\n",
       "      <td>0.881051</td>\n",
       "      <td>0.902106</td>\n",
       "      <td>0.124081</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4400</td>\n",
       "      <td>0.381300</td>\n",
       "      <td>0.543871</td>\n",
       "      <td>0.795591</td>\n",
       "      <td>0.794461</td>\n",
       "      <td>0.600007</td>\n",
       "      <td>0.803758</td>\n",
       "      <td>0.796295</td>\n",
       "      <td>0.875728</td>\n",
       "      <td>0.891139</td>\n",
       "      <td>0.157644</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4600</td>\n",
       "      <td>0.370500</td>\n",
       "      <td>0.407038</td>\n",
       "      <td>0.842351</td>\n",
       "      <td>0.842209</td>\n",
       "      <td>0.685189</td>\n",
       "      <td>0.843036</td>\n",
       "      <td>0.842154</td>\n",
       "      <td>0.885686</td>\n",
       "      <td>0.905623</td>\n",
       "      <td>0.121952</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4800</td>\n",
       "      <td>0.366600</td>\n",
       "      <td>0.439576</td>\n",
       "      <td>0.838343</td>\n",
       "      <td>0.837958</td>\n",
       "      <td>0.678624</td>\n",
       "      <td>0.840646</td>\n",
       "      <td>0.837983</td>\n",
       "      <td>0.886729</td>\n",
       "      <td>0.904691</td>\n",
       "      <td>0.128268</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5000</td>\n",
       "      <td>0.374800</td>\n",
       "      <td>0.412229</td>\n",
       "      <td>0.839679</td>\n",
       "      <td>0.839276</td>\n",
       "      <td>0.681449</td>\n",
       "      <td>0.842148</td>\n",
       "      <td>0.839308</td>\n",
       "      <td>0.886582</td>\n",
       "      <td>0.906631</td>\n",
       "      <td>0.122272</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5200</td>\n",
       "      <td>0.382500</td>\n",
       "      <td>0.403850</td>\n",
       "      <td>0.841683</td>\n",
       "      <td>0.841682</td>\n",
       "      <td>0.683582</td>\n",
       "      <td>0.841812</td>\n",
       "      <td>0.841770</td>\n",
       "      <td>0.879543</td>\n",
       "      <td>0.902785</td>\n",
       "      <td>0.121574</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5400</td>\n",
       "      <td>0.360500</td>\n",
       "      <td>0.396596</td>\n",
       "      <td>0.847695</td>\n",
       "      <td>0.847695</td>\n",
       "      <td>0.695583</td>\n",
       "      <td>0.847807</td>\n",
       "      <td>0.847777</td>\n",
       "      <td>0.890194</td>\n",
       "      <td>0.908638</td>\n",
       "      <td>0.118197</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5600</td>\n",
       "      <td>0.371500</td>\n",
       "      <td>0.398870</td>\n",
       "      <td>0.838343</td>\n",
       "      <td>0.837845</td>\n",
       "      <td>0.679382</td>\n",
       "      <td>0.841466</td>\n",
       "      <td>0.837925</td>\n",
       "      <td>0.887848</td>\n",
       "      <td>0.906833</td>\n",
       "      <td>0.120744</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5800</td>\n",
       "      <td>0.358100</td>\n",
       "      <td>0.401455</td>\n",
       "      <td>0.843019</td>\n",
       "      <td>0.843009</td>\n",
       "      <td>0.686019</td>\n",
       "      <td>0.843006</td>\n",
       "      <td>0.843013</td>\n",
       "      <td>0.890046</td>\n",
       "      <td>0.907624</td>\n",
       "      <td>0.119256</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6000</td>\n",
       "      <td>0.404000</td>\n",
       "      <td>0.400857</td>\n",
       "      <td>0.845023</td>\n",
       "      <td>0.844633</td>\n",
       "      <td>0.692183</td>\n",
       "      <td>0.847537</td>\n",
       "      <td>0.844652</td>\n",
       "      <td>0.892913</td>\n",
       "      <td>0.908568</td>\n",
       "      <td>0.120159</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6200</td>\n",
       "      <td>0.342600</td>\n",
       "      <td>0.395788</td>\n",
       "      <td>0.845023</td>\n",
       "      <td>0.844654</td>\n",
       "      <td>0.692038</td>\n",
       "      <td>0.847380</td>\n",
       "      <td>0.844664</td>\n",
       "      <td>0.894780</td>\n",
       "      <td>0.910396</td>\n",
       "      <td>0.118424</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6400</td>\n",
       "      <td>0.346400</td>\n",
       "      <td>0.424834</td>\n",
       "      <td>0.842351</td>\n",
       "      <td>0.842138</td>\n",
       "      <td>0.685607</td>\n",
       "      <td>0.843512</td>\n",
       "      <td>0.842096</td>\n",
       "      <td>0.891758</td>\n",
       "      <td>0.905116</td>\n",
       "      <td>0.125389</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6600</td>\n",
       "      <td>0.338400</td>\n",
       "      <td>0.406614</td>\n",
       "      <td>0.845023</td>\n",
       "      <td>0.844939</td>\n",
       "      <td>0.690243</td>\n",
       "      <td>0.845359</td>\n",
       "      <td>0.844884</td>\n",
       "      <td>0.896108</td>\n",
       "      <td>0.910716</td>\n",
       "      <td>0.118662</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6800</td>\n",
       "      <td>0.346600</td>\n",
       "      <td>0.403133</td>\n",
       "      <td>0.848363</td>\n",
       "      <td>0.848363</td>\n",
       "      <td>0.696818</td>\n",
       "      <td>0.848402</td>\n",
       "      <td>0.848416</td>\n",
       "      <td>0.895024</td>\n",
       "      <td>0.911130</td>\n",
       "      <td>0.117892</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7000</td>\n",
       "      <td>0.368200</td>\n",
       "      <td>0.400495</td>\n",
       "      <td>0.849031</td>\n",
       "      <td>0.848929</td>\n",
       "      <td>0.698371</td>\n",
       "      <td>0.849502</td>\n",
       "      <td>0.848869</td>\n",
       "      <td>0.895913</td>\n",
       "      <td>0.911256</td>\n",
       "      <td>0.117877</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7200</td>\n",
       "      <td>0.358200</td>\n",
       "      <td>0.408240</td>\n",
       "      <td>0.847695</td>\n",
       "      <td>0.847695</td>\n",
       "      <td>0.695583</td>\n",
       "      <td>0.847807</td>\n",
       "      <td>0.847777</td>\n",
       "      <td>0.896552</td>\n",
       "      <td>0.910557</td>\n",
       "      <td>0.119646</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7400</td>\n",
       "      <td>0.343600</td>\n",
       "      <td>0.401088</td>\n",
       "      <td>0.845691</td>\n",
       "      <td>0.845475</td>\n",
       "      <td>0.692350</td>\n",
       "      <td>0.846921</td>\n",
       "      <td>0.845430</td>\n",
       "      <td>0.896815</td>\n",
       "      <td>0.911380</td>\n",
       "      <td>0.118422</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='375' max='375' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [375/375 01:09]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.39563965797424316, 'eval_accuracy': 0.8463593854375417, 'eval_f1': 0.8460691778911151, 'eval_matthews_correlation': 0.6941895991914795, 'eval_precision': 0.8481466515812726, 'eval_recall': 0.8460461255600578, 'eval_pr_auc': 0.8947050788330689, 'eval_roc_auc': 0.9103943164170579, 'eval_brier_score': 0.11823498775671602, 'eval_runtime': 70.0168, 'eval_samples_per_second': 21.381, 'eval_steps_per_second': 5.356, 'epoch': 5.0}\n"
     ]
    }
   ],
   "source": [
    "trainer.train()\n",
    "results = trainer.evaluate()\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d33618f4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4b5ab55",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dna_llama",
   "language": "python",
   "name": "dna_llama"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}