ProteinBERT demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The examples in this notebook use a set of nine benchmarks described in our publication.\n",
    "# These benchmarks can be downloaded from: https://github.com/nadavbra/proteinbert_data_files/tree/master/protein_benchmarks\n",
    "# Download the benchmarks into a directory on your machine and set the following variable to the path of that directory.\n",
    "BENCHMARKS_DIR = '/path/to/protein_benchmarks'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-tune the model for the signal peptide benchmark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import pandas as pd\n",
    "from IPython.display import display\n",
    "\n",
    "from tensorflow import keras\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from proteinbert import OutputType, OutputSpec, FinetuningModelGenerator, load_pretrained_model, finetune, evaluate_by_len\n",
    "from proteinbert.conv_and_global_attention_model import get_model_with_hidden_layers_as_outputs\n",
    "\n",
    "BENCHMARK_NAME = 'signalP_binary'\n",
    "\n",
    "# A local (non-global) binary output\n",
    "OUTPUT_TYPE = OutputType(False, 'binary')\n",
    "UNIQUE_LABELS = [0, 1]\n",
    "OUTPUT_SPEC = OutputSpec(OUTPUT_TYPE, UNIQUE_LABELS)\n",
    "\n",
    "\n",
    "# Loading the dataset\n",
    "\n",
    "train_set_file_path = os.path.join(BENCHMARKS_DIR, '%s.train.csv' % BENCHMARK_NAME)\n",
    "train_set = pd.read_csv(train_set_file_path).dropna().drop_duplicates()\n",
    "train_set, valid_set = train_test_split(train_set, stratify = train_set['label'], test_size = 0.1, random_state = 0)\n",
    "\n",
    "test_set_file_path = os.path.join(BENCHMARKS_DIR, '%s.test.csv' % BENCHMARK_NAME)\n",
    "test_set = pd.read_csv(test_set_file_path).dropna().drop_duplicates()\n",
    "\n",
    "print(f'{len(train_set)} training set records, {len(valid_set)} validation set records, {len(test_set)} test set records.')\n",
    "\n",
    "\n",
    "# Loading the pre-trained model and fine-tuning it on the loaded dataset\n",
    "\n",
    "pretrained_model_generator, input_encoder = load_pretrained_model()\n",
    "\n",
    "# get_model_with_hidden_layers_as_outputs gives the model output access to the hidden layers (on top of the output)\n",
    "model_generator = FinetuningModelGenerator(pretrained_model_generator, OUTPUT_SPEC, pretraining_model_manipulation_function = \\\n",
    "        get_model_with_hidden_layers_as_outputs, dropout_rate = 0.5)\n",
    "\n",
    "training_callbacks = [\n",
    "    keras.callbacks.ReduceLROnPlateau(patience = 1, factor = 0.25, min_lr = 1e-05, verbose = 1),\n",
    "    keras.callbacks.EarlyStopping(patience = 2, restore_best_weights = True),\n",
    "]\n",
    "\n",
    "finetune(model_generator, input_encoder, OUTPUT_SPEC, train_set['seq'], train_set['label'], valid_set['seq'], valid_set['label'], \\\n",
    "        seq_len = 512, batch_size = 32, max_epochs_per_stage = 40, lr = 1e-04, begin_with_frozen_pretrained_layers = True, \\\n",
    "        lr_with_frozen_pretrained_layers = 1e-02, n_final_epochs = 1, final_seq_len = 1024, final_lr = 1e-05, callbacks = training_callbacks)\n",
    "\n",
    "\n",
    "# Evaluating the performance on the test-set\n",
    "\n",
    "results, confusion_matrix = evaluate_by_len(model_generator, input_encoder, OUTPUT_SPEC, test_set['seq'], test_set['label'], \\\n",
    "        start_seq_len = 512, start_batch_size = 32)\n",
    "\n",
    "print('Test-set performance:')\n",
    "display(results)\n",
    "\n",
    "print('Confusion matrix:')\n",
    "display(confusion_matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run all benchmarks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import pandas as pd\n",
    "from IPython.display import display\n",
    "\n",
    "from tensorflow import keras\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from proteinbert import OutputType, OutputSpec, FinetuningModelGenerator, load_pretrained_model, finetune, evaluate_by_len, log\n",
    "from proteinbert.conv_and_global_attention_model import get_model_with_hidden_layers_as_outputs\n",
    "\n",
    "BENCHMARKS = [\n",
    "    # name, output_type\n",
    "    ('signalP_binary', OutputType(False, 'binary')),\n",
    "    ('fluorescence', OutputType(False, 'numeric')),\n",
    "    ('remote_homology', OutputType(False, 'categorical')),\n",
    "    ('stability', OutputType(False, 'numeric')),\n",
    "    ('scop', OutputType(False, 'categorical')),\n",
    "    ('secondary_structure', OutputType(True, 'categorical')),\n",
    "    ('disorder_secondary_structure', OutputType(True, 'binary')),\n",
    "    ('ProFET_NP_SP_Cleaved', OutputType(False, 'binary')),\n",
    "    ('PhosphositePTM', OutputType(True, 'binary')),\n",
    "]\n",
    "\n",
    "settings = {\n",
    "    'max_dataset_size': None,\n",
    "    'max_epochs_per_stage': 40,\n",
    "    'seq_len': 512,\n",
    "    'batch_size': 32,\n",
    "    'final_epoch_seq_len': 1024,\n",
    "    'initial_lr_with_frozen_pretrained_layers': 1e-02,\n",
    "    'initial_lr_with_all_layers': 1e-04,\n",
    "    'final_epoch_lr': 1e-05,\n",
    "    'dropout_rate': 0.5,\n",
    "    'training_callbacks': [\n",
    "        keras.callbacks.ReduceLROnPlateau(patience = 1, factor = 0.25, min_lr = 1e-05, verbose = 1),\n",
    "        keras.callbacks.EarlyStopping(patience = 2, restore_best_weights = True),\n",
    "    ],\n",
    "}\n",
    "\n",
    "####### Uncomment for debug mode\n",
    "# settings['max_dataset_size'] = 500\n",
    "# settings['max_epochs_per_stage'] = 1\n",
    "\n",
    "def run_benchmark(benchmark_name, pretraining_model_generator, input_encoder, pretraining_model_manipulation_function = None):\n",
    "    \n",
    "    log('========== %s ==========' % benchmark_name)  \n",
    "    \n",
    "    output_type = get_benchmark_output_type(benchmark_name)\n",
    "    log('Output type: %s' % output_type)\n",
    "    \n",
    "    train_set, valid_set, test_set = load_benchmark_dataset(benchmark_name)        \n",
    "    log(f'{len(train_set)} training set records, {len(valid_set)} validation set records, {len(test_set)} test set records.')\n",
    "    \n",
    "    if settings['max_dataset_size'] is not None:\n",
    "        log('Limiting the training, validation and test sets to %d records each.' % settings['max_dataset_size'])\n",
    "        train_set = train_set.sample(min(settings['max_dataset_size'], len(train_set)), random_state = 0)\n",
    "        valid_set = valid_set.sample(min(settings['max_dataset_size'], len(valid_set)), random_state = 0)\n",
    "        test_set = test_set.sample(min(settings['max_dataset_size'], len(test_set)), random_state = 0)\n",
    "    \n",
    "    if output_type.is_seq or output_type.is_categorical:\n",
    "        train_set['label'] = train_set['label'].astype(str)\n",
    "        valid_set['label'] = valid_set['label'].astype(str)\n",
    "        test_set['label'] = test_set['label'].astype(str)\n",
    "    else:\n",
    "        train_set['label'] = train_set['label'].astype(float)\n",
    "        valid_set['label'] = valid_set['label'].astype(float)\n",
    "        test_set['label'] = test_set['label'].astype(float)\n",
    "        \n",
    "    if output_type.is_categorical:\n",
    "        \n",
    "        if output_type.is_seq:\n",
    "            unique_labels = sorted(set.union(*train_set['label'].apply(set)) | set.union(*valid_set['label'].apply(set)) | \\\n",
    "                    set.union(*test_set['label'].apply(set)))\n",
    "        else:\n",
    "            unique_labels = sorted(set(train_set['label'].unique()) | set(valid_set['label'].unique()) | set(test_set['label'].unique()))\n",
    "            \n",
    "        log('%d unique lebels.' % len(unique_labels))\n",
    "    elif output_type.is_binary:\n",
    "        unique_labels = [0, 1]\n",
    "    else:\n",
    "        unique_labels = None\n",
    "        \n",
    "    output_spec = OutputSpec(output_type, unique_labels)\n",
    "    model_generator = FinetuningModelGenerator(pretraining_model_generator, output_spec, pretraining_model_manipulation_function = \\\n",
    "            pretraining_model_manipulation_function, dropout_rate = settings['dropout_rate'])\n",
    "    finetune(model_generator, input_encoder, output_spec, train_set['seq'], train_set['label'], valid_set['seq'], valid_set['label'], \\\n",
    "            seq_len = settings['seq_len'], batch_size = settings['batch_size'], max_epochs_per_stage = settings['max_epochs_per_stage'], \\\n",
    "            lr = settings['initial_lr_with_all_layers'], begin_with_frozen_pretrained_layers = True, lr_with_frozen_pretrained_layers = \\\n",
    "            settings['initial_lr_with_frozen_pretrained_layers'], n_final_epochs = 1, final_seq_len = settings['final_epoch_seq_len'], \\\n",
    "            final_lr = settings['final_epoch_lr'], callbacks = settings['training_callbacks'])\n",
    "    \n",
    "    for dataset_name, dataset in [('Training-set', train_set), ('Validation-set', valid_set), ('Test-set', test_set)]:\n",
    "        \n",
    "        log('*** %s performance: ***' % dataset_name)\n",
    "        results, confusion_matrix = evaluate_by_len(model_generator, input_encoder, output_spec, dataset['seq'], dataset['label'], \\\n",
    "                start_seq_len = settings['seq_len'], start_batch_size = settings['batch_size'])\n",
    "    \n",
    "        with pd.option_context('display.max_rows', None, 'display.max_columns', None):\n",
    "            display(results)\n",
    "        \n",
    "        if confusion_matrix is not None:\n",
    "            with pd.option_context('display.max_rows', 16, 'display.max_columns', 10):\n",
    "                log('Confusion matrix:')\n",
    "                display(confusion_matrix)\n",
    "                \n",
    "    return model_generator\n",
    "\n",
    "def load_benchmark_dataset(benchmark_name):\n",
    "    \n",
    "    train_set_file_path = os.path.join(BENCHMARKS_DIR, '%s.train.csv' % benchmark_name)\n",
    "    valid_set_file_path = os.path.join(BENCHMARKS_DIR, '%s.valid.csv' % benchmark_name)\n",
    "    test_set_file_path = os.path.join(BENCHMARKS_DIR, '%s.test.csv' % benchmark_name)\n",
    "    \n",
    "    train_set = pd.read_csv(train_set_file_path).dropna().drop_duplicates()\n",
    "    test_set = pd.read_csv(test_set_file_path).dropna().drop_duplicates()\n",
    "          \n",
    "    if os.path.exists(valid_set_file_path):\n",
    "        valid_set = pd.read_csv(valid_set_file_path).dropna().drop_duplicates()\n",
    "    else:\n",
    "        log(f'Validation set {valid_set_file_path} missing. Splitting training set instead.')\n",
    "        train_set, valid_set = train_test_split(train_set, stratify = train_set['label'], test_size = 0.1, random_state = 0)\n",
    "    \n",
    "    return train_set, valid_set, test_set\n",
    "\n",
    "def get_benchmark_output_type(benchmark_name):\n",
    "    for name, output_type in BENCHMARKS:\n",
    "        if name == benchmark_name:\n",
    "            return output_type\n",
    "        \n",
    "pretrained_model_generator, input_encoder = load_pretrained_model()\n",
    "\n",
    "for benchmark_name, _ in BENCHMARKS:\n",
    "    run_benchmark(benchmark_name, pretrained_model_generator, input_encoder, pretraining_model_manipulation_function = \\\n",
    "            get_model_with_hidden_layers_as_outputs)\n",
    "        \n",
    "log('Done.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualizing the attention layers\n",
    "\n",
    "You can run this only after you have fine-tuned the model on a benchmark (e.g. signal peptide) and obtained *model_generator*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "BENCHMARK_DISPLAY_NAME = 'Signal peptide'\n",
    "\n",
    "TEST_SET_FILE_PATH = '/cs/phd/nadavb/my_public_ftp_site/protein_bert/protein_benchmarks/signalP_binary.train.csv'\n",
    "IDEAL_LEN = 80\n",
    "\n",
    "def calculate_attentions(model, input_encoder, seq, seq_len = None):\n",
    "    \n",
    "    from tensorflow.keras import backend as K\n",
    "    from proteinbert.tokenization import index_to_token\n",
    "    \n",
    "    if seq_len is None:\n",
    "        seq_len = len(seq) + 2\n",
    "    \n",
    "    X = input_encoder.encode_X([seq], seq_len)\n",
    "    (X_seq,), _ = X\n",
    "    seq_tokens = list(map(index_to_token.get, X_seq))\n",
    "\n",
    "    model_inputs = [layer.input for layer in model.layers if 'InputLayer' in str(type(layer))][::-1]\n",
    "    model_attentions = [layer.calculate_attention(layer.input) for layer in model.layers if 'GlobalAttention' in str(type(layer))]\n",
    "    invoke_model_attentions = K.function(model_inputs, model_attentions)\n",
    "    attention_values = invoke_model_attentions(X)\n",
    "    \n",
    "    attention_labels = []\n",
    "    merged_attention_values = []\n",
    "\n",
    "    for attention_layer_index, attention_layer_values in enumerate(attention_values):\n",
    "        for head_index, head_values in enumerate(attention_layer_values):\n",
    "            attention_labels.append('Attention %d - head %d' % (attention_layer_index + 1, head_index + 1))\n",
    "            merged_attention_values.append(head_values)\n",
    "\n",
    "    attention_values = np.array(merged_attention_values)\n",
    "    \n",
    "    return attention_values, seq_tokens, attention_labels\n",
    "\n",
    "def plot_attention(attention_values, seq_tokens, attention_labels, ax, cmap = 'Reds', vmin = 0, vmax = None, text_value_threshold = 0.1):\n",
    "\n",
    "    heatmap = ax.pcolor(attention_values.transpose(), cmap = cmap, vmin = vmin, vmax = vmax)\n",
    "\n",
    "    ax.set_xticks(np.arange(len(attention_labels)) + 0.5)\n",
    "    ax.set_xticklabels(attention_labels, rotation = 45, ha = 'right', fontsize = 12)\n",
    "    ax.set_yticks(np.arange(len(seq_tokens)) + 0.5)\n",
    "    ax.set_yticklabels(seq_tokens, fontsize = 12)\n",
    "\n",
    "    for i, row in enumerate(attention_values):\n",
    "        for j, value in enumerate(row):\n",
    "            if abs(value) >= text_value_threshold:\n",
    "                add_plus_sign = attention_values.min() < 0 and value > 0\n",
    "                plus_sign = '+' if add_plus_sign else ''\n",
    "                ax.text(i + 0.5, j + 0.5, plus_sign + '%d%%' % (100 * value), color = 'white', ha = 'center', va = 'center', \\\n",
    "                        fontsize = 9, fontweight = 'bold', fontstretch = 'condensed')\n",
    "                \n",
    "test_set = pd.read_csv(TEST_SET_FILE_PATH)\n",
    "chosen_index = ((test_set['seq'].str.len() - IDEAL_LEN).abs()).sort_values().index[0]\n",
    "seq = test_set.loc[chosen_index, 'seq']\n",
    "label = test_set.loc[chosen_index, 'label']\n",
    "                \n",
    "seq_len = len(seq) + 2\n",
    "\n",
    "pretrained_model_generator, input_encoder = load_pretrained_model()\n",
    "model = pretrained_model_generator.create_model(seq_len)\n",
    "pretrained_attention_values, pretrained_seq_tokens, pretrained_attention_labels = calculate_attentions(model, input_encoder, seq, \\\n",
    "        seq_len = seq_len)\n",
    "\n",
    "model = model_generator.create_model(seq_len)\n",
    "finetuned_attention_values, finetuned_seq_tokens, finetuned_attention_labels = calculate_attentions(model, input_encoder, seq, \\\n",
    "        seq_len = seq_len)\n",
    "assert finetuned_seq_tokens == pretrained_seq_tokens\n",
    "assert finetuned_attention_labels == pretrained_attention_labels[:len(finetuned_attention_labels)]\n",
    "\n",
    "fig, axes = plt.subplots(ncols = 4, figsize = (20, 0.2 * seq_len), gridspec_kw = dict(width_ratios = [1, 5, 1, 5]))\n",
    "fig.subplots_adjust(wspace = 0.3)\n",
    "\n",
    "axes[0].barh(np.arange(seq_len), 100 * pretrained_attention_values.sum(axis = 0), color = '#EC7063')\n",
    "axes[0].set_ylim((-0.5, seq_len - 0.5))\n",
    "axes[0].set_yticks([])\n",
    "axes[0].invert_xaxis()\n",
    "axes[0].set_xlabel('Total atten. %', fontsize = 14)\n",
    "\n",
    "vmax = pretrained_attention_values.max()\n",
    "plot_attention(pretrained_attention_values, pretrained_seq_tokens, pretrained_attention_labels, axes[1], cmap = 'Reds', vmax = vmax, \\\n",
    "        text_value_threshold = 0.05)\n",
    "axes[1].set_title('Only pre-training', fontsize = 16)\n",
    "\n",
    "axes[2].barh(np.arange(seq_len), 100 * (finetuned_attention_values - pretrained_attention_values).sum(axis = 0), color = '#28B463')\n",
    "axes[2].set_ylim((-0.5, seq_len - 0.5))\n",
    "axes[2].set_yticks([])\n",
    "axes[2].invert_xaxis()\n",
    "axes[2].set_xlabel('Total atten. % diff', fontsize = 14)\n",
    "\n",
    "attention_diff = finetuned_attention_values - pretrained_attention_values[:len(finetuned_attention_labels), :]\n",
    "vmax = np.abs(attention_diff).max()\n",
    "plot_attention(attention_diff, finetuned_seq_tokens, finetuned_attention_labels, axes[3], cmap = 'PiYG', vmin = -vmax, vmax = vmax, \\\n",
    "        text_value_threshold = 0.03)\n",
    "axes[3].set_title('%s fine-tuning' % BENCHMARK_DISPLAY_NAME, fontsize = 16)\n",
    "\n",
    "print(seq, label)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "my_python",
   "language": "python",
   "name": "my_python"
  },
  "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}