Update README.md

README.md

@@ -21,33 +21,26 @@ Raw format of training data (in Llama 3.1 chat template format):
 <|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\nDominant secondary structure of < V V F D V V F D V V F D V V F D V V F D V V F D V V F D V V F D ><|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\nUNSTRUCTURED<|eot_id|>
 ```
 
-## Sample inference
+Here is a visual representation of what the model predicts:
+
+![image/png](https://cdn-uploads.huggingface.co/production/uploads/623ce1c6b66fedf374859fe7/xOwcZGs6Q1NRV7ln5oP-m.png)
+
+## How to load the model 
 
 ```
-base_model_name = "lamm-mit/BioinspiredLlama-3-1-8B-128k"
-
-# Fine-tuned model name
-FT_model = "lamm-mit/BioinspiredLlama-3-1-8B-128k-dominant-protein-SS-structure"
-
-bnb_config4bit = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16,
-    bnb_4bit_use_double_quant=True,
-    use_nested_quant = False,
-)
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
 model = AutoModelForCausalLM.from_pretrained(
-    base_model_name,
+    'lamm-mit/BioinspiredLlama-3-1-8B-128k-dominant-protein-SS-structure',
     trust_remote_code=True,
     device_map="auto",
-    quantization_config= bnb_config4bit,
     torch_dtype =torch.bfloat16,
-
 ) 
-model = PeftModel.from_pretrained(model, FT_model,  )
 
-tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+tokenizer = AutoTokenizer.from_pretrained('lamm-mit/BioinspiredLlama-3-1-8B-128k-dominant-protein-SS-structure',)
 ```
+
 ## Example
 
 Inference function for convenience:
@@ -111,4 +104,19 @@ A visualization of the protein, to check:
 
 ![image/png](https://cdn-uploads.huggingface.co/production/uploads/623ce1c6b66fedf374859fe7/aO-0BbS8Sp_dV796w-Hm0.png)
 
-As predicted, this protein (PDB ID 6N7P, https://www.rcsb.org/structure/6N7P) is primarily alpha-helical. 
+As predicted, this protein (PDB ID 6N7P, https://www.rcsb.org/structure/6N7P) is primarily alpha-helical. 
+
+## Notes
+
+This model has been trained on sequences shorter than 128 amino acids. 
+
+## Reference
+
+```bibtex
+@article{Buehler_2024,
+  title={Fine-tuning LLMs for protein feature predictions},
+  author={Markus J. Buehler},
+  journal={},
+  year={2024}
+}
+```