Update README.md

README.md

@@ -4,7 +4,7 @@ license: mit
 We provide two ways to use SaProt, including through huggingface class and 
 through the same way as in [esm github](https://github.com/facebookresearch/esm). Users can choose either one to use.
 
-### Huggingface model
+## Huggingface model
 The following code shows how to load the model.
 ```
 from transformers import EsmTokenizer, EsmForMaskedLM
@@ -33,11 +33,88 @@ torch.Size([1, 11, 446])
 """
 ```
 
-### esm model
+## esm model
 The esm version is also stored in the same folder, named `SaProt_650M_AF2.pt`. We provide a function to load the model.
 ```
 from utils.esm_loader import load_esm_saprot
 
 model_path = "/your/path/to/SaProt_650M_AF2.pt"
 model, alphabet = load_esm_saprot(model_path)
+```
+
+## Predict mutational effect
+We provide a function to predict the mutational effect of a protein sequence. The example below shows how to predict
+the mutational effect at a specific position. If using the AF2 structure, we strongly recommend that you add pLDDT mask (see below). 
+```python
+from model.saprot.saprot_foldseek_mutation_model import SaprotFoldseekMutationModel
+
+
+config = {
+    "foldseek_path": None,
+    "config_path": "/your/path/to/SaProt_650M_AF2", # Note this is the directory path of SaProt, not the ".pt" file
+    "load_pretrained": True,
+}
+model = SaprotFoldseekMutationModel(**config)
+tokenizer = model.tokenizer
+
+device = "cuda"
+model.eval()
+model.to(device)
+
+seq = "M#EvVpQpL#VyQdYaKv" # Here "#" represents lower plDDT regions (plddt < 70)
+
+# Predict the effect of mutating the 3rd amino acid to A
+mut_info = "V3A"
+mut_value = model.predict_mut(seq, mut_info)
+print(mut_value)
+
+# Predict all effects of mutations at 3rd position
+mut_pos = 3
+mut_dict = model.predict_pos_mut(seq, mut_pos)
+print(mut_dict)
+
+# Predict probabilities of all amino acids at 3rd position
+mut_pos = 3
+mut_dict = model.predict_pos_prob(seq, mut_pos)
+print(mut_dict)
+
+"""
+0.7908501625061035
+
+{'V3A': 0.7908501625061035, 'V3C': -0.9117952585220337, 'V3D': 2.7700226306915283, 'V3E': 2.3255627155303955, 'V3F': 0.2094242423772812, 'V3G': 2.699633836746216, 'V3H': 1.240191102027893, 'V3I': 0.10231903940439224, 'V3K': 1.804598093032837,
+'V3L': 1.3324960470199585, 'V3M': -0.18938277661800385, 'V3N': 2.8249857425689697, 'V3P': 0.40185314416885376, 'V3Q': 1.8361762762069702, 'V3R': 1.1899691820144653, 'V3S': 2.2159857749938965, 'V3T': 0.8813426494598389, 'V3V': 0.0, 'V3W': 0.5853186249732971, 'V3Y': 0.17449656128883362}
+
+{'A': 0.021275954321026802, 'C': 0.0038764977362006903, 'D': 0.15396881103515625, 'E': 0.0987202599644661, 'F': 0.011895398609340191, 'G': 0.14350374042987823, 'H': 0.03334535285830498, 'I': 0.010687196627259254, 'K': 0.058634623885154724, 'L': 0.03656982257962227, 'M': 0.00798324216157198, 'N': 0.16266827285289764, 'P': 0.014419485814869404, 'Q': 0.06051575019955635, 'R': 0.03171204403042793, 'S': 0.08847439289093018, 'T': 0.023291070014238358, 'V': 0.009647775441408157, 'W': 0.017323188483715057, 'Y': 0.011487090960144997}
+"""
+```
+
+## Get protein embeddings
+If you want to generate protein embeddings, you could refer to the following code. The embeddings are the average of
+the hidden states of the last layer.
+```python
+from model.saprot.base import SaprotBaseModel
+from transformers import EsmTokenizer
+
+
+config = {
+    "task": "base",
+    "config_path": "/your/path/to/SaProt_650M_AF2", # Note this is the directory path of SaProt, not the ".pt" file
+    "load_pretrained": True,
+}
+
+model = SaprotBaseModel(**config)
+tokenizer = EsmTokenizer.from_pretrained(config["config_path"])
+
+device = "cuda"
+model.to(device)
+
+seq = "M#EvVpQpL#VyQdYaKv" # Here "#" represents lower plDDT regions (plddt < 70)
+tokens = tokenizer.tokenize(seq)
+print(tokens)
+
+inputs = tokenizer(seq, return_tensors="pt")
+inputs = {k: v.to(device) for k, v in inputs.items()}
+
+embeddings = model.get_hidden_states(inputs, reduction="mean")
+print(embeddings[0].shape)
 ```