update readme

README.md

@@ -40,21 +40,24 @@ SeqScreen computes cosine similarities between protein and molecule embeddings,
 
 ```python
 from transformers import AutoTokenizer, AutoModel
+from peft import PeftModel
 import torch
 
-# proteins
+# proteins: ESM2 + LoRA adapter
 tokenizer_prot = AutoTokenizer.from_pretrained('facebook/esm2_t36_3B_UR50D')
-encoder_prot = AutoModel.from_pretrained('facebook/esm2_t36_3B_UR50D').eval()
+backbone = AutoModel.from_pretrained(
+    'facebook/esm2_t36_3B_UR50D',
+    torch_dtype=torch.bfloat16
+)
+backbone = PeftModel.from_pretrained(backbone, 'SaeedLab/SeqScreen-lora').eval()
 
-proteins = ["MKTFFVLLL", "ABCDE"]
-proteins = [" ".join(i) for i in proteins]
+proteins = ["MKTFFVLLL", "ACDEFGHIKLM"]
 inputs_prot = tokenizer_prot(proteins, return_tensors="pt", padding=True)
 
 with torch.no_grad():
-  outputs = encoder_prot(**inputs_prot)
-  hidden = outputs.last_hidden_state[:, :]
-  mask = inputs_prot['attention_mask'].unsqueeze(-1).float()
-  prot_rep = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-8)
+    hidden = backbone(**inputs_prot).last_hidden_state
+    mask = inputs_prot['attention_mask'].unsqueeze(-1).float()
+    prot_emb = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-8)
 
 # molecules
 tokenizer_mol = AutoTokenizer.from_pretrained('SaeedLab/MolDeBERTa-base-123M-mlc')
@@ -64,16 +67,15 @@ molecules = ["NCCc1nc(-c2ccccc2)cs1", "CC(=O)OCC(C)C"]
 inputs_mol = tokenizer_mol(molecules, return_tensors="pt", padding=True)
 
 with torch.no_grad():
-  outputs = encoder_mol(**inputs_mol)
-  hidden = outputs.last_hidden_state[:, :]
-  mask = inputs_mol['attention_mask'].unsqueeze(-1).float()
-  mol_rep = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-8)
+    hidden = encoder_mol(**inputs_mol).last_hidden_state
+    mask = inputs_mol['attention_mask'].unsqueeze(-1).float()
+    mol_emb = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-8)
 
 # seqscreen
 seqscreen = AutoModel.from_pretrained('SaeedLab/SeqScreen-Finetuning', trust_remote_code=True).eval()
 
 with torch.no_grad():
-  outputs = seqscreen(prot=prot_rep, mol=mol_rep)
+    outputs = seqscreen(prot=prot_emb, mol=mol_emb)
 
 print('Protein embeddings projected:', outputs.prot_rep)
 print('Molecule embeddings projected:', outputs.mol_rep)