Update README.md

README.md

@@ -54,7 +54,7 @@ esm_dict = {
     "VESM_3B": 'facebook/esm2_t36_3B_UR50D', 
     "VESM3": "esm3_sm_open_v1"
 }
-def load_vesm(model_name="VESM_35M", local_dir="vesm", device='cuda'):
+def load_vesm(model_name="VESM_3B", local_dir="vesm", device='cuda'):
     if model_name in esm_dict:
         ckt = esm_dict[model_name]
     else:
@@ -66,7 +66,7 @@ def load_vesm(model_name="VESM_35M", local_dir="vesm", device='cuda'):
     # load base model
     if model_name == "VESM3":
       from esm.models.esm3 import ESM3
-      model = ESM3.from_pretrained("esm3_sm_open_v1", device=device).to(torch.float)
+      model = ESM3.from_pretrained(ckt, device=device).to(torch.float)
       tokenizer = model.tokenizers.sequence
     else:
       model = EsmForMaskedLM.from_pretrained(ckt).to(device)
@@ -131,7 +131,7 @@ def inference(model, tokenizer, sequence, device):
     Prediction with VESM models
 """
 # load vesm models
-model_name = 'VESM_35M'
+model_name = 'VESM_3B'
 model, tokenizer = load_vesm(model_name, local_dir=local_dir, device=device)
 sequence_vocabs = tokenizer.get_vocab()
 # inference
@@ -145,9 +145,9 @@ print(f"Predicted score by {model_name}: ", mutant_score)
 ```py
 from esm.sdk.api import ESMProtein
 
-# A sample structure pdb
-# !wget https://alphafold.ebi.ac.uk/files/AF-P32245-F1-model_v4.pdb
-pdb_file = "AF-P32245-F1-model_v4.pdb"
+# A sample structure pdb: download the latest version
+# !wget https://alphafold.ebi.ac.uk/files/AF-P32245-F1-model_v6.pdb
+pdb_file = "AF-P32245-F1-model_v6.pdb"
 protein = ESMProtein.from_pdb(pdb_file)
 mutant = "M1Y:V2T"
 ```
@@ -166,7 +166,7 @@ with torch.no_grad():
   logits = outs.sequence_logits[0, :, :]
   input_ids = tokens.sequence
 
-# calcualte log-likelihood ratio from the logits 
+# calculate log-likelihood ratio from the logits 
 llrs = get_llrs(logits, input_ids)
 # compute mutant score
 mutant_score = score_mutant(llrs, mutant, sequence_vocabs)