SaeedLab/BindScreen-Finetuning-LIT_PCBA

1import torch2import torch.nn as nn3import torch.nn.functional as F4from dataclasses import dataclass5import torch6from transformers.utils import ModelOutput7from transformers import PreTrainedModel8 9from .configuration_seqscreen import SeqScreenConfig10 11@dataclass12class SeqScreenModelOutput(ModelOutput):13  prot_rep: torch.FloatTensor = None14  mol_rep: torch.FloatTensor = None15  similarity: torch.FloatTensor = None16 17class ProjectionLayer(nn.Module):18  def __init__(self, in_dim, out_dim, dropout):19    super().__init__()20    self.projection = nn.Sequential(21      nn.Linear(in_dim, out_dim),22      nn.LayerNorm(out_dim),23      nn.GELU(),24      nn.Dropout(dropout),25      nn.Linear(out_dim, out_dim)26    )27 28  def forward(self, x):29    x = self.projection(x)30    return F.normalize(x, dim=-1)31 32 33class SeqScreenModel(PreTrainedModel):34  config_class = SeqScreenConfig35  base_model_prefix = "seqscreen"36 37  def __init__(self, config: SeqScreenConfig):38    super().__init__(config)39 40    self.proj_prot = ProjectionLayer(config.prot_dim, config.proj_dim, dropout=config.dropout)41    self.proj_mol = ProjectionLayer(config.mol_dim, config.proj_dim, dropout=config.dropout)42 43    self.post_init()44 45  def forward(self, prot: torch.Tensor, mol: torch.Tensor):46    prot_rep = self.proj_prot(prot)47    mol_rep = self.proj_mol(mol)48    similarity = prot_rep @ mol_rep.T49 50    return SeqScreenModelOutput(51      prot_rep=prot_rep,52      mol_rep=mol_rep,53      similarity=similarity54    )55