ESM2AbCLIP: Antibody structure aware ESM2 model¶

!pip install torch transformers accelerate &> /dev/null
!pip install --upgrade huggingface_hub &> /dev/null

import os
from pathlib import Path
import json

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import pandas as pd

import pandas as pd
import torch
from torch.utils.data import Dataset

class AntibodyDataset(Dataset):
    """
    Initialize the dataset.

    Args:
        data_path (str): Path to the pickle file containing data.
        tokenizer (transformers.PreTrainedTokenizer): Tokenizer to process the sequences.
    """
    def __init__(self, data_path, tokenizer):
        self.data = pd.read_pickle(data_path)
        self.tokenizer = tokenizer

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        """
        Get item by index.
        """
        row = self.data.iloc[idx]
        sequence = row['sequence']
        embedding = torch.tensor(row['embedding'], dtype=torch.float32)

        inputs = self.tokenizer(sequence, return_tensors='pt', padding=False,
                                truncation=False)

        return {
            'input_ids': inputs['input_ids'].squeeze(),
            'attention_mask': inputs['attention_mask'].squeeze(),
            'labels': embedding
        }

# Initialize DataCollator
data_collator = DataCollatorWithPadding(tokenizer)

# Initialize DataLoader with DataCollator
batch_size=5
train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True, collate_fn=data_collator)
valid_dl = DataLoader(valid_ds, batch_size=batch_size, shuffle=False, collate_fn=data_collator)
test_dl = DataLoader(test_ds, batch_size=batch_size, shuffle=False, collate_fn=data_collator)

import numpy as np
import torch
import torch.nn as nn

class ESMCLIPHead(nn.Module):
    """Head for CLIP multimodal tasks."""

    def __init__(self, config, projection_dim=512, dropout_prob=0.0):
        super().__init__()
        self.sequence_projection = nn.Sequential(
            nn.Linear(config.hidden_size, projection_dim),
            nn.GELU(),
            nn.Dropout(dropout_prob)
        )
        self.structure_projection = nn.Sequential(
            nn.Linear(projection_dim, projection_dim),
            nn.GELU(),
            nn.Dropout(dropout_prob)
        )

    def forward(self, sequence_output, structure_embeddings):
        sequence_projected = self.sequence_projection(sequence_output[:, 0, :]) # [CLS] token pooling
        structure_projected = self.structure_projection(structure_embeddings)
        return sequence_projected, structure_projected

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import EsmPreTrainedModel, EsmModel, EsmConfig

class ESMForCLIP(EsmPreTrainedModel):
    config_class = EsmConfig

    def __init__(self, config):
        super().__init__(config)
        self.esm = EsmModel(config, add_pooling_layer=False)
        self.clip_head = ESMCLIPHead(config)
        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, labels=None):
        outputs = self.esm(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
        sequence_output = outputs.last_hidden_state
        sequence_projected, structure_projected = self.clip_head(sequence_output, labels)
        # Clamp the logit scale value to ensure it does not exceed log(100)
        self.logit_scale.data.clamp_(max=np.log(100.0))

        return {
            'logits': (sequence_projected, structure_projected)
        }

import torch
import torch.nn as nn
import torch.nn.functional as F

class ContrastiveLoss(nn.Module):
    def __init__(self):
        super(ContrastiveLoss, self).__init__()

    def forward(self, seq_embeddings, struct_embeddings, logit_scale):
        # Normalize embeddings to unit vectors
        seq_embeddings = F.normalize(seq_embeddings, dim=1)
        struct_embeddings = F.normalize(struct_embeddings, dim=1)

        # Compute pairwise cosine similarities and scale with temperature
        temperature = torch.exp(logit_scale)
        logits_per_seq = torch.mm(seq_embeddings, struct_embeddings.t()) * temperature
        logits_per_struct = logits_per_seq.t()

        # Labels for contrastive loss
        labels = torch.arange(seq_embeddings.size(0)).to(seq_embeddings.device)

        # Contrastive loss as described in the paper
        loss_seq = F.cross_entropy(logits_per_seq, labels)
        loss_struct = F.cross_entropy(logits_per_struct, labels)
        loss = (loss_seq + loss_struct) / 2

        return loss

import torch
import torch.nn.functional as F

# Function to compute the Alignment metric
def compute_alignment(seq_embeddings, struct_embeddings):
    distances = (seq_embeddings - struct_embeddings).pow(2).sum(dim=1)
    alignment = distances.mean().item()
    return alignment

# Function to compute the Uniformity metric
# This function is computationally intensive and may take a while to run
def compute_uniformity(embeddings):
    pairwise_distances = torch.cdist(embeddings, embeddings, p=2).pow(2)
    uniformity = torch.log(torch.exp(-2 * pairwise_distances).mean()).item()
    return uniformity

# Function to compute the Cosine Similarity metric
def compute_cosine_similarity(seq_embeddings, struct_embeddings):
    """
    This function normalizes the sequence and structure embeddings to unit
    vectors and then computes the cosine similarity between each pair using
    matrix multiplication.
    """
    cosine_sim = torch.mm(seq_embeddings, struct_embeddings.t())
    return cosine_sim

# Function to compute the Contrastive Accuracy
def compute_contrastive_accuracy(cosine_sim):
    """
    This function finds the index of the maximum cosine similarity for each
    sequence embedding and compares it to the correct index.
    It then computes the mean accuracy.
    """
    correct_preds = cosine_sim.argmax(dim=1)
    correct = correct_preds == torch.arange(cosine_sim.size(0)).to(cosine_sim)
    return correct.float().mean().item()

# Function to compute the Top-K Accuracy
def compute_top_k_accuracy(cosine_sim, k=1):
    """
    This function finds the top K predictions for each sequence embedding and
    checks if the correct match is within these top K predictions.
    It then computes the mean accuracy.
    """
    top_k_preds = cosine_sim.topk(k, dim=1)[1]
    correct = torch.arange(cosine_sim.size(0)).unsqueeze(1).expand_as(top_k_preds)
    correct = correct == top_k_preds
    return correct.any(dim=1).float().mean().item()

def compute_metrics(eval_pred):
    seq_embeddings, struct_embeddings = eval_pred.predictions
    seq_embeddings = torch.tensor(seq_embeddings)
    struct_embeddings = torch.tensor(struct_embeddings)

    # Normalize embeddings
    seq_embeddings = F.normalize(seq_embeddings, dim=1)
    struct_embeddings = F.normalize(struct_embeddings, dim=1)

    # Compute metrics
    alignment = compute_alignment(seq_embeddings, struct_embeddings)
    combined_embeddings = torch.cat((seq_embeddings, struct_embeddings), dim=0)
    uniformity = compute_uniformity(combined_embeddings)
    cosine_sim = compute_cosine_similarity(seq_embeddings, struct_embeddings)
    contrastive_accuracy = compute_contrastive_accuracy(cosine_sim)
    top_3_accuracy = compute_top_k_accuracy(cosine_sim, k=3)

    metrics = {
        "alignment": alignment,
        "uniformity": uniformity,
        "contrastive_accuracy": contrastive_accuracy,
        "top_3_accuracy": top_3_accuracy
    }
    return metrics

from transformers import Trainer

class ContrastiveTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        outputs = model(**inputs)  # Forward pass to get logits
        logits = outputs['logits']

        # Access logit_scale from the underlying model
        logit_scale = model.logit_scale

        loss_fct = ContrastiveLoss()  # Instantiate the loss function
        loss = loss_fct(logits[0], logits[1], logit_scale)  # Compute the loss

        return (loss, outputs) if return_outputs else loss

import gc # Python's garbage collection module

def clear_memory():
    gc.collect() # explicitly triggers garbage collection, free up memory
    if torch.cuda.is_available(): torch.cuda.empty_cache() # clears the PyTorch CUDA memory cache

from transformers import Trainer
# Initialize DataCollator
data_collator = DataCollatorWithPadding(tokenizer)

trainer = ContrastiveTrainer(
    model=model,
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=valid_ds,
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics
)

# Optionally set the max_split_size_mb to avoid fragmentation issues
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'