update files

README.md

@@ -70,7 +70,7 @@ with torch.no_grad():
   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()
+seqscreen = AutoModel.from_pretrained('SaeedLab/SeqScreen-Finetuning', trust_remote_code=True).eval()
 
 with torch.no_grad():
   outputs = seqscreen(prot=prot_rep, mol=mol_rep)

config.json

@@ -0,0 +1,18 @@
+{
+  "architectures": [
+    "SeqScreenModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_seqscreen.SeqScreenConfig",
+    "AutoModel": "modeling_seqscreen.SeqScreenModel"
+  },
+  "dropout": 0.1,
+  "dtype": "float32",
+  "esm2_model_name": "facebook/esm2_t36_3B_UR50D",
+  "lora_adapter_repo": "SaeedLab/SeqScreen-lora",
+  "model_type": "seqscreen",
+  "mol_dim": 768,
+  "proj_dim": 512,
+  "prot_dim": 2560,
+  "transformers_version": "4.57.3"
+}

configuration_seqscreen.py

@@ -0,0 +1,23 @@
+from transformers import PretrainedConfig
+
+
+class SeqScreenConfig(PretrainedConfig):
+  model_type = "seqscreen"
+
+  def __init__(
+    self,
+    prot_dim: int = 2560,
+    mol_dim: int = 768,
+    proj_dim: int = 512,
+    dropout: float = 0.1,
+    esm2_model_name: str = "facebook/esm2_t36_3B_UR50D",
+    lora_adapter_repo: str = None,
+    **kwargs,
+  ):
+    super().__init__(**kwargs)
+    self.prot_dim = prot_dim
+    self.mol_dim = mol_dim
+    self.proj_dim = proj_dim
+    self.dropout = dropout
+    self.esm2_model_name = esm2_model_name
+    self.lora_adapter_repo = lora_adapter_repo

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c4a55beefcc43242daf84849856898dcf03039d21a5849c1878e7a2edc05042
+size 8930448

modeling_seqscreen.py

@@ -0,0 +1,54 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import dataclass
+import torch
+from transformers.utils import ModelOutput
+from transformers import PreTrainedModel
+
+from .configuration_seqscreen import SeqScreenConfig
+
+@dataclass
+class SeqScreenModelOutput(ModelOutput):
+  prot_rep: torch.FloatTensor = None
+  mol_rep: torch.FloatTensor = None
+  similarity: torch.FloatTensor = None
+
+class ProjectionLayer(nn.Module):
+  def __init__(self, in_dim, out_dim, dropout):
+    super().__init__()
+    self.projection = nn.Sequential(
+      nn.Linear(in_dim, out_dim),
+      nn.LayerNorm(out_dim),
+      nn.GELU(),
+      nn.Dropout(dropout),
+      nn.Linear(out_dim, out_dim)
+    )
+
+  def forward(self, x):
+    x = self.projection(x)
+    return F.normalize(x, dim=-1)
+
+
+class SeqScreenModel(PreTrainedModel):
+  config_class = SeqScreenConfig
+  base_model_prefix = "seqscreen"
+
+  def __init__(self, config: SeqScreenConfig):
+    super().__init__(config)
+
+    self.proj_prot = ProjectionLayer(config.prot_dim, config.proj_dim, dropout=config.dropout)
+    self.proj_mol = ProjectionLayer(config.mol_dim, config.proj_dim, dropout=config.dropout)
+
+    self.post_init()
+
+  def forward(self, prot: torch.Tensor, mol: torch.Tensor):
+    prot_rep = self.proj_prot(prot)
+    mol_rep = self.proj_mol(mol)
+    similarity = prot_rep @ mol_rep.T
+
+    return SeqScreenModelOutput(
+      prot_rep=prot_rep,
+      mol_rep=mol_rep,
+      similarity=similarity
+    )