Spaces:
Sleeping
Sleeping
| """ | |
| 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 | |