Update README.md

README.md

@@ -1,6 +1,34 @@
 ---
 license: bsd-3-clause
 ---
-Base ProGen2-small model.
+Mirror of the base ProGen2-small model.
 
-[Original repo](https://github.com/salesforce/progen/tree/main/progen2)
+[Original repo](https://github.com/salesforce/progen/tree/main/progen2)
+
+See also my github [repo](https://github.com/hugohrban/ProGen2-finetuning/tree/main) for an example of finetuning this model.
+
+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", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("hugohrban/progen2-small", trust_remote_code=True)
+
+# prepare input
+prompt = "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}%")
+```