Spaces:
Runtime error
Runtime error
File size: 3,224 Bytes
ce95901 0e50f5d d253cf5 937669e ce95901 0e50f5d 937669e 0e50f5d 937669e 0e50f5d ce95901 0e50f5d 937669e d253cf5 937669e d253cf5 937669e d253cf5 937669e 0e50f5d ce95901 937669e d253cf5 937669e d253cf5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import torch
import torch.nn as nn
import torch.nn.functional as F
import esm
import joblib # ✅ added
class BiTransformerAttentionModel(nn.Module):
def __init__(self, input_dim, seq_len=8, embed_dim=512, n_heads=8, ff_hidden=256, dropout=0.2):
super().__init__()
assert input_dim % seq_len == 0
self.seq_len = seq_len
self.token_dim = input_dim // seq_len
self.fc_in = nn.Linear(self.token_dim, embed_dim)
enc = lambda: nn.TransformerEncoderLayer(
d_model=embed_dim, nhead=n_heads,
dim_feedforward=ff_hidden, dropout=dropout,
batch_first=True
)
self.fwd = nn.TransformerEncoder(enc(), num_layers=2)
self.bwd = nn.TransformerEncoder(enc(), num_layers=2)
self.attn = nn.Linear(embed_dim*2,1)
self.ln = nn.LayerNorm(embed_dim*2)
self.fc = nn.Sequential(
nn.Linear(embed_dim*2, ff_hidden),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(ff_hidden, ff_hidden//2),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(ff_hidden//2,1)
)
self.pos = nn.Parameter(torch.randn(1, seq_len, embed_dim)*0.01)
def forward(self,x):
x = x.view(x.size(0), self.seq_len, self.token_dim)
x = self.fc_in(x)
f = self.fwd(x+self.pos)
b = self.bwd(torch.flip(x+self.pos,[1]))
x = torch.cat([f,b],dim=-1)
x = self.ln(x)
w = torch.softmax(self.attn(x),dim=1)
x = (w*x).sum(dim=1)
return self.fc(x).squeeze()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# ✅ Load checkpoint
checkpoint = torch.load("BiTransformer_BEST.pt", map_location=device)
model = BiTransformerAttentionModel(input_dim=checkpoint["input_dim"])
model.load_state_dict(checkpoint["model_state"])
model.to(device)
model.eval()
# ✅ Load scaler
scaler = joblib.load("scaler.pkl")
# -----------------------------
# ESM embedding
# -----------------------------
esm_model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
esm_model = esm_model.eval().to(device)
batch_converter = alphabet.get_batch_converter()
def sequence_to_embedding(sequences):
batch_labels, batch_strs, batch_tokens = batch_converter(sequences)
batch_tokens = batch_tokens.to(device)
with torch.no_grad():
results = esm_model(batch_tokens, repr_layers=[33], return_contacts=False)
token_reps = results["representations"][33]
mask = batch_tokens != alphabet.padding_idx
embeddings = (token_reps * mask.unsqueeze(-1)).sum(1) / mask.sum(1, keepdim=True)
return embeddings
# -----------------------------
# Prediction
# -----------------------------
def predict(sequence: str):
seq_tuple = [("seq", sequence)]
X_embed = sequence_to_embedding(seq_tuple)
# match input size
X_embed = X_embed[:, :checkpoint["input_dim"]]
# ✅ APPLY SCALER (CRITICAL FIX)
X_embed = scaler.transform(X_embed.cpu().numpy())
X_embed = torch.tensor(X_embed, dtype=torch.float32).to(device)
model.eval()
with torch.no_grad():
logits = model(X_embed)
prob = torch.sigmoid(logits)
return float(prob.cpu().item())
|