![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\"){kind=icon}
"
]
},
{
"cell_type": "markdown",
"source": [
"#ESM2AbCLIP: Antibody structure aware ESM2 model"
],
"metadata": {
"id": "JKV_0K5QjIzu"
}
},
{
"cell_type": "markdown",
"source": [
"OpenAI's CLIP model is pretty nifty connecting text with images. It efficiently learns visual concepts from natural language supervision. Can we do the same for protein sequences and structure? Instead of text and images, we're integrating protein sequences with their structural information.\n",
"\n",
"In this notebook we go through multimodal training using a contrastive learning framework, tailor-made for sequence-structure pretraining. By aligning antibody sequence data with their corresponding structural embeddings, we aim to develop a modeling approach that is better suited for antibody therapeutic property optimization such as antibody binding to its cognate antigen.\n",
"\n",
"**Key Steps**\n",
"\n",
"1. **Data Preparation**: Load and preprocess antibody sequence and structural data. Obtain structure embedding using frozen ESM-IF1 inverse-folding model.\n",
"2. **Model Definition**: Use the pre-trained ESM2 model as the base model and add a custom head for contrastive learning. The custom head projects both sequence and structural embeddings into a shared latent space.\n",
"3. **Contrastive Learning**: Implement a contrastive loss function and train the model to align sequence and structural embeddings. This process brings corresponding sequence and structure embeddings closer in the embedding space while pushing non-corresponding pairs further apart.\n",
"4. **Performance Evaluation**: Monitor performance with evaluation metrics such as Alignment, Uniformity, Contrastive Accuracy, and Top-K Accuracy.\n",
"5. **Training Setup**: Define training arguments and set up the training process using the HuggingFace `Trainer` class.\n"
],
"metadata": {
"id": "4Auwd1AR8Uuf"
}
},
{
"cell_type": "markdown",
"source": [
"## Setup"
],
"metadata": {
"id": "p9vnHkEAfBEx"
}
},
{
"cell_type": "code",
"source": [
"!pip install torch transformers accelerate &> /dev/null\n",
"!pip install --upgrade huggingface_hub &> /dev/null"
],
"metadata": {
"id": "LNeTW6Psdk5X"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"# To share your model with the community\n",
"# First store your authentication token from the Hugging Face website and then execute this cell\n",
"# Make sure to get token with WRITE access\n",
"from huggingface_hub import notebook_login\n",
"\n",
"notebook_login()"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 145,
"referenced_widgets": [
"5fb226d1c8b84c0aa28dd793248cd93f",
"fc0daa9f0dd04017a70d4400b7264865",
"fb1c9c255993422296bb681b3d31364f",
"69b99ce64af64ba2b14041195bbb586e",
"a7fae46c2fb148dc80df42cd467a2075",
"5972b4dab6ba4810b737cf9d84d7b08e",
"026da66a07cb4625a581bf314eddbcb5",
"ddb84a134bf64999bfec40f648a86d5f",
"e26636b6fc1c4942a55b7e6539490328",
"298cbb3879da48e7baeb1991fe92b405",
"94516a22592546cf8ab3af6be2a4d378",
"fbfa69fe2ef248868652df91af1292ff",
"eb35679b87e140d498aa7f92b8dd0dc9",
"64cb926c24764c0fa4a7e2a64bb45fe0",
"20d233b04f394b1ebe92070fe334df50",
"33b7cbc27df540f8b6f28dc312cdefc4",
"0aa118113fa1451ab718d932d595bb2c",
"c3e98f6300194894b2a01112162936aa",
"f198bcb8e0e54cd1aae44395433da88d",
"e21eb0ba33ec490f9219284fcdb925e8",
"f4522ef4f8ca42f395e5bb30c35c5a57",
"0a6ec3bfaa644dd9961fb40ea2e5e9bd",
"8590a0cfb42345bfb1a02bf3d724e834",
"f596eb7843d549f68313b3a760259e3e",
"84d85e4d38f445cea3091e9728603a84",
"ae413fe609e44b1cb8f9cddd07d69167",
"82b950ae3ba34f82b4ded1b0faa9c285",
"4f725058157b46bcb8de8aef17601a77",
"5953565c7df54d3b9c4fa93a3d9c77c6",
"c6f35e4424404209bd1023c7313c3d3d",
"6628874b8d7140a6a690ed845c6317fd",
"3ab630a7286344d8985dc05e2a40fe59"
]
},
"id": "sbjvpAIkCRoG",
"outputId": "c41ab10d-5795-4f3f-8f12-0f9f27e566c9"
},
"execution_count": null,
"outputs": [
{
"output_type": "display_data",
"data": {
"text/plain": [
"VBox(children=(HTML(value='| Epoch | \n", "Training Loss | \n", "Validation Loss | \n", "Alignment | \n", "Uniformity | \n", "Contrastive Accuracy | \n", "Top 3 Accuracy | \n", "
|---|---|---|---|---|---|---|
| 1 | \n", "1.015300 | \n", "1.295794 | \n", "1.272869 | \n", "-2.411410 | \n", "0.068027 | \n", "0.183673 | \n", "
| 2 | \n", "0.462300 | \n", "1.072861 | \n", "1.175534 | \n", "-2.800844 | \n", "0.095238 | \n", "0.265306 | \n", "
| 3 | \n", "0.264600 | \n", "1.086700 | \n", "1.151074 | \n", "-2.963047 | \n", "0.136054 | \n", "0.326531 | \n", "
| 4 | \n", "0.181500 | \n", "1.076758 | \n", "1.149029 | \n", "-3.040295 | \n", "0.129252 | \n", "0.367347 | \n", "
| 5 | \n", "0.140300 | \n", "1.052683 | \n", "1.142588 | \n", "-3.066312 | \n", "0.122449 | \n", "0.360544 | \n", "
"
]
},
"metadata": {}
},
{
"output_type": "execute_result",
"data": {
"text/plain": [
"CommitInfo(commit_url='https://huggingface.co/arjan-hada/esm2_t6_8M_UR50D-Ab-CLIP-v0/commit/8b840c6c73f563ccece9b4c599846cfeaa1dc2f2', commit_message='Training completed!', commit_description='', oid='8b840c6c73f563ccece9b4c599846cfeaa1dc2f2', pr_url=None, pr_revision=None, pr_num=None)"
],
"application/vnd.google.colaboratory.intrinsic+json": {
"type": "string"
}
},
"metadata": {},
"execution_count": 25
}
]
},
{
"cell_type": "code",
"source": [
"# Evaluate the model on the test set\n",
"test_result = trainer.evaluate(eval_dataset=test_ds)"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 37
},
"id": "YFc4m6xMOyLp",
"outputId": "4b995092-68d1-43a3-f4e9-c37357e0db03"
},
"execution_count": null,
"outputs": [
{
"output_type": "display_data",
"data": {
"text/plain": [
"\n",
" \n",
" \n",
" [19/19 00:01]\n",
" \n",
" "
]
},
"metadata": {}
}
]
},
{
"cell_type": "code",
"source": [
"# Print the results\n",
"print(f\"Test Loss: {test_result['eval_loss']}\")\n",
"for key, value in test_result.items():\n",
" if key != 'eval_loss':\n",
" print(f\"{key}: {value}\")"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "SB_IicqAO68l",
"outputId": "ca4697d5-3911-4a52-c669-cde210591941"
},
"execution_count": null,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Test Loss: 0.8743131756782532\n",
"eval_alignment: 1.021303653717041\n",
"eval_uniformity: -3.077335834503174\n",
"eval_contrastive_accuracy: 0.2266666740179062\n",
"eval_top_3_accuracy: 0.5066666603088379\n",
"eval_runtime: 1.4628\n",
"eval_samples_per_second: 102.546\n",
"eval_steps_per_second: 12.989\n",
"epoch: 5.0\n"
]
}
]
},
{
"cell_type": "code",
"source": [
"# Saves the best model\n",
"trainer.save_model(\"models/esm2_t6_8M_UR50D-Ab-CLIP-v0\")\n",
"model.config.save_pretrained(\"models/esm2_t6_8M_UR50D-Ab-CLIP-v0\")\n",
"tokenizer.save_pretrained(\"models/esm2_t6_8M_UR50D-Ab-CLIP-v0\")"
],
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 118,
"referenced_widgets": [
"70c245ec6efb47e780f1d09aa53dc1ff",
"86e74cf7d5844608ba13afb144e88e38",
"ab208a36766242afa41e20d62f828823",
"07265564a070462d96d88368ea3ae9c7",
"72cc3e413b1d4c028811b22536cc3fa4",
"d884a9f9646d4e0bad5df3f4f03e3e1b",
"0828c0b1150c4336b39a326716b79cc4",
"54c7bdd2c5a94d719b12f51bd09809db",
"766b91b655c2420998fd2d454a00ec3c",
"aae6adff761746709662a908efd706cf",
"6a750a9635a844ab947518efb1a1e7b5"
]
},
"id": "qJbjeV2HJIkX",
"outputId": "20efbc8a-a50c-4527-8339-dd95e1d86fb5"
},
"execution_count": null,
"outputs": [
{
"output_type": "display_data",
"data": {
"text/plain": [
"events.out.tfevents.1720548652.a3d6e2bd7b02.796.1: 0%| | 0.00/588 [00:00