Update README.md

README.md

@@ -66,3 +66,62 @@ with torch.no_grad():
 print({"peptide": peptide, "bind_prob": round(prob_bind, 6), "label": pred})
 ```
 
+
+## Batch Inference (Python)
+
+```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 and tokenizer
+model_id = "SkywalkerLu/TransHLA2.0"  # replace with your model id if different
+model = AutoModel.from_pretrained(model_id, trust_remote_code=True).to(device).eval()
+tok = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
+
+# 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]
+
+# Example batch (use real HLA pseudosequences in your data)
+batch = [
+    {"peptide": "GILGFVFTL", "hla_pseudo": "YOUR_HLA_PSEUDOSEQ_1"},
+    {"peptide": "NLVPMVATV", "hla_pseudo": "YOUR_HLA_PSEUDOSEQ_2"},
+    {"peptide": "SIINFEKL",  "hla_pseudo": "YOUR_HLA_PSEUDOSEQ_3"},
+]
+
+# Tokenize and pad/truncate
+pep_ids_batch, hla_ids_batch = [], []
+for item in batch:
+    pep_ids = tok(item["peptide"], add_special_tokens=True)["input_ids"]
+    hla_ids = tok(item["hla_pseudo"], add_special_tokens=True)["input_ids"]
+    pep_ids_batch.append(pad_to_len(pep_ids, PEP_LEN, PAD_ID))
+    hla_ids_batch.append(pad_to_len(hla_ids, HLA_LEN, PAD_ID))
+
+# To tensors
+pep_tensor = torch.tensor(pep_ids_batch, dtype=torch.long, device=device)  # [B, PEP_LEN]
+hla_tensor = torch.tensor(hla_ids_batch, dtype=torch.long, device=device)  # [B, HLA_LEN]
+
+# Forward
+with torch.no_grad():
+    logits, _ = model(pep_tensor, hla_tensor)   # logits shape: [B, 2]
+    probs = F.softmax(logits, dim=1)[:, 1]      # binding probability for class-1
+
+# Threshold to labels (0/1)
+labels = (probs >= 0.5).long().tolist()
+
+# Print results
+for i, item in enumerate(batch):
+    print({
+        "peptide": item["peptide"],
+        "bind_prob": float(probs[i].item()),
+        "label": labels[i]
+    })
+```