Views
No views yet
logits, features = model(epitope_ids, hla_ids), then apply softmax to get the binding probability.pip install torch transformers peft1import torch
2import torch.nn.functional as F
3from transformers import AutoModel, AutoTokenizer
4model_id = "SkywalkerLu/TriStageHLA-BIND"
5model = AutoModel.from_pretrained(model_id, trust_remote_code=True).to(device).eval()## How to use TriStageHLA-BIND
```python
import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer
# Device
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load model (replace with your model id if different)
model_id = "SkywalkerLu/TriStageHLA-BIND"
model = AutoModel.from_pretrained(model_id, trust_remote_code=True).to(device).eval()
# Load tokenizer used in training (ESM2 650M)
tok = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
# Example inputs
peptide = "GILGFVFTL" # 9-mer example
# Fake placeholder pseudosequence for demo; replace with a real one from your mapping/data
hla_pseudoseq = (
"YYSEYRNIYAQTDESNLYLSYDYYTWAERAYEWY"
)
# Fixed lengths (must match training)
PEP_LEN = 16
HLA_LEN = 36
PAD_ID = tok.pad_token_id if tok.pad_token_id is not None else 1
def pad_to_len(ids_list, target_len, pad_id):
return ids_list + [pad_id] * (target_len - len(ids_list)) if len(ids_list) < target_len else ids_list[:target_len]
# Tokenize
pep_ids = tok(peptide, add_special_tokens=True)["input_ids"]
hla_ids = tok(hla_pseudoseq, add_special_tokens=True)["input_ids"]
# Pad/truncate
pep_ids = pad_to_len(pep_ids, PEP_LEN, PAD_ID)
hla_ids = pad_to_len(hla_ids, HLA_LEN, PAD_ID)
# Tensors (batch=1)
pep_tensor = torch.tensor([pep_ids], dtype=torch.long, device=device)
hla_tensor = torch.tensor([hla_ids], dtype=torch.long, device=device)
# Forward + probability
with torch.no_grad():
logits, features = model(pep_tensor, hla_tensor)
prob_bind = F.softmax(logits, dim=1)[0, 1].item()
pred = int(prob_bind >= 0.5)
print({"peptide": peptide, "bind_prob": round(prob_bind, 6), "label": pred})1import torch
2import torch.nn.functional as F
3from transformers import AutoModel, AutoTokenizer
4
5# Device
6device = "cuda" if torch.cuda.is_available() else "cpu"
7
8# Load model and tokenizer
9model_id = "SkywalkerLu/TriStageHLA-BIND" # replace with your model id if different
10model = AutoModel.from_pretrained(model_id, trust_remote_code=True).to(device).eval()
11tok = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
12
13# Fixed lengths (must match training)
14PEP_LEN = 16
15HLA_LEN = 36
16PAD_ID = tok.pad_token_id if tok.pad_token_id is not None else 1
17
18def pad_to_len(ids_list, target_len, pad_id):
19 return ids_list + [pad_id] * (target_len - len(ids_list)) if len(ids_list) < target_len else ids_list[:target_len]
20
21# Example batch (use real HLA pseudosequences in your data)
22batch = [
23 {"peptide": "GILGFVFTL", "hla_pseudo": "YYSEYRNIYAQTDESNLYLSYDYYTWAERAYEWY"},
24 {"peptide": "NLVPMVATV", "hla_pseudo": "YYSEYRNIYAQTDESNLYLSYDYYTWAERAYEWY"},
25 {"peptide": "SIINFEKL", "hla_pseudo": "YYSEYRNIYAQTDESNLYLSYDYYTWAERAYEWY"},
26]
27
28# Tokenize and pad/truncate
29pep_ids_batch, hla_ids_batch = [], []
30for item in batch:
31 pep_ids = tok(item["peptide"], add_special_tokens=True)["input_ids"]
32 hla_ids = tok(item["hla_pseudo"], add_special_tokens=True)["input_ids"]
33 pep_ids_batch.append(pad_to_len(pep_ids, PEP_LEN, PAD_ID))
34 hla_ids_batch.append(pad_to_len(hla_ids, HLA_LEN, PAD_ID))
35
36# To tensors
37pep_tensor = torch.tensor(pep_ids_batch, dtype=torch.long, device=device) # [B, PEP_LEN]
38hla_tensor = torch.tensor(hla_ids_batch, dtype=torch.long, device=device) # [B, HLA_LEN]
39
40# Forward
41with torch.no_grad():
42 logits, _ = model(pep_tensor, hla_tensor) # logits shape: [B, 2]
43 probs = F.softmax(logits, dim=1)[:, 1] # binding probability for class-1
44
45# Threshold to labels (0/1)
46labels = (probs >= 0.5).long().tolist()
47
48# Print results
49for i, item in enumerate(batch):
50 print({
51 "peptide": item["peptide"],
52 "bind_prob": float(probs[i].item()),
53 "label": labels[i]
54 })