AbLang model for light chains
This is a 🤗 version of AbLang: A language model for antibodies. It was introduced in this paper and first released in this repository. This model is trained on uppercase amino acids: it only works with capital letter amino acids.
Intended uses & limitations
The model could be used for protein feature extraction or to be fine-tuned on downstream tasks (TBA).
How to use
Here is how to use this model to get the features of a given antibody sequence in PyTorch:
from transformers import AutoModel, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained('qilowoq/AbLang_light')
model = AutoModel.from_pretrained('qilowoq/AbLang_light', trust_remote_code=True)
sequence_Example = ' '.join("GSELTQDPAVSVALGQTVRITCQGDSLRNYYASWYQQKPRQAPVLVFYGKNNRPSGIPDRFSGSSSGNTASLTISGAQAEDEADYYCNSRDSSSNHLVFGGGTKLTVLSQ")
encoded_input = tokenizer(sequence_Example, return_tensors='pt')
model_output = model(**encoded_input)
Sequence embeddings can be produced as follows:
def get_sequence_embeddings(encoded_input, model_output):
mask = encoded_input['attention_mask'].float()
d = {k: v for k, v in torch.nonzero(mask).cpu().numpy()} # dict of sep tokens
# make sep token invisible
for i in d:
mask[i, d[i]] = 0
mask[:, 0] = 0.0 # make cls token invisible
mask = mask.unsqueeze(-1).expand(model_output.last_hidden_state.size())
sum_embeddings = torch.sum(model_output.last_hidden_state * mask, 1)
sum_mask = torch.clamp(mask.sum(1), min=1e-9)
return sum_embeddings / sum_mask
seq_embeds = get_sequence_embeddings(encoded_input, model_output)
Fine-tune
To save memory we recomend using LoRA:
pip install git+https://github.com/huggingface/peft.git
pip install loralib
LoRA greatly reduces the number of trainable parameters and performs on-par or better than fine-tuning full model.
from peft import LoraConfig, get_peft_model
def apply_lora_bert(model):
config = LoraConfig(
r=8, lora_alpha=32,
lora_dropout=0.3,
target_modules=['query', 'value']
)
for param in model.parameters():
param.requires_grad = False # freeze the model - train adapters later
if param.ndim == 1:
# cast the small parameters (e.g. layernorm) to fp32 for stability
param.data = param.data.to(torch.float32)
model.gradient_checkpointing_enable() # reduce number of stored activations
model.enable_input_require_grads()
model = get_peft_model(model, config)
return model
model = apply_lora_bert(model)
model.print_trainable_parameters()
# trainable params: 294912 || all params: 85493760 || trainable%: 0.3449514911965505
Citation
@article{Olsen2022,
title={AbLang: An antibody language model for completing antibody sequences},
author={Tobias H. Olsen, Iain H. Moal and Charlotte M. Deane},
journal={bioRxiv},
doi={https://doi.org/10.1101/2022.01.20.477061},
year={2022}
}