Update README.md

README.md

@@ -12,14 +12,14 @@ tags:
 - OAS
 ---
 
-# AbLang model for heavy chains
+### AbLang model for heavy chains
 
 This is a huggingface version of AbLang: A language model for antibodies. It was introduced in
 [this paper](https://doi.org/10.1101/2022.01.20.477061) and first released in
 [this repository](https://github.com/oxpig/AbLang). This model is trained on uppercase amino acids: it only works with capital letter amino acids.
 
 
-# Intended uses & limitations
+### Intended uses & limitations
 
 The model could be used for protein feature extraction or to be fine-tuned on downstream tasks (TBA).
 
@@ -42,6 +42,42 @@ Sequence embeddings can be produced as follows:
 
 TBA (just mean pool not including special tokens)
 
+### Fine-tune
+
+To save memory we recomend using [LoRA](https://doi.org/10.48550/arXiv.2106.09685):
+
+```python
+pip install git+https://github.com/huggingface/peft.git
+pip install loralib
+```
+
+LoRA greatly reduces the number of trainable parameters and performs on-par or better than fine-tuning full model.
+
+```python
+from peft import LoraConfig, get_peft_model
+
+def apply_lora_bert(model):
+    config = LoraConfig(
+        r=8, lora_alpha=32, 
+        lora_dropout=0.3,
+        target_modules=['query', 'value']
+    )
+    for param in model.parameters():
+        param.requires_grad = False  # freeze the model - train adapters later
+        if param.ndim == 1:
+        # cast the small parameters (e.g. layernorm) to fp32 for stability
+            param.data = param.data.to(torch.float32)
+    model.gradient_checkpointing_enable()  # reduce number of stored activations
+    model.enable_input_require_grads()
+    model = get_peft_model(model, config)
+    return model
+
+model = apply_lora_bert(model)
+
+model.print_trainable_parameters()
+# trainable params: 294912 || all params: 85493760 || trainable%: 0.3449514911965505
+```
+
 ### Citation
 ```
 @article{Olsen2022,