Update classifier_code/nonfouling_wt.py

classifier_code/nonfouling_wt.py

@@ -1,98 +1,78 @@
-import sys
-import os
+import numpy as np
 import xgboost as xgb
 import torch
-import numpy as np
-import warnings
-import numpy as np
-from rdkit import Chem, rdBase, DataStructs
-from transformers import AutoTokenizer, EsmModel
+from transformers import EsmModel, AutoTokenizer
+import torch.nn as nn
+import pdb
 
-rdBase.DisableLog('rdApp.error')
-warnings.filterwarnings("ignore", category=DeprecationWarning)
-warnings.filterwarnings("ignore", category=UserWarning)
-warnings.filterwarnings("ignore", category=FutureWarning)
+# ======================== MLP =========================================
+# Still need mean pooling along lengths
+class MaskedMeanPool(nn.Module):
+    def forward(self, X, M):  # X: (B,L,H), M: (B,L)
+        Mf = M.unsqueeze(-1).float()
+        denom = Mf.sum(dim=1).clamp(min=1.0)
+        return (X * Mf).sum(dim=1) / denom  # (B,H)
 
-class Nonfouling:
-    def __init__(self):
-        # change model path
-        self.predictor = xgb.Booster(model_file='../classifier_ckpt/best_model_nonfouling.json')
-        
-        # Load ESM model and tokenizer
-        self.tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
-        self.model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D")
-        self.model.eval()
+class MLPClassifier(nn.Module):
+    def __init__(self, in_dim, hidden=512, dropout=0.1):
+        super().__init__()
+        self.pool = MaskedMeanPool()
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, hidden),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden, 1),
+        )
+    def forward(self, X, M):
+        z = self.pool(X, M)
+        return self.net(z).squeeze(-1)  # logits
+# ======================== MLP =========================================
+
+
+class NonfoulingModel:
+    def __init__(self, device):
+        ckpt = torch.load('../classifier_ckpt/wt_nonfouling.pt', weights_only=False, map_location=device)
+        best_params = ckpt["best_params"]
+        self.predictor = MLPClassifier(in_dim=1280, hidden=int(best_params["hidden"]), dropout=float(best_params.get("dropout", 0.1)))
+        self.predictor.load_state_dict(ckpt["state_dict"])
+        self.predictor = self.predictor.to(device)
+        self.predictor.eval()
+
+        self.model = EsmModel.from_pretrained("facebook/esm2_t33_650M_UR50D").to(device)
+        # self.model.eval()
+
+        self.device = device
     
-    def generate_embeddings(self, sequences):
+    def generate_embeddings(self, input_ids, attention_mask):
         """Generate ESM embeddings for protein sequences"""
-        embeddings = []
-        
-        # Process sequences in batches to avoid memory issues
-        batch_size = 8
-        for i in range(0, len(sequences), batch_size):
-            batch_sequences = sequences[i:i + batch_size]
-            
-            inputs = self.tokenizer(
-                batch_sequences, 
-                padding=True, 
-                truncation=True, 
-                return_tensors="pt"
-            )
-            
-            if torch.cuda.is_available():
-                inputs = {k: v.cuda() for k, v in inputs.items()}
-                self.model = self.model.cuda()
-            
-            # Generate embeddings
-            with torch.no_grad():
-                outputs = self.model(**inputs)
-                
-                # Get last hidden states
-                last_hidden_states = outputs.last_hidden_state
-                
-                # Compute mean pooling (excluding padding tokens)
-                attention_mask = inputs['attention_mask'].unsqueeze(-1)
-                masked_hidden_states = last_hidden_states * attention_mask
-                sum_hidden_states = masked_hidden_states.sum(dim=1)
-                seq_lengths = attention_mask.sum(dim=1)
-                batch_embeddings = sum_hidden_states / seq_lengths
-                
-                batch_embeddings = batch_embeddings.cpu().numpy()
-                embeddings.append(batch_embeddings)
-        
-        if embeddings:
-            return np.vstack(embeddings)
-        else:
-            return np.array([])
-    
-    def get_scores(self, input_seqs: list):
-        scores = np.zeros(len(input_seqs))
-        features = self.generate_embeddings(input_seqs)
+        with torch.no_grad():
+            embeddings = self.model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
         
-        if len(features) == 0:
-            return scores
-        
-        features = np.nan_to_num(features, nan=0.)
-        features = np.clip(features, np.finfo(np.float32).min, np.finfo(np.float32).max)
-        
-        features = xgb.DMatrix(features)
-        
-        scores = self.predictor.predict(features)
-        return scores
+        return embeddings
     
-    def __call__(self, input_seqs: list):
-        scores = self.get_scores(input_seqs)
-        return scores
+    def get_scores(self, input_ids, attention_mask):
+        features = self.generate_embeddings(input_ids, attention_mask)
+
+        keep = (input_ids != 0) & (input_ids != 1) & (input_ids != 2)
+        attention_mask[keep==False] = 0
+        scores = self.predictor(features, attention_mask)
+        return scores.detach().cpu().numpy()
     
+    def __call__(self, input_ids, attention_mask):
+        scores = self.get_scores(input_ids, attention_mask)
+        return 1.0 / (1.0 + np.exp(-scores))
+
+
 def unittest():
-    nonfouling = Nonfouling()
-    sequences = [
-        "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG",
-        "MSEGIRQAFVLAKSIWPARVARFTVDNRIRSLVKTYEAIKVDPYNPAFLEVLD"
-    ]    
-    
-    scores = nonfouling(input_seqs=sequences)
+    device = 'cuda:0'
+    nf = NonfoulingModel(device=device)
+    seq = ["HAIYPRH", "HAEGTFTSDVSSYLEGQAAKEFIAWLVKGR"]
+
+    tokenizer = AutoTokenizer.from_pretrained('facebook/esm2_t33_650M_UR50D')
+    seq_tokens = tokenizer(seq, padding=True, return_tensors='pt').to(device)
+    scores = nf(**seq_tokens)
     print(scores)
-    
+
+
 if __name__ == '__main__':
     unittest()