13_TCR_bert.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use of Transfer Learning and Large Language Models methods: TCR-BERT "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First exploration of TCR-BERT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Different model classes needed for different models\n",
    "from transformers import BertModel\n",
    "# This model was pretrained on MAA and TRB classification\n",
    "tcrbert_model = BertModel.from_pretrained(\"wukevin/tcr-bert\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import model_utils  # found under the \"tcr\" folder that can be found in the GitHub repo: https://github.com/wukevin/tcr-bert\n",
    "\n",
    "# Load a TextClassificationPipeline using TCR-BERT\n",
    "tcrbert_trb_cls = model_utils.load_classification_pipeline(\"wukevin/tcr-bert\", device=-1) # device=0 if cuda is available\n",
    "# For the pipeline, input amino acids are expected to be spaced\n",
    "results = model_utils.reformat_classification_pipeline_preds(tcrbert_trb_cls([\n",
    "    \"C A S S P V T G G I Y G Y T F\",  # Binds to NLVPMVATV CMV antigen\n",
    "    \"C A T S G R A G V E Q F F\",      # Binds to GILGFVFTL flu antigen\n",
    "]))  # Return a dataframe where each column is an antigen, each row corresponds to an input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get highest probability antigen\n",
    "results.idxmax(axis=1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Application to our dataset to perform prediction (without finetuning) and evaluation of its **raw** performance "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "df = pd.read_csv('./peptide-detail-ci_filtered_aligned_merged.csv')\n",
    "df[\"Amino Acids 1\"] = df[\"Amino Acids\"].apply(lambda x: x.split(\",\")[0])\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_unique = df.drop_duplicates(subset=['CDR3'])\n",
    "df_unique.reset_index(inplace=True, drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Count the number of occurences of each antigen\n",
    "print(df_unique[\"Amino Acids 1\"].value_counts())\n",
    "# select only antigens with between 1000 and 5000 occurences\n",
    "df_unique_balanced = df_unique[df_unique[\"Amino Acids 1\"].isin(df_unique[\"Amino Acids 1\"].value_counts().index[(df_unique[\"Amino Acids 1\"].value_counts() > 2000) & (df_unique[\"Amino Acids 1\"].value_counts() < 5000)])]\n",
    "print(len(df_unique_balanced))\n",
    "print(len(df_unique_balanced[\"Amino Acids 1\"].unique()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 1000\n",
    "store = []\n",
    "true_labels = []\n",
    "for i in range(N):\n",
    "    # Change format, by putting spaces between each letter in CDR3\n",
    "    cdr3 = df_unique[\"CDR3\"][i]\n",
    "    cdr3 = \" \".join(cdr3)\n",
    "    store.append(cdr3)\n",
    "    true_labels.append(df_unique[\"Amino Acids 1\"][i])\n",
    "\n",
    "results = model_utils.reformat_classification_pipeline_preds(tcrbert_trb_cls(store))  # Return a dataframe where each column is an antigen, each row corresponds to an input\n",
    "print(len(store))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 1000\n",
    "sample_unique_balanced = df_unique_balanced.sample(N)\n",
    "store_balanced = []\n",
    "true_labels_balanced = []\n",
    "for i in range(N):\n",
    "    # Change format, by putting spaces between each letter in CDR3\n",
    "    cdr3 = sample_unique_balanced[\"CDR3\"].iloc[i]\n",
    "    cdr3 = \" \".join(cdr3)\n",
    "    store_balanced.append(cdr3)\n",
    "    true_labels_balanced.append(sample_unique_balanced[\"Amino Acids 1\"].iloc[i])\n",
    "\n",
    "results_balanced = model_utils.reformat_classification_pipeline_preds(tcrbert_trb_cls(store_balanced))  # Return a dataframe where each column is an antigen, each row corresponds to an input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preds = results.idxmax(axis=1)\n",
    "preds_unique = results.idxmax(axis=1).unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = model_utils.get_transformer_embeddings(\n",
    "            model_dir=\"wukevin/tcr-bert\",\n",
    "            seqs=store,\n",
    "            layers=[-1],\n",
    "            method=\"mean\",\n",
    "            device=3,\n",
    "        )\n",
    "\n",
    "# use UMAP to reduce dimensionality\n",
    "import umap\n",
    "reducer = umap.UMAP(random_state=42)\n",
    "embedding_umap = reducer.fit_transform(embeddings)\n",
    "\n",
    "# Associate each prediction with a number for plotting\n",
    "preds_color = []\n",
    "for i in range(len(preds)):\n",
    "    preds_color.append(list(preds_unique).index(preds[i]))\n",
    "\n",
    "# plot with labels as colors \n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "sns.set(style='white', context='poster', rc={'figure.figsize':(14,10)})\n",
    "plt.scatter(embedding_umap[:, 0], embedding_umap[:, 1], c=preds_color, s=20, cmap='Spectral')\n",
    "plt.gca().set_aspect('equal', 'datalim')\n",
    "plt.colorbar(boundaries=np.arange(len(preds_unique)+1)-0.5).set_ticks(np.arange(len(preds_unique)))\n",
    "plt.title('UMAP projection of the TCR-BERT embeddings with labels from TCR-BERT (without finetuning)', fontsize=16);\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get unique of true labels\n",
    "true_labels_unique = np.unique(true_labels)\n",
    "\n",
    "# Associate each prediction with a number for plotting\n",
    "true_labels_color = []\n",
    "for i in range(len(true_labels)):\n",
    "    true_labels_color.append(list(true_labels_unique).index(true_labels[i]))\n",
    "\n",
    "# plot with labels as colors \n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "sns.set(style='white', context='poster', rc={'figure.figsize':(14,10)})\n",
    "plt.scatter(embedding_umap[:, 0], embedding_umap[:, 1], c=true_labels_color, s=20, cmap='Spectral')\n",
    "plt.gca().set_aspect('equal', 'datalim')\n",
    "plt.colorbar(boundaries=np.arange(len(true_labels_unique)+1)-0.5).set_ticks(np.arange(len(true_labels_unique)))\n",
    "plt.title('UMAP projection of the TCR-BERT embeddings with true labels (without finetuning)', fontsize=16);\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = model_utils.get_transformer_embeddings(\n",
    "            model_dir=\"wukevin/tcr-bert\",\n",
    "            seqs=store_balanced,\n",
    "            layers=[-1],\n",
    "            method=\"mean\",\n",
    "            device=3,\n",
    "        )\n",
    "\n",
    "# use UMAP to reduce dimensionality\n",
    "import umap\n",
    "reducer = umap.UMAP(random_state=42)\n",
    "embedding_umap = reducer.fit_transform(embeddings)\n",
    "\n",
    "true_labels_balanced_unique = np.unique(true_labels_balanced)\n",
    "# Associate each prediction with a number for plotting\n",
    "true_labels_balanced_color = []\n",
    "for i in range(len(true_labels_balanced)):\n",
    "    true_labels_balanced_color.append(list(true_labels_balanced_unique).index(true_labels_balanced[i]))\n",
    "\n",
    "# plot with labels as colors \n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "sns.set(style='white', context='poster', rc={'figure.figsize':(14,10)})\n",
    "plt.scatter(embedding_umap[:, 0], embedding_umap[:, 1], c=true_labels_balanced_color, s=20, cmap='Spectral')\n",
    "plt.gca().set_aspect('equal', 'datalim')\n",
    "plt.colorbar(boundaries=np.arange(len(preds_unique)+1)-0.5).set_ticks(np.arange(len(preds_unique)))\n",
    "plt.title('UMAP projection of the TCR-BERT embeddings with labels from TCR-BERT (without finetuning)', fontsize=16);\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use PCA to reduce dimensionality\n",
    "from sklearn.decomposition import PCA\n",
    "pca = PCA(n_components=2)\n",
    "embedding_pca = pca.fit_transform(embeddings)\n",
    "\n",
    "\n",
    "\n",
    "true_labels_balanced_unique = np.unique(true_labels_balanced)\n",
    "# Associate each prediction with a number for plotting\n",
    "true_labels_balanced_color = []\n",
    "for i in range(len(true_labels_balanced)):\n",
    "    true_labels_balanced_color.append(list(true_labels_balanced_unique).index(true_labels_balanced[i]))\n",
    "\n",
    "# plot with labels as colors \n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "sns.set(style='white', context='poster', rc={'figure.figsize':(14,10)})\n",
    "plt.scatter(embedding_pca[:, 0], embedding_pca[:, 1], c=true_labels_balanced_color, s=20, cmap='Spectral')\n",
    "plt.gca().set_aspect('equal', 'datalim')\n",
    "plt.colorbar(boundaries=np.arange(len(preds_unique)+1)-0.5).set_ticks(np.arange(len(preds_unique)))\n",
    "plt.title('PCA projection of the TCR-BERT embeddings with labels from TCR-BERT (without finetuning)', fontsize=16);\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_unique_antigen = df.drop_duplicates(subset=['Amino Acids 1'])\n",
    "df_unique_antigen.reset_index(inplace=True, drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = len(df_unique_antigen)\n",
    "store = []\n",
    "true_labels = []\n",
    "for i in range(N):\n",
    "    # Change format, by putting spaces between each letter in CDR3\n",
    "    cdr3 = df_unique_antigen[\"CDR3\"][i]\n",
    "    cdr3 = \" \".join(cdr3)\n",
    "    store.append(cdr3)\n",
    "    true_labels.append(df_unique_antigen[\"Amino Acids 1\"][i])\n",
    "\n",
    "results = model_utils.reformat_classification_pipeline_preds(tcrbert_trb_cls(store))  # Return a dataframe where each column is an antigen, each row corresponds to an input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Associate each prediction with a number for plotting\n",
    "true_labels_color = []\n",
    "for i in range(len(true_labels)):\n",
    "    true_labels_color.append(list(true_labels_unique).index(true_labels[i]))\n",
    "print(len(true_labels_color))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = model_utils.get_transformer_embeddings(\n",
    "            model_dir=\"wukevin/tcr-bert\",\n",
    "            seqs=store,\n",
    "            layers=[-1],\n",
    "            method=\"mean\",\n",
    "            device=3,\n",
    "        )\n",
    "\n",
    "# use UMAP to reduce dimensionality\n",
    "import umap\n",
    "reducer = umap.UMAP(random_state=42)\n",
    "embedding_umap = reducer.fit_transform(embeddings)\n",
    "\n",
    "# plot with labels as colors \n",
    "sns.set(style='white', context='poster', rc={'figure.figsize':(14,10)})\n",
    "plt.scatter(embedding_umap[:, 0], embedding_umap[:, 1], c=true_labels_color, s=20, cmap='Spectral')\n",
    "plt.gca().set_aspect('equal', 'datalim')\n",
    "plt.title('UMAP projection of the TCR-BERT embeddings with labels from TCR-BERT (without finetuning)', fontsize=16);\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preds = results.idxmax(axis=1)\n",
    "preds_unique = results.idxmax(axis=1).unique()\n",
    "\n",
    "embeddings = model_utils.get_transformer_embeddings(\n",
    "            model_dir=\"wukevin/tcr-bert\",\n",
    "            seqs=store,\n",
    "            layers=[-1],\n",
    "            method=\"mean\",\n",
    "            device=3,\n",
    "        )\n",
    "\n",
    "# use UMAP to reduce dimensionality\n",
    "import umap\n",
    "reducer = umap.UMAP(random_state=42)\n",
    "embedding_umap = reducer.fit_transform(embeddings)\n",
    "\n",
    "# Associate each prediction with a number for plotting\n",
    "preds_color = []\n",
    "for i in range(len(preds)):\n",
    "    preds_color.append(list(preds_unique).index(preds[i]))\n",
    "\n",
    "# plot with labels as colors \n",
    "sns.set(style='white', context='poster', rc={'figure.figsize':(14,10)})\n",
    "plt.scatter(embedding_umap[:, 0], embedding_umap[:, 1], c=preds_color, s=20, cmap='Spectral')\n",
    "plt.gca().set_aspect('equal', 'datalim')\n",
    "plt.colorbar(boundaries=np.arange(len(preds_unique)+1)-0.5).set_ticks(np.arange(len(preds_unique)))\n",
    "plt.title('UMAP projection of the TCR-BERT embeddings with labels from TCR-BERT (without finetuning)', fontsize=16);\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Perform k-means clustering\n",
    "from sklearn.cluster import KMeans\n",
    "kmeans = KMeans(n_clusters=10, random_state=0).fit(embedding_umap)\n",
    "kmeans.labels_\n",
    "\n",
    "# plot with labels as colors\n",
    "sns.set(style='white', context='poster', rc={'figure.figsize':(14,10)})\n",
    "plt.scatter(embedding_umap[:, 0], embedding_umap[:, 1], c=kmeans.labels_, s=20, cmap='Spectral')\n",
    "plt.gca().set_aspect('equal', 'datalim')\n",
    "plt.colorbar(boundaries=np.arange(11)-0.5).set_ticks(np.arange(10))\n",
    "plt.title('UMAP projection of the TCR-BERT embeddings with k-means clustering (without finetuning)', fontsize=16);\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add small plots to big plot for labels of groups \n",
    "def add_subplot_axes(ax,rect,axisbg='w'):\n",
    "    fig = plt.gcf()\n",
    "    box = ax.get_position()\n",
    "    width = box.width\n",
    "    height = box.height\n",
    "    inax_position  = ax.transAxes.transform(rect[0:2])\n",
    "    transFigure = fig.transFigure.inverted()\n",
    "    infig_position = transFigure.transform(inax_position)    \n",
    "    x = infig_position[0]\n",
    "    y = infig_position[1]\n",
    "    width *= rect[2]\n",
    "    height *= rect[3]  # <= Typo was here\n",
    "    #subax = fig.add_axes([x,y,width,height],facecolor=facecolor)  # matplotlib 2.0+\n",
    "    subax = fig.add_axes([x,y,width,height])\n",
    "    x_labelsize = subax.get_xticklabels()[0].get_size()\n",
    "    y_labelsize = subax.get_yticklabels()[0].get_size()\n",
    "    x_labelsize *= rect[2]**0.5\n",
    "    y_labelsize *= rect[3]**0.5\n",
    "    subax.xaxis.set_tick_params(labelsize=x_labelsize)\n",
    "    subax.yaxis.set_tick_params(labelsize=y_labelsize)\n",
    "    return subax\n",
    "\n",
    "# Perform k-means clustering\n",
    "from sklearn.cluster import KMeans\n",
    "import pyrepseq as rs\n",
    "import pyrepseq.plotting as rsp\n",
    "n_clusters = 10\n",
    "kmeans = KMeans(n_clusters=n_clusters, random_state=0).fit(embedding_umap)\n",
    "\n",
    "# plot with labels as colors\n",
    "fig = plt.figure(figsize=(20,10))\n",
    "sns.set(style='white', context='poster', rc={'figure.figsize':(14,10)})\n",
    "ax = fig.add_subplot(111)\n",
    "ax.scatter(embedding_umap[:, 0], embedding_umap[:, 1], c=kmeans.labels_, s=20, cmap='Spectral')\n",
    "ax.set_title('UMAP projection of the TCR-BERT embeddings with k-means clustering (without finetuning)', fontsize=24);\n",
    "\n",
    "\n",
    "for k in range(n_clusters):\n",
    "    # Select only the rows with the k means group k\n",
    "    df_unique_antigen_k = df_unique_antigen[kmeans.labels_ == k]\n",
    "    # get coordinate of center of cluster \n",
    "    center = kmeans.cluster_centers_[k]\n",
    "    # plot small plots by specifying size of axes\n",
    "    subax1 = add_subplot_axes(ax, [center[0]/10-0.8, center[1]/10-0.5, 0.14, 0.08])\n",
    "    rsp.seqlogos(df_unique_antigen_k[\"CDR3_al\"], ax=subax1)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fine-tuning of TCR-BERT on our dataset and evaluation of its **fine-tuned** performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sample 1000 rows \n",
    "df_sample = df.drop_duplicates(subset=['CDR3'])\n",
    "print(len(df_sample))\n",
    "df_sample = df_sample.sample(n=100000, random_state=1)\n",
    "df_sample.reset_index(inplace=True, drop=True)\n",
    "# Rename 'CDR3' to 'TRB' \n",
    "df_sample.rename(columns={'CDR3': 'TRB'}, inplace=True)\n",
    "# Rename 'Amino Acids 1' to 'label'\n",
    "df_sample.rename(columns={'Amino Acids 1': 'label'}, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save cdr3 and labels as tsv file \n",
    "df_sample[[\"TRB\", \"label\"]].to_csv(\"tcrbert_trb_cls_100k.tsv\", sep=\"\\t\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Finetune transformer that accepts a SINGLE input (either TRA or TRB)\n",
    "as opposed to two inputs (i.e. TRA and TRB) (see finetune_transformer.py)\n",
    "\"\"\"\n",
    "\n",
    "### Useful links:\n",
    "### https://huggingface.co/transformers/custom_datasets.html\n",
    "\n",
    "from enum import unique\n",
    "import os, sys\n",
    "import logging\n",
    "import json\n",
    "import itertools\n",
    "import argparse\n",
    "from typing import *\n",
    "\n",
    "import tqdm\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn import metrics\n",
    "from scipy.special import softmax\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "import skorch\n",
    "import skorch.helper\n",
    "import neptune\n",
    "\n",
    "from transformers import (\n",
    "    BertForSequenceClassification,\n",
    "    BertConfig,\n",
    "    BertTokenizer,\n",
    "    DataCollatorWithPadding,\n",
    "    TrainingArguments,\n",
    "    Trainer,\n",
    "    EvalPrediction,\n",
    "    EarlyStoppingCallback,\n",
    ")\n",
    "\n",
    "import git\n",
    "import data_loader as dl\n",
    "import featurization as ft\n",
    "import model_utils\n",
    "import utils\n",
    "from transformer_custom import BertForSequenceClassificationMulti\n",
    "\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "\n",
    "METRICS = []  # Global tracker for metrics states\n",
    "\n",
    "\n",
    "def sigmoid(x):\n",
    "    \"\"\"Sigmoid\"\"\"\n",
    "    return 1 / (1 + np.exp(-x))\n",
    "\n",
    "\n",
    "def compute_metrics(pred: EvalPrediction) -> Dict[str, float]:\n",
    "    \"\"\"Compute binary metrics to report\"\"\"\n",
    "    labels = pred.label_ids\n",
    "    labels_expanded = np.zeros(pred.predictions.shape)\n",
    "    labels_expanded[np.arange(labels.size), labels.squeeze()] += 1\n",
    "    preds = pred.predictions.argmax(-1)\n",
    "    preds_probs = softmax(pred.predictions, axis=1)\n",
    "    if preds_probs.shape[1] == 2:\n",
    "        preds_probs = preds_probs[:, 1:]\n",
    "        labels_expanded = labels_expanded[:, 1:]\n",
    "    assert (\n",
    "        labels_expanded.shape == preds_probs.shape\n",
    "    ), f\"Got differing shapes: {labels_expanded.shape} {preds_probs.shape}\"\n",
    "    acc = metrics.accuracy_score(labels, preds)\n",
    "\n",
    "    # Compute averages per category\n",
    "    auroc_values, auprc_values = [], []\n",
    "    for i in range(preds_probs.shape[-1]):\n",
    "        l = labels_expanded[:, i]\n",
    "        p = preds_probs[:, i]\n",
    "        if len(np.unique(l)) < 2:\n",
    "            continue\n",
    "        auroc_values.append(metrics.roc_auc_score(l, p))\n",
    "        auprc_values.append(metrics.average_precision_score(l, p))\n",
    "\n",
    "    # Update global metrics tracker\n",
    "    global METRICS\n",
    "    METRICS.append(\n",
    "        {\"auroc_per_label\": auroc_values, \"auprc_per_label\": auprc_values,}\n",
    "    )\n",
    "\n",
    "    auroc = np.mean(auroc_values)\n",
    "    auprc = np.mean(auprc_values)\n",
    "    return {\n",
    "        \"accuracy\": acc,\n",
    "        \"auroc\": auroc,\n",
    "        \"auprc\": auprc,\n",
    "    }\n",
    "\n",
    "\n",
    "def compute_metrics_multi(pred: EvalPrediction) -> Dict[str, float]:\n",
    "    \"\"\"Compute multi-label metrics\"\"\"\n",
    "    labels = pred.label_ids\n",
    "    preds = pred.predictions\n",
    "    preds_sigmoid = sigmoid(preds)\n",
    "    assert labels.shape == preds.shape == preds_sigmoid.shape\n",
    "    # For both, this calculates per-class metric and averages across classes\n",
    "    auroc_values, auprc_values, used_classes, pos_rate = [], [], [], []\n",
    "    for j in range(labels.shape[1]):\n",
    "        if len(set(labels[:, j])) == 1:  # These are undefined\n",
    "            continue\n",
    "        used_classes.append(j)\n",
    "        auroc = metrics.roc_auc_score(\n",
    "            labels[:, j], preds_sigmoid[:, j], average=\"macro\"\n",
    "        )\n",
    "        auprc = metrics.average_precision_score(\n",
    "            labels[:, j], preds_sigmoid[:, j], average=\"macro\"\n",
    "        )\n",
    "        auroc_values.append(auroc)\n",
    "        auprc_values.append(auprc)\n",
    "        pos_rate.append(float(np.mean(labels[:, j])))\n",
    "    METRICS.append(\n",
    "        {\n",
    "            \"auroc_per_label\": auroc_values,\n",
    "            \"auprc_per_label\": auprc_values,\n",
    "            \"pos_per_class\": pos_rate,\n",
    "            \"used_classes\": used_classes,\n",
    "        }\n",
    "    )\n",
    "    return {\"auroc\": np.mean(auroc_values), \"auprc\": np.mean(auprc_values)}\n",
    "\n",
    "\n",
    "def load_data_single(\n",
    "    keyword: Literal[\"lcmv\", \"vdjdb\", \"pird\", \"covid\", \"nsclc\"],\n",
    "    segment: Literal[\"TRB\", \"TRA\"] = \"TRB\",\n",
    ") -> List[Tuple[str, str]]:\n",
    "    \"\"\"\n",
    "    Load the data\n",
    "    Returns pairs of (antigen, seq) or more generally (label, seq)\n",
    "    Comma (,) can be included in antigen/label string to indicate multi-label\n",
    "    Antigen/label can also be empty string to indicate multi-label\n",
    "    \"\"\"\n",
    "    logging.info(f\"Loading {keyword} {segment} data\")\n",
    "    assert segment == \"TRA\" or segment == \"TRB\"\n",
    "    if keyword.lower() == \"lcmv\":\n",
    "        if segment != \"TRB\":\n",
    "            raise NotImplementedError\n",
    "        lcmv_data = dl.load_lcmv_table()\n",
    "        lcmv_trb, lcmv_labels = dl.dedup_lcmv_table_trb_only(lcmv_data)\n",
    "        lcmv_labels = np.array(\n",
    "            [\"TetPos\" in t or \"TetMid\" in t for t in lcmv_data[\"tetramer\"]]\n",
    "        )\n",
    "        lcmv_antigen = [\n",
    "            lcmv_data.iloc[i][\"antigen.sequence\"] if lcmv_labels[i] else \"\"\n",
    "            for i in range(len(lcmv_data))\n",
    "        ]\n",
    "        return list(zip(lcmv_antigen, lcmv_trb))\n",
    "    elif keyword.lower() == \"vdjdb\":  # Evey sequence here has a corresponding epitope\n",
    "        vdjdb_tab = dl.load_vdjdb(tra_trb_filter=[segment])\n",
    "        vdjdb_antigens = vdjdb_tab[\"antigen.epitope\"]\n",
    "        vdjdb_aa = vdjdb_tab[\"cdr3\"]\n",
    "        return list(zip(vdjdb_antigens, vdjdb_aa))\n",
    "    elif keyword.lower() == \"pird\":\n",
    "        tcr_key = f\"CDR3.{'beta' if segment == 'TRB' else 'alpha'}.aa\"\n",
    "        pird_tab = dl.load_pird(with_antigen_only=True)\n",
    "        pird_tab_tcr_nonnull = pird_tab.loc[~pd.isnull(pird_tab[tcr_key])]\n",
    "        pird_antigens = pird_tab_tcr_nonnull[\"Antigen.sequence\"]\n",
    "        pird_tcrs = pird_tab_tcr_nonnull[tcr_key]\n",
    "        logging.info(\n",
    "            f\"Loaded PIRD {segment} dataset with {len(pird_tcrs)} antigen-TCR pairs\"\n",
    "        )\n",
    "        return list(zip(pird_antigens, pird_tcrs))\n",
    "    elif keyword.lower() == \"covid\":\n",
    "        if segment == \"TRA\":\n",
    "            raise NotImplementedError\n",
    "        logging.info(\"Loading immuneACCESS SARS-CoV-2 TRBs with TCRdb background\")\n",
    "        covid_table = dl.load_immuneaccess_mira_covid()\n",
    "        covid_trbs = list(covid_table[\"TCR_aa\"])\n",
    "        covid_antigens = list(covid_table[\"Amino Acids\"])\n",
    "        logging.info(f\"Loaded {len(covid_trbs)} SARS-COV-2 TRBs\")\n",
    "\n",
    "        # Dedup and merge TRBs with multiple labels\n",
    "        uniq_trbs, uniq_antigens = dl.dedup_and_merge_labels(covid_trbs, covid_antigens)\n",
    "\n",
    "        # Sample negatives\n",
    "        tcrdb = dl.load_tcrdb()\n",
    "        random_tcr_sampler = np.random.default_rng(seed=12345)\n",
    "        random_tcr_idx = random_tcr_sampler.choice(\n",
    "            np.arange(len(tcrdb)), size=len(uniq_trbs), replace=False\n",
    "        )\n",
    "        random_tcr = [tcrdb[\"AASeq\"][i] for i in random_tcr_idx]\n",
    "        random_antigens = [\"\"] * len(random_tcr)  # Blank or unknown\n",
    "        logging.info(f\"Sampled {len(random_antigens)} random TRBs\")\n",
    "\n",
    "        return list(zip(uniq_antigens + random_antigens, uniq_trbs + random_tcr))\n",
    "    elif keyword.lower() == \"nsclc\":\n",
    "        if segment == \"TRA\":\n",
    "            raise NotImplementedError\n",
    "        logging.info(\"Loading NSCLC TRBs\")\n",
    "        nsclc_table = dl.load_reuben_nsclc()\n",
    "        return list(zip(nsclc_table[\"label\"], nsclc_table[\"aminoAcid\"]))\n",
    "    elif os.path.isfile(keyword):\n",
    "        df = pd.read_csv(keyword, sep=\"\\t\")\n",
    "        tcrs = df[segment]\n",
    "        labels = df[\"label\"]\n",
    "        return list(zip(labels, tcrs))\n",
    "    else:\n",
    "        raise ValueError(f\"Unrecognized data keyword or file not found: {keyword}\")\n",
    "\n",
    "\n",
    "def load_data(\n",
    "    keywords: Iterable[Literal[\"lcmv\", \"VDJdb\", \"PIRD\"]],\n",
    "    segments: Iterable[Literal[\"TRA\", \"TRB\"]],\n",
    "    blacklist: Optional[Iterable[str]] = None,\n",
    "    use_multilabel: bool = False,\n",
    ") -> Dataset:\n",
    "    \"\"\"Load and concatenate the data\"\"\"\n",
    "    # First, load blacklist\n",
    "    blacklist_seqs = set()\n",
    "    if blacklist is not None:\n",
    "        for blacklist_fname in blacklist:\n",
    "            seqs = utils.read_newline_file(blacklist_fname)\n",
    "            blacklist_seqs.update(seqs)\n",
    "    if blacklist_seqs:\n",
    "        logging.info(f\"{len(blacklist_seqs)} blacklisted sequences: {blacklist_seqs}\")\n",
    "\n",
    "    # Load in data\n",
    "    antigens_seqs, aa_seqs = [], []\n",
    "    for kw in keywords:\n",
    "        for seg in segments:\n",
    "            d = load_data_single(kw, segment=seg)\n",
    "            antigen, aa = zip(*d)\n",
    "            antigens_seqs.extend(antigen)\n",
    "            aa_seqs.extend(aa)\n",
    "            assert len(antigens_seqs) == len(aa_seqs)\n",
    "    # If we have multiple labels or cases with no labels, we are in multilabel setting\n",
    "    if use_multilabel:\n",
    "        is_multilabel = any([\",\" in a or not a for a in antigens_seqs])\n",
    "        if not is_multilabel:\n",
    "            raise ValueError(\"Failed to interpret labels as multi labels\")\n",
    "        multilabels = [tuple(a.split(\",\")) for a in antigens_seqs]\n",
    "\n",
    "        # Check for overlaps\n",
    "        uniq_combos = utils.dedup(multilabels)\n",
    "        has_overlaps = False\n",
    "        for i, j in itertools.product(uniq_combos, uniq_combos):\n",
    "            i, j = set(i), set(j)\n",
    "            if i != j and i.intersection(j):\n",
    "                logging.debug(f\"Found overlap: {i}, {j}\")\n",
    "                has_overlaps = True\n",
    "        if not has_overlaps:\n",
    "            logging.warning(\n",
    "                \"Labels are multi labels, but no overlaps between labels! Is multilabel necessary?\"\n",
    "            )\n",
    "\n",
    "        unique_labels = [\n",
    "            seq\n",
    "            for seq in utils.dedup(itertools.chain.from_iterable(multilabels))\n",
    "            if seq\n",
    "        ]  # Excludes empty string\n",
    "        label_to_idx = {l: i for i, l in enumerate(unique_labels)}\n",
    "        labels = np.zeros((len(multilabels), len(unique_labels)), dtype=np.float32)\n",
    "        for i, label_set in enumerate(multilabels):\n",
    "            idx = [label_to_idx[a] for a in label_set if a]\n",
    "            labels[i, idx] = 1\n",
    "    else:\n",
    "        onehot = ft.one_hot(antigens_seqs, alphabet=None)\n",
    "        labels = np.where(onehot)[0]  # Index encoding, e.g. [0, 4, 2, 3]\n",
    "        unique_labels = utils.dedup(antigens_seqs)\n",
    "    logging.info(f\"Generated labels of shape {labels.shape}\")\n",
    "\n",
    "    # Create dataset to return\n",
    "    dset = dl.TcrFineTuneSingleDataset(\n",
    "        aa_seqs, labels, label_continuous=False, label_labels=unique_labels\n",
    "    )\n",
    "    return dset\n",
    "\n",
    "\n",
    "def get_bert_classifier(\n",
    "    path: str, labels: Sequence[str], problem_type: str = \"single_label_classification\"\n",
    ") -> BertForSequenceClassification:\n",
    "    \"\"\"Get BERT classifier model form the path\"\"\"\n",
    "    bert_class = (\n",
    "        BertForSequenceClassification\n",
    "        if problem_type == \"single_label_classification\"\n",
    "        else BertForSequenceClassificationMulti\n",
    "    )\n",
    "    logging.info(f\"Loading BERT classifier for {problem_type}: {bert_class}\")\n",
    "    if os.path.isdir(path) or path.startswith(\"wukevin/\"):  # Poor heuristic\n",
    "        logging.info(f\"Loading BERT classifier from {path} with {len(labels)} labels\")\n",
    "        # Remake config with correct number of labels\n",
    "        cfg = BertConfig.from_pretrained(\n",
    "            path,\n",
    "            num_labels=len(labels),\n",
    "            id2label={str(i): l for i, l in enumerate(labels)},\n",
    "            label2id={l: i for i, l in enumerate(labels)},\n",
    "        )\n",
    "        retval = bert_class.from_pretrained(path)\n",
    "        # Manually define the classifier layer to avoid shape mismatches\n",
    "        retval.config = cfg\n",
    "        retval.num_labels = len(labels)\n",
    "        retval.classifier = nn.Linear(\n",
    "            in_features=retval.classifier.in_features, out_features=len(labels)\n",
    "        )\n",
    "    elif os.path.isfile(path) and path.split(\".\")[-1] == \"json\":\n",
    "        logging.info(\n",
    "            f\"Loading newly initialized BERT from {path} with {len(labels)} labels\"\n",
    "        )\n",
    "        json_args = utils.load_json_params(path)\n",
    "        cfg = BertConfig(\n",
    "            **json_args,\n",
    "            vocab_size=len(ft.AMINO_ACIDS_WITH_ALL_ADDITIONAL),\n",
    "            pad_token_id=ft.AMINO_ACIDS_WITH_ALL_ADDITIONAL_TO_IDX[ft.PAD],\n",
    "            num_labels=len(labels),\n",
    "            id2label=dict(enumerate(labels)),\n",
    "        )\n",
    "        retval = bert_class(cfg)\n",
    "    else:\n",
    "        raise ValueError(f\"Unrecognized value: {path}\")\n",
    "    return retval\n",
    "\n",
    "\n",
    "def build_parser():\n",
    "    parser = argparse.ArgumentParser(\n",
    "        description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"-p\", \"--pretrained\", type=str, required=True, help=\"Pretrained network\"\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--data\",\n",
    "        type=str,\n",
    "        nargs=\"*\",\n",
    "        help=\"Datasets to train, keywords or filenames (files should be formatted as csv)\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"-b\",\n",
    "        \"--blacklist\",\n",
    "        type=str,\n",
    "        nargs=\"*\",\n",
    "        required=False,\n",
    "        help=\"File containing labels to ignore\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--multilabel\",\n",
    "        action=\"store_true\",\n",
    "        help=\"Try to interpret labels as multilabel\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"-s\",\n",
    "        \"--segment\",\n",
    "        type=str,\n",
    "        choices=[\"TRA\", \"TRB\"],\n",
    "        required=True,\n",
    "        nargs=\"*\",\n",
    "        help=\"TRA or TRB\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"-o\", \"--outdir\", type=str, default=os.getcwd(), help=\"Directory to save model\"\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"-d\", \"--downsample\", type=float, default=1.0, help=\"Downsample training data\"\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"-m\",\n",
    "        \"--monitor\",\n",
    "        type=str,\n",
    "        default=\"auprc\",\n",
    "        choices=[\"auroc\", \"auprc\", \"acc\", \"loss\"],\n",
    "        help=\"Metric to monitor for best model\",\n",
    "    )\n",
    "    parser.add_argument(\"--bs\", type=int, default=128, help=\"Batch size\")\n",
    "    parser.add_argument(\"--lr\", type=float, default=5e-5, help=\"Learning rate\")\n",
    "    parser.add_argument(\"-e\", \"--epochs\", type=int, default=25, help=\"Max num epochs\")\n",
    "    parser.add_argument(\n",
    "        \"-w\", \"--warmup\", type=float, default=0.1, help=\"Proportion of steps to warmup\"\n",
    "    )\n",
    "    parser.add_argument(\"--debug\", action=\"store_true\")\n",
    "    return parser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "args = {\n",
    "    \"--pretrained\": \"wukevin/tcr-bert\",\n",
    "    \"--data\": [\"tcrbert_trb_cls.tsv\"],\n",
    "    \"--segment\": [\"TRB\"],\n",
    "    \"--blacklist\": None,\n",
    "    \"--multilabel\": False,\n",
    "    \"--outdir\": \"finetuned_test1\",\n",
    "    \"--downsample\": 1.0,\n",
    "    \"--monitor\": \"loss\",\n",
    "    \"--bs\": 128,\n",
    "    \"--lr\": 5e-5,\n",
    "    \"--epochs\": 25,\n",
    "    \"--warmup\": 0.1,\n",
    "    \"--debug\": False,\n",
    "}\n",
    "if not os.path.isdir(args[\"--outdir\"]):\n",
    "    os.makedirs(args[\"--outdir\"])\n",
    "\n",
    "# Setup logging\n",
    "logger = logging.getLogger()\n",
    "fh = logging.FileHandler(os.path.join(args[\"--outdir\"], \"classifier_training.log\"), \"w\")\n",
    "fh.setLevel(logging.INFO)\n",
    "logger.addHandler(fh)\n",
    "\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    logging.info(f\"PyTorch CUDA version: {torch.version.cuda}\")\n",
    "for arg in args:\n",
    "    logging.info(f\"Parameter {arg}: {args[arg]}\")\n",
    "\n",
    "with open(os.path.join(args[\"--outdir\"], \"params.json\"), \"w\") as sink:\n",
    "    json.dump(args, sink, indent=4)\n",
    "\n",
    "full_dset = load_data(args[\"--data\"], args[\"--segment\"], args[\"--blacklist\"], args[\"--multilabel\"])\n",
    "if full_dset.label_labels is not None:\n",
    "    with open(os.path.join(args[\"--outdir\"], \"trained_labels.txt\"), \"w\") as sink:\n",
    "        for l in full_dset.label_labels:\n",
    "            sink.write(l + \"\\n\")\n",
    "train_dataset = dl.DatasetSplit(full_dset, split=\"train\")\n",
    "valid_dataset = dl.DatasetSplit(full_dset, split=\"valid\")\n",
    "test_dataset = dl.DatasetSplit(full_dset, split=\"test\")\n",
    "\n",
    "# Write datasets, or downsample\n",
    "if not args[\"--debug\"]:\n",
    "    # DownsampledDataset doesn't support writing\n",
    "    logging.info(\"Writing train/valid/test dataset to disk\")\n",
    "    train_dataset.to_file(os.path.join(args[\"--outdir\"], \"train_dataset.json\"))\n",
    "    valid_dataset.to_file(os.path.join(args[\"--outdir\"], \"valid_dataset.json\"))\n",
    "    test_dataset.to_file(os.path.join(args[\"--outdir\"], \"test_dataset.json\"))\n",
    "else:\n",
    "    train_dataset = dl.DownsampledDataset(train_dataset, 0.05)\n",
    "\n",
    "if args[\"--downsample\"] < 1.0:\n",
    "    logging.info(f\"Downsampling training set to {args['--downsample']}\")\n",
    "    train_dataset = dl.DownsampledDataset(train_dataset, downsample=args[\"--downsample\"])\n",
    "\n",
    "classifier = get_bert_classifier(\n",
    "    args[\"--pretrained\"],\n",
    "    labels=full_dset.label_labels,\n",
    "    problem_type=\"multi_label_classification\"\n",
    "    if full_dset.is_multilabel\n",
    "    else \"single_label_classification\",\n",
    ")\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=args[\"--outdir\"],\n",
    "    overwrite_output_dir=True,\n",
    "    num_train_epochs=args[\"--epochs\"],\n",
    "    per_device_train_batch_size=args[\"--bs\"],\n",
    "    per_device_eval_batch_size=args[\"--bs\"],\n",
    "    learning_rate=args[\"--lr\"],\n",
    "    warmup_ratio=args[\"--warmup\"],\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    logging_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    save_total_limit=1,\n",
    "    dataloader_num_workers=8,\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=f\"eval_{args['--monitor']}\",\n",
    "    no_cuda=args[\"--debug\"],  # Useful for debugging\n",
    "    skip_memory_metrics=True,\n",
    "    disable_tqdm=False,\n",
    "    logging_dir=os.path.join(args[\"--outdir\"], \"logs\"),\n",
    "    use_mps_device=True,\n",
    ")\n",
    "\n",
    "# Early stop uses metric_for_best_model from above\n",
    "callbacks = [EarlyStoppingCallback(early_stopping_patience=5)]\n",
    "trainer = Trainer(\n",
    "    model=classifier,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=valid_dataset,\n",
    "    compute_metrics=compute_metrics_multi\n",
    "    if full_dset.is_multilabel\n",
    "    else compute_metrics,\n",
    "    callbacks=callbacks,\n",
    ")\n",
    "trainer.train()\n",
    "\n",
    "# Save the global metrics before additioanl eval calls\n",
    "with open(os.path.join(args[\"--outdir\"], \"custom_metrics.json\"), \"w\") as sink:\n",
    "    json.dump(METRICS, sink, indent=4)\n",
    "\n",
    "# Save model and perform final evaluation\n",
    "trainer.save_model(args[\"--outdir\"])\n",
    "valid_eval_metrics = trainer.evaluate()\n",
    "for k in sorted(valid_eval_metrics.keys()):\n",
    "    logging.info(f\"{k}\\t{valid_eval_metrics[k]:.4f}\")\n",
    "logging.info(\"Test set:\")\n",
    "test_eval_metrics = trainer.evaluate(test_dataset)\n",
    "for k in sorted(test_eval_metrics.keys()):\n",
    "    logging.info(f\"{k}\\t{test_eval_metrics[k]:.4f}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "m4r",
   "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.9.2"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}