Views
No views yet
1import torch
2from transformers import AutoModel, AutoTokenizer, AutoConfig
3from esm.tokenization import get_esm3_model_tokenizers
4
5model_name = "area-science-park/ESM3-PPISites"
6model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
7
8tokenizers = get_esm3_model_tokenizers("esm3_sm_open_v1")
9sequence_tokenizer = tokenizers.sequence
10
11# move model to device
12device = "cuda" if torch.cuda.is_available() else "cpu"
13model = model.to(device)
14
15# run over a sample sequence
16sequence = "MKTVRQERLKSIVRILEAAKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"
17# tokenize
18tokens = sequence_tokenizer.encode(sequence)
19inputs = torch.tensor(tokens).unsqueeze(0).to(device)
20# inference
21logits = model(inputs)["logits"]
22probabilities = torch.sigmoid(logits)
23
24# get predictions
25probabilities