| --- |
| 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-small").to(torch.float16).to(device).eval() |
| tokenizer = AutoTokenizer.from_pretrained("jinyuan22/ProGen2-small") |
| |
| # 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 - 2.4) < 0.1 |
| ``` |
