Spaces:
Sleeping
Sleeping
File size: 2,197 Bytes
21c9bd3 | 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 | """
SERAPH model architecture definition and weight-loading helper.
"""
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from transformers import EsmModel, EsmTokenizer
from config import HF_REPO_ID, WEIGHTS_FILE, ESM_MODEL_ID
class SERAPH(nn.Module):
"""ESM-2 backbone + Conv1D + BiLSTM head for 3-state secondary
structure prediction (Helix / Sheet / Coil)."""
def __init__(self, esm_model, conv_channels=256, kernel_size=7,
lstm_hidden=256, num_classes=3, dropout=0.3, freeze_esm=True):
super().__init__()
self.esm = esm_model
if freeze_esm:
for param in self.esm.encoder.layer[:-2].parameters():
param.requires_grad = False
esm_embed_dim = self.esm.config.hidden_size
self.conv = nn.Conv1d(esm_embed_dim, conv_channels, kernel_size=kernel_size, padding=kernel_size // 2)
self.bn = nn.BatchNorm1d(conv_channels)
self.dropout = nn.Dropout(dropout)
self.bilstm = nn.LSTM(conv_channels, lstm_hidden, num_layers=2, batch_first=True, bidirectional=True)
self.fc = nn.Linear(lstm_hidden * 2, num_classes)
def forward(self, input_ids, attention_mask=None):
x = self.esm(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
x = x.transpose(1, 2)
x = torch.relu(self.bn(self.conv(x)))
x = self.dropout(x)
x = x.transpose(1, 2)
x, _ = self.bilstm(x)
x = self.dropout(x)
return self.fc(x)
def load_model():
"""Downloads the ESM-2 backbone + tokenizer and the fine-tuned SERAPH
checkpoint, and returns a ready-to-use (model, tokenizer) pair."""
print("Loading ESM2 backbone...")
esm = EsmModel.from_pretrained(ESM_MODEL_ID)
tokenizer = EsmTokenizer.from_pretrained(ESM_MODEL_ID)
print("Downloading SERAPH weights...")
weights_path = hf_hub_download(repo_id=HF_REPO_ID, filename=WEIGHTS_FILE)
checkpoint = torch.load(weights_path, map_location="cpu")
model = SERAPH(esm_model=esm)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()
print("SERAPH model ready.")
return model, tokenizer
|