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