Spaces:
Sleeping
Sleeping
File size: 2,741 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 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """
Loads the SERAPH model once at import time and exposes the prediction
and preset-change handlers used by the Gradio UI.
"""
import spaces
import torch
from config import IDX_TO_LABEL
from dataset import PROTEIN_MAP
from html_builder import build_result_html
from model import load_model
model, tokenizer = load_model()
@spaces.GPU
def predict_structure(sequence_input: str, selected_protein_name: str) -> tuple:
"""Runs ESM-2 + BiLSTM inference and builds the HTML visualization."""
sequence = sequence_input.upper().strip()
# Sanitize inputs to amino acid alphabet
valid_aas = set("ACDEFGHIKLMNPQRSTVWY")
sequence = "".join([c for c in sequence if c in valid_aas])
if not sequence:
empty_html = """
<div style='padding: 24px; text-align: center; background: rgba(255,255,255,0.02); border: 1px solid rgba(255,255,255,0.08); border-radius: 16px;'>
<p style='color: #8e8e93; font-family: "JetBrains Mono", monospace;'>Please enter a valid amino acid sequence (e.g., M K W V T F I S L L L L F S S A)...</p>
</div>
"""
return empty_html, "", ""
# Model inference
tokens = tokenizer(sequence, return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
output = model(input_ids=tokens["input_ids"], attention_mask=tokens["attention_mask"])
preds = output.argmax(dim=-1)[0]
# Extract predicted structure labels
pred_labels = [IDX_TO_LABEL[p.item()] for p in preds[1:-1]]
# Handle truncation alignment safely
pred_str = "".join(pred_labels[:len(sequence)])
if len(pred_str) < len(sequence):
pred_str += "C" * (len(sequence) - len(pred_str))
# Match with dataset ground truth if selected
meta = PROTEIN_MAP.get(selected_protein_name, None)
has_truth = False
true_ss = ""
description = ""
fun_fact = ""
if meta and meta["sequence"].upper() == sequence:
has_truth = True
true_ss = meta["true_ss"]
description = meta["description"]
fun_fact = meta["fun_fact"]
# Calculate Q3 Metric if Ground Truth matches
q3_score = None
if has_truth and len(true_ss) == len(sequence):
matches = sum(1 for i in range(len(sequence)) if pred_str[i] == true_ss[i])
q3_score = (matches / len(sequence)) * 100.0
html_output = build_result_html(sequence, pred_str, true_ss if has_truth else None, q3_score, description, fun_fact)
return html_output, sequence, pred_str
def on_preset_change(selected_name):
"""Populates the sequence textbox when a preset dropdown item is chosen."""
if selected_name in PROTEIN_MAP:
item = PROTEIN_MAP[selected_name]
return item["sequence"]
return ""
|