Update README.md

README.md

@@ -1,3 +1,43 @@
----
-license: bsd-3-clause
----
+---
+license: bsd-3-clause
+---
+# ProGen2-small
+
+HF mirror for ProGen2-small for **Protein Engineering**
+
+[Official GitHub](https://github.com/salesforce/progen/tree/main/progen2) of [ProGen2 by Nijkamp et al.](https://www.cell.com/cell-systems/fulltext/S2405-4712(23)00272-7).
+
+- The ProGen2 suite of protein language models are scaled to 6.4B parameters
+- Models with increased scale better capture the distribution of protein sequences
+- ProGen2 models generate novel protein sequences adopting natural folds
+- ProGen2 model likelihoods are effective for zero-shot fitness prediction
+
+```python
+import torch
+from faesm.progen2 import ProGenForCausalLM
+from transformers import AutoTokenizer
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+model = ProGenForCausalLM.from_pretrained("jinyuan22/ProGen2-xlarge").to(torch.float16).to(device).eval()
+tokenizer = AutoTokenizer.from_pretrained("jinyuan22/ProGen2-xlarge")
+
+# sequence = "1" + "ACDEFGHIKLMNPQRSTVWY" * 50 + "2" # 1002 token
+
+sequence = "2GFLPFRGADEGLAAREAATLAARGTAARAYREDSWAVPVPRGLLGDLTARVAALGAASPPPADPLAVTLDLHHVTAEVALTTVLDAATLVHGQTRVLSAEDAAEAATAAAAATEAYLERLQDFVLFMSASVRVWRRGNAAGATGPEWDQWYTVADRDALGSAPTHLAVLGRQADALCHFVLDRVAWGTCGTPLWSGDEDLGNVVATFAGYADRLATAPRDLIM1"
+
+inputs = tokenizer(sequence, return_tensors="pt").to(device)
+
+with torch.no_grad():
+  logits = model(inputs.input_ids, labels=inputs.input_ids).logits
+
+logits = logits[0][:-1, ...]
+target = inputs.input_ids[0, 1:]
+
+# remove unused logits
+first_token, last_token = 5, 29
+logits = logits[:, first_token:(last_token+1)]
+target = target - first_token
+
+ce_eval = torch.nn.functional.cross_entropy(input=logits.view(-1, logits.size(-1)), target=target.view(-1), reduction="mean").item()
+print(ce_eval)
+assert abs(ce_eval - 1.0) < 0.1
+```