esm2_vanilla.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Errno 2] No such file or directory: 'code'\n",
      "/mnt/aiongpfs/users/egomez/Projects/GSoC24/ESMbind/code\n"
     ]
    }
   ],
   "source": [
    "cd code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# I / O, helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "#import wandb\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import pickle\n",
    "import xml.etree.ElementTree as ET\n",
    "from datetime import datetime\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.utils.class_weight import compute_class_weight\n",
    "from sklearn.metrics import (\n",
    "    accuracy_score,\n",
    "    precision_recall_fscore_support,\n",
    "    roc_auc_score,\n",
    "    matthews_corrcoef\n",
    ")\n",
    "from transformers import (\n",
    "    AutoModelForTokenClassification,\n",
    "    AutoTokenizer,\n",
    "    DataCollatorForTokenClassification,\n",
    "    TrainingArguments,\n",
    "    Trainer\n",
    ")\n",
    "from datasets import Dataset\n",
    "from accelerate import Accelerator\n",
    "# Imports specific to the custom peft lora model\n",
    "from peft import get_peft_config, PeftModel, PeftConfig, get_peft_model, LoraConfig, TaskType\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Helper Functions and Data Preparation\n",
    "def truncate_labels(labels, max_length):\n",
    "    \"\"\"Truncate labels to the specified max_length.\"\"\"\n",
    "    return [label[:max_length] for label in labels]\n",
    "\n",
    "def compute_metrics_train(p):\n",
    "    \"\"\"Compute metrics for evaluation.\"\"\"\n",
    "    predictions, labels = p\n",
    "    predictions = np.argmax(predictions, axis=2)\n",
    "    \n",
    "    # Remove padding (-100 labels)\n",
    "    predictions = predictions[labels != -100].flatten()\n",
    "    labels = labels[labels != -100].flatten()\n",
    "    \n",
    "    # Compute accuracy\n",
    "    accuracy = accuracy_score(labels, predictions)\n",
    "    \n",
    "    # Compute precision, recall, F1 score, and AUC\n",
    "    precision, recall, f1, _ = precision_recall_fscore_support(labels, predictions, average='binary')\n",
    "    auc = roc_auc_score(labels, predictions)\n",
    "    \n",
    "    # Compute MCC\n",
    "    mcc = matthews_corrcoef(labels, predictions) \n",
    "    \n",
    "    return {'accuracy': accuracy, 'precision': precision, 'recall': recall, 'f1': f1, 'auc': auc, 'mcc': mcc} \n",
    "\n",
    "def compute_loss(model, inputs):\n",
    "    \"\"\"Custom compute_loss function.\"\"\"\n",
    "    logits = model(**inputs).logits\n",
    "    labels = inputs[\"labels\"]\n",
    "    loss_fct = nn.CrossEntropyLoss(weight=class_weights)\n",
    "    active_loss = inputs[\"attention_mask\"].view(-1) == 1\n",
    "    active_logits = logits.view(-1, model.config.num_labels)\n",
    "    active_labels = torch.where(\n",
    "        active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)\n",
    "    )\n",
    "    loss = loss_fct(active_logits, active_labels)\n",
    "    return loss\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data load & embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Load the data from pickle files (replace with your local paths)\n",
    "with open(\"../data/train_sequences_chunked_by_family.pkl\", \"rb\") as f:\n",
    "    train_sequences = pickle.load(f)\n",
    "\n",
    "with open(\"../data/test_sequences_chunked_by_family.pkl\", \"rb\") as f:\n",
    "    test_sequences = pickle.load(f)\n",
    "\n",
    "with open(\"../data/train_labels_chunked_by_family.pkl\", \"rb\") as f:\n",
    "    train_labels = pickle.load(f)\n",
    "\n",
    "with open(\"../data/test_labels_chunked_by_family.pkl\", \"rb\") as f:\n",
    "    test_labels = pickle.load(f)\n",
    "\n",
    "# make dataset smaller:\n",
    "train_sequences = train_sequences[0:300]\n",
    "test_sequences = test_sequences[0:100]\n",
    "train_labels = train_labels[0:300]\n",
    "test_labels = test_labels[0:100]\n",
    "\n",
    "# Tokenization\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"facebook/esm2_t6_8M_UR50D\")  # esm2_t12_35M_UR50D\n",
    "max_sequence_length = 500 # 1000\n",
    "\n",
    "train_tokenized = tokenizer(train_sequences, padding=True, truncation=True, max_length=max_sequence_length, return_tensors=\"pt\", is_split_into_words=False)\n",
    "test_tokenized = tokenizer(test_sequences, padding=True, truncation=True, max_length=max_sequence_length, return_tensors=\"pt\", is_split_into_words=False)\n",
    "\n",
    "# Directly truncate the entire list of labels\n",
    "train_labels = truncate_labels(train_labels, max_sequence_length)\n",
    "test_labels = truncate_labels(test_labels, max_sequence_length)\n",
    "\n",
    "train_dataset = Dataset.from_dict({k: v for k, v in train_tokenized.items()}).add_column(\"labels\", train_labels)\n",
    "test_dataset = Dataset.from_dict({k: v for k, v in test_tokenized.items()}).add_column(\"labels\", test_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of EsmForTokenClassification were not initialized from the model checkpoint at facebook/esm2_t6_8M_UR50D and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModelForSequenceClassification\n",
    "\n",
    "model = AutoModelForTokenClassification.from_pretrained(\"facebook/esm2_t6_8M_UR50D\", num_labels=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install evaluate\n",
    "import evaluate\n",
    "metric = evaluate.load(\"accuracy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics_train(eval_pred):\n",
    "    \"\"\"Compute metrics for evaluation.\"\"\"\n",
    "    predictions, labels = eval_pred\n",
    "    predictions = np.argmax(predictions, axis=2)\n",
    "    \n",
    "    # Remove padding (-100 labels)\n",
    "    predictions = predictions[labels != -100].flatten()\n",
    "    labels = labels[labels != -100].flatten()\n",
    "    \n",
    "    # Compute accuracy\n",
    "    accuracy = accuracy_score(labels, predictions)\n",
    "    \n",
    "    # Compute precision, recall, F1 score, and AUC\n",
    "    precision, recall, f1, _ = precision_recall_fscore_support(labels, predictions, average='binary')\n",
    "    auc = roc_auc_score(labels, predictions)\n",
    "    \n",
    "    # Compute MCC\n",
    "    mcc = matthews_corrcoef(labels, predictions) \n",
    "    \n",
    "    return {'accuracy': accuracy, 'precision': precision, 'recall': recall, 'f1': f1, 'auc': auc, 'mcc': mcc} "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "timestamp = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=f\"trained_models/esm2_t6_8M-binding-sites_{timestamp}\", \n",
    "    eval_strategy=\"epoch\",\n",
    "    seed=8893,\n",
    "    eval_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=\"f1\",\n",
    "    greater_is_better=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"facebook/esm2_t6_8M_UR50D\")  # esm2_t12_35M_UR50D\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Detected kernel version 4.18.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
     ]
    }
   ],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=test_dataset,\n",
    "    compute_metrics=compute_metrics_train,\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=DataCollatorForTokenClassification(tokenizer=tokenizer)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='114' max='114' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [114/114 21:54, 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>Accuracy</th>\n",
       "      <th>Precision</th>\n",
       "      <th>Recall</th>\n",
       "      <th>F1</th>\n",
       "      <th>Auc</th>\n",
       "      <th>Mcc</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.152388</td>\n",
       "      <td>0.966500</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.500000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.157112</td>\n",
       "      <td>0.966500</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.500000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.156422</td>\n",
       "      <td>0.966500</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.500000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/users/egomez/Projects/GSoC24/ESMbind/code/.venv/esmbind311/lib64/python3.11/site-packages/sklearn/metrics/_classification.py:1517: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 due to no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, f\"{metric.capitalize()} is\", len(result))\n",
      "/home/users/egomez/Projects/GSoC24/ESMbind/code/.venv/esmbind311/lib64/python3.11/site-packages/sklearn/metrics/_classification.py:1517: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 due to no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, f\"{metric.capitalize()} is\", len(result))\n",
      "/home/users/egomez/Projects/GSoC24/ESMbind/code/.venv/esmbind311/lib64/python3.11/site-packages/sklearn/metrics/_classification.py:1517: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 due to no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, f\"{metric.capitalize()} is\", len(result))\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=114, training_loss=0.03993859207421018, metrics={'train_runtime': 1325.1209, 'train_samples_per_second': 0.679, 'train_steps_per_second': 0.086, 'total_flos': 19977740100000.0, 'train_loss': 0.03993859207421018, 'epoch': 3.0})"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('../results/binding_sites/best_model_esm2_t6_8M_2024-05-27_17-22-09/tokenizer_config.json',\n",
       " '../results/binding_sites/best_model_esm2_t6_8M_2024-05-27_17-22-09/special_tokens_map.json',\n",
       " '../results/binding_sites/best_model_esm2_t6_8M_2024-05-27_17-22-09/vocab.txt',\n",
       " '../results/binding_sites/best_model_esm2_t6_8M_2024-05-27_17-22-09/added_tokens.json')"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Save Model\n",
    "save_path = os.path.join(\"../results/binding_sites\", f\"best_model_esm2_t6_8M_{timestamp}\")\n",
    "trainer.save_model(save_path)\n",
    "tokenizer.save_pretrained(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Detected kernel version 4.18.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
     ]
    }
   ],
   "source": [
    "# Define paths to the tuned and base models\n",
    "base_model_path = \"facebook/esm2_t6_8M_UR50D\"  \n",
    "tuned_model_path = \"../results/binding_sites/best_model_esm2_t6_8M_2024-05-27_17-22-09\" \n",
    "\n",
    "# Load the model\n",
    "tuned_model = AutoModelForTokenClassification.from_pretrained(tuned_model_path)\n",
    "accelerator = Accelerator()\n",
    "model = accelerator.prepare(tuned_model)  # Prepare the model using the accelerator\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Define label mappings\n",
    "id2label = {0: \"No binding site\", 1: \"Binding site\"}\n",
    "label2id = {v: k for k, v in id2label.items()}\n",
    "\n",
    "# Create a data collator\n",
    "data_collator = DataCollatorForTokenClassification(tokenizer)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function to compute the metrics\n",
    "def compute_metrics(dataset):\n",
    "    # Get the predictions using the trained model\n",
    "    trainer = Trainer(model=model, data_collator=data_collator)\n",
    "    predictions, labels, _ = trainer.predict(test_dataset=dataset)\n",
    "    \n",
    "    # Remove padding and special tokens\n",
    "    mask = labels != -100\n",
    "    true_labels = labels[mask].flatten()\n",
    "    flat_predictions = np.argmax(predictions, axis=2)[mask].flatten().tolist()\n",
    "\n",
    "    # Compute the metrics\n",
    "    accuracy = accuracy_score(true_labels, flat_predictions)\n",
    "    precision, recall, f1, _ = precision_recall_fscore_support(true_labels, flat_predictions, average='binary')\n",
    "    auc = roc_auc_score(true_labels, flat_predictions)\n",
    "    mcc = matthews_corrcoef(true_labels, flat_predictions)  # Compute the MCC\n",
    "    \n",
    "    return {\"accuracy\": accuracy, \"precision\": precision, \"recall\": recall, \"f1\": f1, \"auc\": auc, \"mcc\": mcc}  # Include the MCC in the returned dictionary\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Detected kernel version 4.18.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
     ]
    },
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/users/egomez/Projects/GSoC24/ESMbind/code/.venv/esmbind311/lib64/python3.11/site-packages/sklearn/metrics/_classification.py:1517: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 due to no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, f\"{metric.capitalize()} is\", len(result))\n",
      "Detected kernel version 4.18.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
     ]
    },
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/users/egomez/Projects/GSoC24/ESMbind/code/.venv/esmbind311/lib64/python3.11/site-packages/sklearn/metrics/_classification.py:1517: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 due to no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, f\"{metric.capitalize()} is\", len(result))\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "({'accuracy': 0.9938003552301905,\n",
       "  'precision': 0.0,\n",
       "  'recall': 0.0,\n",
       "  'f1': 0.0,\n",
       "  'auc': 0.5,\n",
       "  'mcc': 0.0},\n",
       " {'accuracy': 0.966500023361211,\n",
       "  'precision': 0.0,\n",
       "  'recall': 0.0,\n",
       "  'f1': 0.0,\n",
       "  'auc': 0.5,\n",
       "  'mcc': 0.0})"
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get the metrics for the training and test datasets\n",
    "train_metrics = compute_metrics(train_dataset)\n",
    "test_metrics = compute_metrics(test_dataset)\n",
    "\n",
    "train_metrics, test_metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EsmForTokenClassification(\n",
       "  (esm): EsmModel(\n",
       "    (embeddings): EsmEmbeddings(\n",
       "      (word_embeddings): Embedding(33, 320, padding_idx=1)\n",
       "      (dropout): Dropout(p=0.0, inplace=False)\n",
       "      (position_embeddings): Embedding(1026, 320, padding_idx=1)\n",
       "    )\n",
       "    (encoder): EsmEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0-5): 6 x EsmLayer(\n",
       "          (attention): EsmAttention(\n",
       "            (self): EsmSelfAttention(\n",
       "              (query): Linear(in_features=320, out_features=320, bias=True)\n",
       "              (key): Linear(in_features=320, out_features=320, bias=True)\n",
       "              (value): Linear(in_features=320, out_features=320, bias=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "              (rotary_embeddings): RotaryEmbedding()\n",
       "            )\n",
       "            (output): EsmSelfOutput(\n",
       "              (dense): Linear(in_features=320, out_features=320, bias=True)\n",
       "              (dropout): Dropout(p=0.0, inplace=False)\n",
       "            )\n",
       "            (LayerNorm): LayerNorm((320,), eps=1e-05, elementwise_affine=True)\n",
       "          )\n",
       "          (intermediate): EsmIntermediate(\n",
       "            (dense): Linear(in_features=320, out_features=1280, bias=True)\n",
       "          )\n",
       "          (output): EsmOutput(\n",
       "            (dense): Linear(in_features=1280, out_features=320, bias=True)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "          (LayerNorm): LayerNorm((320,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "      )\n",
       "      (emb_layer_norm_after): LayerNorm((320,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (contact_head): EsmContactPredictionHead(\n",
       "      (regression): Linear(in_features=120, out_features=1, bias=True)\n",
       "      (activation): Sigmoid()\n",
       "    )\n",
       "  )\n",
       "  (dropout): Dropout(p=0.0, inplace=False)\n",
       "  (classifier): Linear(in_features=320, out_features=2, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Set paths and model if not loaded before\n",
    "# Path to the saved tuned model and ESM2 base model\n",
    "base_model_path = \"facebook/esm2_t6_8M_UR50D\"  \n",
    "tuned_model_path = \"../results/binding_sites/best_model_esm2_t6_8M_2024-05-27_17-22-09\" \n",
    "\n",
    "# Load the model\n",
    "tuned_model = AutoModelForTokenClassification.from_pretrained(tuned_model_path)\n",
    "\n",
    "# Ensure the model is in evaluation mode\n",
    "tuned_model.eval()\n",
    "\n",
    "# Load the tokenizer if not loaded already\n",
    "# tokenizer = AutoTokenizer.from_pretrained(base_model_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('M', 'No binding site')\n",
      "('A', 'No binding site')\n",
      "('V', 'No binding site')\n",
      "('P', 'No binding site')\n",
      "('E', 'No binding site')\n",
      "('T', 'No binding site')\n",
      "('R', 'No binding site')\n",
      "('P', 'No binding site')\n",
      "('N', 'No binding site')\n",
      "('H', 'No binding site')\n",
      "('T', 'No binding site')\n",
      "('I', 'No binding site')\n",
      "('Y', 'No binding site')\n",
      "('I', 'No binding site')\n",
      "('N', 'No binding site')\n",
      "('N', 'No binding site')\n",
      "('L', 'No binding site')\n",
      "('N', 'No binding site')\n",
      "('E', 'No binding site')\n",
      "('K', 'No binding site')\n",
      "('I', 'No binding site')\n",
      "('K', 'No binding site')\n",
      "('K', 'No binding site')\n",
      "('D', 'No binding site')\n",
      "('E', 'No binding site')\n",
      "('L', 'No binding site')\n",
      "('K', 'No binding site')\n",
      "('K', 'No binding site')\n",
      "('S', 'No binding site')\n",
      "('L', 'No binding site')\n",
      "('H', 'No binding site')\n",
      "('A', 'No binding site')\n",
      "('I', 'No binding site')\n",
      "('F', 'No binding site')\n",
      "('S', 'No binding site')\n",
      "('R', 'No binding site')\n",
      "('F', 'No binding site')\n",
      "('G', 'No binding site')\n",
      "('Q', 'No binding site')\n",
      "('I', 'No binding site')\n",
      "('L', 'No binding site')\n",
      "('D', 'No binding site')\n",
      "('I', 'No binding site')\n",
      "('L', 'No binding site')\n",
      "('V', 'No binding site')\n",
      "('S', 'No binding site')\n",
      "('R', 'No binding site')\n",
      "('S', 'No binding site')\n",
      "('L', 'No binding site')\n",
      "('K', 'No binding site')\n",
      "('M', 'No binding site')\n",
      "('R', 'No binding site')\n",
      "('G', 'No binding site')\n",
      "('Q', 'No binding site')\n",
      "('A', 'No binding site')\n",
      "('F', 'No binding site')\n",
      "('V', 'No binding site')\n",
      "('I', 'No binding site')\n",
      "('F', 'No binding site')\n",
      "('K', 'No binding site')\n",
      "('E', 'No binding site')\n",
      "('V', 'No binding site')\n",
      "('S', 'No binding site')\n",
      "('S', 'No binding site')\n",
      "('A', 'No binding site')\n",
      "('T', 'No binding site')\n",
      "('N', 'No binding site')\n",
      "('A', 'No binding site')\n",
      "('L', 'No binding site')\n",
      "('R', 'No binding site')\n",
      "('S', 'No binding site')\n",
      "('M', 'No binding site')\n",
      "('Q', 'No binding site')\n",
      "('G', 'No binding site')\n",
      "('F', 'No binding site')\n",
      "('P', 'No binding site')\n",
      "('F', 'No binding site')\n",
      "('Y', 'No binding site')\n",
      "('D', 'No binding site')\n",
      "('K', 'No binding site')\n",
      "('P', 'No binding site')\n",
      "('M', 'No binding site')\n",
      "('R', 'No binding site')\n",
      "('I', 'No binding site')\n",
      "('Q', 'No binding site')\n",
      "('Y', 'No binding site')\n",
      "('A', 'No binding site')\n",
      "('K', 'No binding site')\n",
      "('T', 'No binding site')\n",
      "('D', 'No binding site')\n",
      "('S', 'No binding site')\n",
      "('D', 'No binding site')\n",
      "('I', 'No binding site')\n",
      "('I', 'No binding site')\n",
      "('A', 'No binding site')\n",
      "('K', 'No binding site')\n",
      "('M', 'No binding site')\n",
      "('K', 'No binding site')\n",
      "('G', 'No binding site')\n",
      "('T', 'No binding site')\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# Protein sequence for inference\n",
    "protein_sequence = \"MAVPETRPNHTIYINNLNEKIKKDELKKSLHAIFSRFGQILDILVSRSLKMRGQAFVIFKEVSSATNALRSMQGFPFYDKPMRIQYAKTDSDIIAKMKGT\"  # Replace with your actual sequence\n",
    "\n",
    "# Tokenize the sequence\n",
    "inputs = tokenizer(protein_sequence, return_tensors=\"pt\", truncation=True, max_length=1024, padding='max_length')\n",
    "\n",
    "# Run the model\n",
    "with torch.no_grad():\n",
    "    logits = tuned_model(**inputs).logits\n",
    "\n",
    "# Get predictions\n",
    "tokens = tokenizer.convert_ids_to_tokens(inputs[\"input_ids\"][0])  # Convert input ids back to tokens\n",
    "predictions = torch.argmax(logits, dim=2)\n",
    "\n",
    "# Define labels\n",
    "id2label = {\n",
    "    0: \"No binding site\",\n",
    "    1: \"Binding site\"\n",
    "}\n",
    "\n",
    "# Print the predicted labels for each token\n",
    "for token, prediction in zip(tokens, predictions[0].numpy()):\n",
    "    if token not in ['<pad>', '<cls>', '<eos>']:\n",
    "        print((token, id2label[prediction]))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv_esmbind311",
   "language": "python",
   "name": "venv_esmbind311"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}