update readme

README.md

@@ -39,9 +39,52 @@ SeqScreen computes cosine similarities between protein and molecule embeddings,
 ### Similarity
 
 ```python
-# code here
+from transformers import AutoTokenizer, AutoModel
+import torch
+
+# proteins
+tokenizer_prot = AutoTokenizer.from_pretrained('facebook/esm2_t36_3B_UR50D')
+encoder_prot = AutoModel.from_pretrained('facebook/esm2_t36_3B_UR50D').eval()
+
+proteins = ["MKTFFVLLL", "ABCDE"]
+proteins = [" ".join(i) for i in proteins]
+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)
+
+# molecules
+tokenizer_mol = AutoTokenizer.from_pretrained('SaeedLab/MolDeBERTa-base-123M-mlc')
+encoder_mol = AutoModel.from_pretrained('SaeedLab/MolDeBERTa-base-123M-mlc').eval()
+
+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)
+
+# seqscreen
+seqscreen = AutoModel.from_pretrained('SaeedLab/SeqScreen-Frozen', trust_remote_code=True).eval()
+
+with torch.no_grad():
+  outputs = seqscreen(prot=prot_rep, mol=mol_rep)
+
+print('Protein embeddings projected:', outputs.prot_rep)
+print('Molecule embeddings projected:', outputs.mol_rep)
+print('Cossine similarity:', outputs.similarity)
+
 ```
 
+The returned outputs are:
+- prot_rep: Projected embeddings for protein input in 512 dimension.
+- mol_rep: Projected embeddings for molecule input in 512 dimension.
+- similarity: Cossine similarity between proteins and molecules.
 
 ## Citation