This repository contains the VESM protein language models developed in the paper
"VESM: Compressing the collective knowledge of ESM into a single protein language model" by Tuan Dinh, Seon-Kyeong Jang, Noah Zaitlen and Vasilis Ntranos.
A simple way to get started is to run our notebook directly on a Google Colab instance:
Getting Started with VESM 1 from huggingface_hub import snapshot_download , hf_hub_download
2
3 local_dir = './vesm'
4
5 # Download each model
6 model_offset = 0
7 model_name = [ "VESM_35M" , "VESM_150M" , "VESM_650M" , "VESM_3B" , "VESM3" ] [ model_offset ]
8 hf_hub_download ( repo_id = "ntranoslab/vesm" , filename = f" { model_name } .pth" , local_dir = local_dir )
9
10 # Download all models
11 snapshot_download ( repo_id = "ntranoslab/vesm" , local_dir = local_dir )
We provide a simple usage of our models for predicting variant effects.
1 import torch
2 from huggingface_hub import hf_hub_download
3 from transformers import AutoTokenizer , EsmForMaskedLM
4
5 esm_dict = {
6 "VESM_35M" : 'facebook/esm2_t12_35M_UR50D' ,
7 "VESM_150M" : 'facebook/esm2_t30_150M_UR50D' ,
8 "VESM_650M" : 'facebook/esm2_t33_650M_UR50D' ,
9 "VESM_3B" : 'facebook/esm2_t36_3B_UR50D' ,
10 "VESM3" : "esm3_sm_open_v1"
11 }
12 def load_vesm ( model_name = "VESM_3B" , local_dir = "vesm" , device = 'cuda' ) :
13 if model_name in esm_dict :
14 ckt = esm_dict [ model_name ]
15 else :
16 print ( "Model not found" )
17 return None
18
19 # download weights
20 hf_hub_download ( repo_id = "ntranoslab/vesm" , filename = f" { model_name } .pth" , local_dir = local_dir )
21 # load base model
22 if model_name == "VESM3" :
23 from esm . models . esm3 import ESM3
24 model = ESM3 . from_pretrained ( ckt , device = device ) . to ( torch . float )
25 tokenizer = model . tokenizers . sequence
26 else :
27 model = EsmForMaskedLM . from_pretrained ( ckt ) . to ( device )
28 tokenizer = AutoTokenizer . from_pretrained ( ckt )
29 # load pretrained VESM
30 model . load_state_dict ( torch . load ( f' { local_dir } / { model_name } .pth' ) , strict = False )
31 return model , tokenizer
1 # scoring functions
2 import torch . nn . functional as F
3 # calculate log-likelihood ratio from the logits
4 def get_llrs ( sequence_logits , input_ids ) :
5 token_probs = torch . log_softmax ( sequence_logits , dim = - 1 )
6 wt_positions = F . one_hot ( input_ids , num_classes = token_probs . shape [ - 1 ] )
7 wt_probs = token_probs * wt_positions
8 wt_probs = wt_probs . sum ( dim = - 1 , keepdim = True )
9 # add alpha
10 llrs = token_probs - wt_probs . expand ( token_probs . shape )
11 return llrs
12
13 # compute mutation score
14 def score_mutation ( llrs , mutation , sequence_vocabs ) :
15 mutation_score = 0
16 for mut in mutation . split ( ":" ) :
17 _ , idx , mt = mut [ 0 ] , int ( mut [ 1 : - 1 ] ) , mut [ - 1 ]
18 pred = llrs [ idx , sequence_vocabs [ mt ] ]
19 mutation_score += pred . item ( )
20 return mutation_score
Here, we provide sample scripts to compute mutation scores.
1 # sequence and mutation
2 sequence = "MVNSTHRGMHTSLHLWNRSSYRLHSNASESLGKGYSDGGCYEQLFVSPEVFVTLGVISLLENILV"
3 mutation = "M1Y:V2T"
1 # Setting
2 local_dir = 'vesm'
3 gpu_id = 0
4 device = torch . device ( f'cuda: { gpu_id } ' ) if torch . cuda . is_available ( ) else 'cpu'
5
6 # Helper
7 def inference ( model , tokenizer , sequence , device ) :
8 tokens = tokenizer ( [ sequence ] , return_tensors = 'pt' ) . to ( device )
9 with torch . no_grad ( ) :
10 outputs = model ( ** tokens )
11 logits = outputs [ 'logits' ] [ 0 ]
12 input_ids = tokens [ 'input_ids' ] [ 0 ]
13 # calculate log-likelihood ratio from the logits
14 llrs = get_llrs ( logits , input_ids )
15 return llrs
16
17 # Prediction with VESM
18 model_name = 'VESM_3B'
19 model , tokenizer = load_vesm ( model_name , local_dir = local_dir , device = device )
20 sequence_vocabs = tokenizer . get_vocab ( )
21 # compute mutation score
22 llrs = inference ( model , tokenizer , sequence , device )
23 mutation_score = score_mutation ( llrs , mutation , sequence_vocabs )
24 print ( f"Predicted score by { model_name } : " , mutation_score )
1 from esm . sdk . api import ESMProtein
2
3 # A sample structure pdb: download the latest version
4 # !wget https://alphafold.ebi.ac.uk/files/AF-P32245-F1-model_v6.pdb
5 pdb_file = "AF-P32245-F1-model_v6.pdb"
6 protein = ESMProtein . from_pdb ( pdb_file )
7 mutation = "M1Y:V2T"
1 # load model
2 model , tokenizer = load_vesm ( 'VESM3' , local_dir = local_dir , device = device )
3 sequence_vocabs = tokenizer . get_vocab ( )
4
5 # inference
6 tokens = model . encode ( protein )
7 seq_tokens = tokens . sequence . reshape ( 1 , - 1 )
8 struct_tokens = tokens . structure . reshape ( 1 , - 1 )
9 with torch . no_grad ( ) :
10 outs = model . forward ( sequence_tokens = seq_tokens , structure_tokens = struct_tokens )
11 logits = outs . sequence_logits [ 0 , : , : ]
12 input_ids = tokens . sequence
13
14 # calculate log-likelihood ratio from the logits
15 llrs = get_llrs ( logits , input_ids )
16 # compute mutation score
17 mutation_score = score_mutation ( llrs , mutation , sequence_vocabs )
18 print ( "mutation score: " , mutation_score )
The source code and model weights for VESM models are distributed under the MIT License.
The VESM3 model is a fine-tuned version of ESM3-Open (EvolutionaryScale) and is available under a
non-commercial license agreement .
