Update README.md

README.md

@@ -9,4 +9,33 @@ ProGen2-small finetuned on 7 protein families.
 
 Bidirectional model trained on both N -> C and C -> N directions of protein sequences, specified by tokens "1" and "2" respectively. 
 
-See [github repo](https://github.com/hugohrban/ProGen2-finetuning/tree/main) for more info.
+See my [github repo](https://github.com/hugohrban/ProGen2-finetuning/tree/main) for more information.
+
+Example usage:
+
+```python
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+# optionally use local imports
+# from models.progen.modeling_progen import ProGenForCausalLM
+# from models.progen.configuration_progen import ProGenConfig
+import torch
+import torch.nn.functional as F
+
+# load model and tokenizer
+model = AutoModelForCausalLM.from_pretrained("hugohrban/progen2-small-mix7-bidi", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("hugohrban/progen2-small-mix7-bidi", trust_remote_code=True)
+
+# prepare input
+prompt = "<|pf00125|>2FDDDVSAVKSTGV"
+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}%")
+```