Update README.md

README.md

@@ -7,13 +7,18 @@ license: apache-2.0
 # ProteinForceGPT: Generative strategies for modeling, design and analysis of protein mechanics
 
 
-### Load model
+### Basic information
+
+This protein language model is an autoregressive transformer model in GPT-style, trained to analyze and predict the mechanical properties of a large number of protein sequences. 
 
-This model is an autoregressive transformer model in GPT-style, trained to analyze and predict the mechanical properties of a large number of protein sequences. 
+This protein language foundation model was based on the NeoGPT-X architecture and uses rotary positional embeddings (RoPE). It has 16 attention heads, 36 hidden layers and a hidden size of 1024, an intermediate size of 4086 and uses a GeLU activation function. 
 
 The pretraining task is defined as "Sequence<...>" where ... is an amino acid sequence.
 
-Mechanics-related forward and inverse tasks are:
+Pretraining dataset: https://huggingface.co/datasets/lamm-mit/GPTProteinPretrained 
+Pretrained model: https://huggingface.co/lamm-mit/GPTProteinPretrained
+
+In this fine-tuned model, mechanics-related forward and inverse tasks are:
 
 ```raw
 CalculateForce<GEECDCGSPSNP..>, 
@@ -25,20 +30,21 @@ GenerateForce<0.220>
 GenerateForceEnergy<0.262,0.220> 
 GenerateForceHistory<0.004,0.034,0.125,0.142,0.159,0.102,0.079,0.073,0.131,0.105,0.071,0.058,0.072,0.060,0.049,0.114,0.122,0.108,0.173,0.192,0.208,0.153,0.212,0.222,0.244>
 ```
-Load pretrained model:
+
+### Load model
+
+Load model:
 
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-pretrained_model_name='lamm-mit/ProteinForceGPT'
+ForceGPT_model_name='lamm-mit/ProteinForceGPT'
 
-tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(ForceGPT_model_name, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
 
-model_name = pretrained_model_name
- 
 model = AutoModelForCausalLM.from_pretrained(
-    model_name, 
+    ForceGPT_model_name, 
     trust_remote_code=True
 ).to(device)
 
@@ -91,7 +97,7 @@ Output:
 0: CalculateForce<GEECDCGSPSNPCCDAATCKLRPGAQCADGLCCDQCRFKKKRTICRIARGDFPDDRCTGQSADCPRWN> [0.262]```
 ```
 
-## Citation
+## Citations
 To cite this work:
 ```
 @article{GhafarollahiBuehler_2024,
@@ -103,4 +109,17 @@ To cite this work:
     pages   = {},
     url     = {}
 }
-```
+```
+
+The dataset used to fine-tune the model is available at:
+
+@article{GhafarollahiBuehler_2024,
+    title   = {ForceGen: End-to-end de novo protein generation based on nonlinear mechanical unfolding responses using a protein language diffusion model},
+    author  = {B. Ni, D.L. Kaplan, M.J. Buehler},
+    journal = {Science Advances},
+    year    = {2024},
+    volume  = {},
+    pages   = {},
+    url     = {}
+}
+```