Update README.md

README.md

@@ -2,7 +2,34 @@
 license: bsd-3-clause
 tags:
 - protein
+- progen2
 ---
 This is the one-directional model trained on 7 protein families.
 
-Check out the [github repo](https://github.com/hugohrban/ProGen2-finetuning) for more information.
+Check out the [github repo](https://github.com/hugohrban/ProGen2-finetuning) for more information.
+
+Example usage:
+
+```python
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+import torch
+import torch.nn.functional as F
+
+# load model and tokenizer
+model = AutoModelForCausalLM.from_pretrained("hugohrban/progen2-small-mix7", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("hugohrban/progen2-small-mix7", trust_remote_code=True)
+
+# prepare input
+prompt = "<|pf03668|>1MEVVIVTGMSGAGK"
+input_ids = torch.tensor(tokenizer.encode(prompt)).to(model.device)
+
+# forward pass
+logits = model(input_ids).logits
+
+# print output probabilities
+next_token_logits = logits[-1, :]
+next_token_probs = F.softmax(next_token_logits, dim=-1)
+for i, prob in enumerate(next_token_probs):
+    print(f"{tokenizer.decode(i)}: {100 * prob:.2f}%")
+```