Update app.py

app.py

@@ -5,184 +5,102 @@ import re
 import zipfile
 import tempfile
 from dataclasses import dataclass
-from typing import Dict, List, Tuple, Optional
+from typing import Dict, List, Optional
 
 import gradio as gr
-import numpy as np
-import pandas as pd
 import torch
+from transformers import AutoModel, AutoTokenizer, T5EncoderModel, T5Tokenizer
 
-from transformers import (
-    AutoModel,
-    AutoTokenizer,
-    T5EncoderModel,
-    T5Tokenizer,
-)
-
-# =========================
-# Global config
-# =========================
-APP_TITLE = "Protein Embedding Hub"
-APP_DESC = """
-Input FASTA protein sequences, choose a model, and export residue-level embeddings with shape L*d.
-This app automatically normalizes model outputs such as L+1, L+2, or tokenized variants back to strict residue-level L*d.
-"""
+APP_TITLE = "Protein Embedding"
 
 ALLOWED_AA = set(list("ACDEFGHIKLMNPQRSTVWYXBZJUO"))
 REPLACE_WITH_X = set(list("UZOB"))
 
-# =========================
-# Model registry
-# =========================
+
 @dataclass
 class ModelSpec:
     name: str
-    family: str         # "hf_encoder", "t5_encoder", "esmc"
+    family: str
     model_id: str
     tokenizer_id: Optional[str] = None
-    note: str = ""
 
 
 MODEL_SPECS: Dict[str, ModelSpec] = {
-    # ESM2
     "ESM2-8M": ModelSpec(
         name="ESM2-8M",
         family="hf_encoder",
         model_id="facebook/esm2_t6_8M_UR50D",
         tokenizer_id="facebook/esm2_t6_8M_UR50D",
-        note="Very light."
     ),
     "ESM2-35M": ModelSpec(
         name="ESM2-35M",
         family="hf_encoder",
         model_id="facebook/esm2_t12_35M_UR50D",
         tokenizer_id="facebook/esm2_t12_35M_UR50D",
-        note="Good small baseline."
     ),
     "ESM2-150M": ModelSpec(
         name="ESM2-150M",
         family="hf_encoder",
         model_id="facebook/esm2_t30_150M_UR50D",
         tokenizer_id="facebook/esm2_t30_150M_UR50D",
-        note="Balanced."
     ),
     "ESM2-650M": ModelSpec(
         name="ESM2-650M",
         family="hf_encoder",
         model_id="facebook/esm2_t33_650M_UR50D",
         tokenizer_id="facebook/esm2_t33_650M_UR50D",
-        note="Strong sequence-only baseline."
     ),
-
-    # ESMC
     "ESMC-300M": ModelSpec(
         name="ESMC-300M",
         family="esmc",
         model_id="esmc_300m",
-        note="Representation model; usually better efficiency/performance than similar-size ESM2."
     ),
     "ESMC-600M": ModelSpec(
         name="ESMC-600M",
         family="esmc",
         model_id="esmc_600m",
-        note="Larger ESMC."
     ),
-
-    # Ankh
     "Ankh-Base": ModelSpec(
         name="Ankh-Base",
         family="hf_encoder",
         model_id="ElnaggarLab/ankh-base",
         tokenizer_id="ElnaggarLab/ankh-base",
-        note="Efficient strong general-purpose protein LM."
     ),
     "Ankh-Large": ModelSpec(
         name="Ankh-Large",
         family="hf_encoder",
         model_id="ElnaggarLab/ankh-large",
         tokenizer_id="ElnaggarLab/ankh-large",
-        note="Larger Ankh variant."
     ),
-
-    # ProtT5 encoder
     "ProtT5-XL-Encoder": ModelSpec(
         name="ProtT5-XL-Encoder",
         family="t5_encoder",
         model_id="Rostlab/prot_t5_xl_half_uniref50-enc",
         tokenizer_id="Rostlab/prot_t5_xl_half_uniref50-enc",
-        note="Classic protein embedding model; heavy."
+    ),
+    "ProSST-2048": ModelSpec(
+        name="ProSST-2048",
+        family="prosst",
+        model_id="AI4Protein/ProSST-2048",
+        tokenizer_id="AI4Protein/ProSST-2048",
     ),
 }
 
 
-# =========================
-# Model manager
-# =========================
-class ModelManager:
-    def __init__(self):
-        self.current_key = None
-        self.current_family = None
-        self.model = None
-        self.tokenizer = None
-        self.device = None
-
-    def unload(self):
-        self.model = None
-        self.tokenizer = None
-        self.current_key = None
-        self.current_family = None
-        self.device = None
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-    def load(self, model_key: str, device: str):
-        if self.current_key == model_key and self.device == device and self.model is not None:
-            return
-
-        self.unload()
-
-        spec = MODEL_SPECS[model_key]
-        resolved_device = _resolve_device(device)
-
-        if spec.family == "hf_encoder":
-            self.tokenizer = AutoTokenizer.from_pretrained(spec.tokenizer_id)
-            self.model = AutoModel.from_pretrained(spec.model_id)
-            self.model.to(resolved_device)
-            self.model.eval()
-
-        elif spec.family == "t5_encoder":
-            self.tokenizer = T5Tokenizer.from_pretrained(spec.tokenizer_id, do_lower_case=False)
-            self.model = T5EncoderModel.from_pretrained(spec.model_id)
-            self.model.to(resolved_device)
-            self.model.eval()
-
-        elif spec.family == "esmc":
-            try:
-                from esm.models.esmc import ESMC
-            except Exception as e:
-                raise RuntimeError(
-                    "Failed to import ESMC. Please install the official `esm` package. "
-                    f"Original error: {e}"
-                )
-            self.model = ESMC.from_pretrained(spec.model_id).to(resolved_device)
-            self.model.eval()
-            self.tokenizer = None
-
-        else:
-            raise ValueError(f"Unsupported family: {spec.family}")
-
-        self.current_key = model_key
-        self.current_family = spec.family
-        self.device = resolved_device
+def resolve_device(device: str) -> str:
+    if device == "auto":
+        return "cuda" if torch.cuda.is_available() else "cpu"
+    if device == "cuda" and not torch.cuda.is_available():
+        return "cpu"
+    return device
 
 
-MODEL_MANAGER = ModelManager()
+def safe_filename(x: str) -> str:
+    x = re.sub(r"[^A-Za-z0-9._-]+", "_", x)
+    x = x.strip("._")
+    return x or "item"
 
 
-# =========================
-# FASTA and sequence utils
-# =========================
 def parse_fasta(text: str) -> List[Dict[str, str]]:
     text = text.strip()
     if not text:
@@ -220,10 +138,8 @@ def parse_fasta(text: str) -> List[Dict[str, str]]:
     return records
 
 
-def clean_sequence(seq: str) -> Tuple[str, List[str]]:
+def clean_sequence(seq: str) -> str:
     seq = re.sub(r"\s+", "", seq).upper()
-    warnings = []
-
     if not seq:
         raise ValueError("Empty sequence after cleaning.")
 
@@ -231,70 +147,35 @@ def clean_sequence(seq: str) -> Tuple[str, List[str]]:
     if bad:
         raise ValueError(f"Invalid amino acid letters found: {bad}")
 
-    replaced = sorted({c for c in seq if c in REPLACE_WITH_X})
-    if replaced:
-        for c in replaced:
-            seq = seq.replace(c, "X")
-        warnings.append(f"Replaced uncommon residues {replaced} with X.")
-
-    return seq, warnings
+    for c in REPLACE_WITH_X:
+        seq = seq.replace(c, "X")
+    return seq
 
 
 def protein_to_spaced(seq: str) -> str:
     return " ".join(list(seq))
 
 
-def safe_filename(x: str) -> str:
-    x = re.sub(r"[^A-Za-z0-9._-]+", "_", x)
-    x = x.strip("._")
-    return x or "sequence"
-
-
-def _resolve_device(device: str) -> str:
-    if device == "auto":
-        return "cuda" if torch.cuda.is_available() else "cpu"
-    if device == "cuda" and not torch.cuda.is_available():
-        return "cpu"
-    return device
-
-
-# =========================
-# Embedding normalization
-# =========================
-def normalize_to_residue_level(
+def normalize_to_Ld(
     hidden: torch.Tensor,
     expected_len: int,
     special_tokens_mask: Optional[torch.Tensor] = None,
     attention_mask: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
-    """
-    Convert model output to strict residue-level shape [L, d].
-
-    Priority:
-    1) If special_tokens_mask exists, remove special tokens.
-    2) If already exactly L, keep.
-    3) If L+2, assume BOS/EOS and slice [1:-1].
-    4) If L+1, trim one token from the end.
-    5) Else crop to first L after best effort.
-    """
     if hidden.ndim != 2:
-        raise ValueError(f"Expected hidden shape [T, d], got {tuple(hidden.shape)}")
+        raise ValueError(f"Expected [T, d], got {tuple(hidden.shape)}")
 
-    T, d = hidden.shape
+    T = hidden.shape[0]
 
     if special_tokens_mask is not None:
-        mask = special_tokens_mask.bool().view(-1)
+        keep = ~special_tokens_mask.bool().view(-1)
         if attention_mask is not None:
-            attn = attention_mask.bool().view(-1)
-            keep = (~mask) & attn
-        else:
-            keep = ~mask
-        if keep.numel() == T:
-            filtered = hidden[keep]
-            if filtered.shape[0] == expected_len:
-                return filtered
-            if filtered.shape[0] > expected_len:
-                return filtered[:expected_len]
+            keep = keep & attention_mask.bool().view(-1)
+        filtered = hidden[keep]
+        if filtered.shape[0] == expected_len:
+            return filtered
+        if filtered.shape[0] > expected_len:
+            return filtered[:expected_len]
 
     if T == expected_len:
         return hidden
@@ -305,247 +186,270 @@ def normalize_to_residue_level(
     if T > expected_len:
         return hidden[:expected_len]
 
-    raise ValueError(
-        f"Could not normalize token length {T} to residue length {expected_len}."
-    )
+    raise ValueError(f"Cannot normalize token length {T} to residue length {expected_len}.")
+
+
+class SingleModelRunner:
+    def __init__(self):
+        self.model_key = None
+        self.family = None
+        self.device = None
+        self.model = None
+        self.tokenizer = None
+        self.sst_predictor = None
+
+    def unload(self):
+        self.model_key = None
+        self.family = None
+        self.device = None
+        self.model = None
+        self.tokenizer = None
+        self.sst_predictor = None
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    def load(self, model_key: str, device: str):
+        target_device = resolve_device(device)
+        if self.model_key == model_key and self.device == target_device and self.model is not None:
+            return
+
+        self.unload()
+        spec = MODEL_SPECS[model_key]
+
+        if spec.family == "hf_encoder":
+            self.tokenizer = AutoTokenizer.from_pretrained(spec.tokenizer_id)
+            self.model = AutoModel.from_pretrained(spec.model_id)
+            self.model.to(target_device)
+            self.model.eval()
+
+        elif spec.family == "t5_encoder":
+            self.tokenizer = T5Tokenizer.from_pretrained(spec.tokenizer_id, do_lower_case=False)
+            self.model = T5EncoderModel.from_pretrained(spec.model_id)
+            self.model.to(target_device)
+            self.model.eval()
+
+        elif spec.family == "esmc":
+            from esm.models.esmc import ESMC
+            self.model = ESMC.from_pretrained(spec.model_id).to(target_device)
+            self.model.eval()
+
+        elif spec.family == "prosst":
+            self.tokenizer = AutoTokenizer.from_pretrained(spec.tokenizer_id, trust_remote_code=True)
+            self.model = AutoModel.from_pretrained(
+                spec.model_id,
+                trust_remote_code=True,
+                output_hidden_states=True,
+            )
+            self.model.to(target_device)
+            self.model.eval()
+
+            # Official ProSST sequence-only route:
+            # predict structure tokens from sequence, then feed them into ProSST.
+            from prosst.structure.get_sst_seq import SSTPredictor
+            self.sst_predictor = SSTPredictor()
+
+        else:
+            raise ValueError(f"Unsupported family: {spec.family}")
+
+        self.model_key = model_key
+        self.family = spec.family
+        self.device = target_device
+
+
+RUNNER = SingleModelRunner()
 
 
-# =========================
-# Model-specific embedding
-# =========================
 @torch.no_grad()
-def embed_one_hf_encoder(seq: str, model, tokenizer, device: str) -> np.ndarray:
-    enc = tokenizer(
+def embed_hf_encoder(seq: str) -> torch.Tensor:
+    enc = RUNNER.tokenizer(
         seq,
         return_tensors="pt",
         add_special_tokens=True,
         return_special_tokens_mask=True,
         truncation=False,
     )
-    enc = {k: v.to(device) for k, v in enc.items()}
-    out = model(**{k: v for k, v in enc.items() if k != "special_tokens_mask"})
-    hidden = out.last_hidden_state[0]  # [T, d]
-
-    special_tokens_mask = enc.get("special_tokens_mask", None)
-    attention_mask = enc.get("attention_mask", None)
+    enc = {k: v.to(RUNNER.device) for k, v in enc.items()}
+    out = RUNNER.model(**{k: v for k, v in enc.items() if k != "special_tokens_mask"})
+    hidden = out.last_hidden_state[0]
 
-    residue_hidden = normalize_to_residue_level(
+    emb = normalize_to_Ld(
         hidden=hidden,
         expected_len=len(seq),
-        special_tokens_mask=special_tokens_mask[0] if special_tokens_mask is not None else None,
-        attention_mask=attention_mask[0] if attention_mask is not None else None,
+        special_tokens_mask=enc.get("special_tokens_mask", None)[0] if enc.get("special_tokens_mask", None) is not None else None,
+        attention_mask=enc.get("attention_mask", None)[0] if enc.get("attention_mask", None) is not None else None,
     )
-    return residue_hidden.detach().cpu().float().numpy()
+    return emb.detach().cpu().float()
 
 
 @torch.no_grad()
-def embed_one_t5_encoder(seq: str, model, tokenizer, device: str) -> np.ndarray:
-    # ProtT5 style preprocessing: uppercase residues separated by spaces.
+def embed_t5_encoder(seq: str) -> torch.Tensor:
     spaced = protein_to_spaced(seq)
-    enc = tokenizer(
+    enc = RUNNER.tokenizer(
         spaced,
         return_tensors="pt",
         add_special_tokens=True,
         return_special_tokens_mask=True,
         truncation=False,
     )
-    enc = {k: v.to(device) for k, v in enc.items()}
-    out = model(**{k: v for k, v in enc.items() if k != "special_tokens_mask"})
+    enc = {k: v.to(RUNNER.device) for k, v in enc.items()}
+    out = RUNNER.model(**{k: v for k, v in enc.items() if k != "special_tokens_mask"})
     hidden = out.last_hidden_state[0]
 
-    special_tokens_mask = enc.get("special_tokens_mask", None)
-    attention_mask = enc.get("attention_mask", None)
-
-    residue_hidden = normalize_to_residue_level(
+    emb = normalize_to_Ld(
         hidden=hidden,
         expected_len=len(seq),
-        special_tokens_mask=special_tokens_mask[0] if special_tokens_mask is not None else None,
-        attention_mask=attention_mask[0] if attention_mask is not None else None,
+        special_tokens_mask=enc.get("special_tokens_mask", None)[0] if enc.get("special_tokens_mask", None) is not None else None,
+        attention_mask=enc.get("attention_mask", None)[0] if enc.get("attention_mask", None) is not None else None,
     )
-    return residue_hidden.detach().cpu().float().numpy()
+    return emb.detach().cpu().float()
 
 
 @torch.no_grad()
-def embed_one_esmc(seq: str, model, device: str) -> np.ndarray:
+def embed_esmc(seq: str) -> torch.Tensor:
     from esm.sdk.api import ESMProtein, LogitsConfig
 
     protein = ESMProtein(sequence=seq)
-    protein_tensor = model.encode(protein)
-    out = model.logits(
+    protein_tensor = RUNNER.model.encode(protein)
+    out = RUNNER.model.logits(
         protein_tensor,
         LogitsConfig(sequence=True, return_embeddings=True)
     )
 
     emb = out.embeddings
-    if isinstance(emb, np.ndarray):
-        arr = emb
-    else:
-        arr = emb.detach().cpu().float().numpy()
-
-    # Expected shape is typically [1, T, d] or [T, d]
-    if arr.ndim == 3:
-        arr = arr[0]
-
-    if arr.shape[0] == len(seq):
-        return arr
-    if arr.shape[0] == len(seq) + 2:
-        return arr[1:-1]
-    if arr.shape[0] == len(seq) + 1:
-        return arr[:len(seq)]
-    if arr.shape[0] > len(seq):
-        return arr[:len(seq)]
-
-    raise ValueError(
-        f"ESMC returned incompatible shape {arr.shape} for sequence length {len(seq)}."
-    )
+    if not isinstance(emb, torch.Tensor):
+        emb = torch.tensor(emb)
 
+    if emb.ndim == 3:
+        emb = emb[0]
 
-def embed_sequences(
-    fasta_text: str,
-    model_key: str,
-    device: str,
-    progress=gr.Progress(track_tqdm=False),
-):
-    records = parse_fasta(fasta_text)
-
-    cleaned_records = []
-    global_warnings = []
+    if emb.shape[0] == len(seq):
+        return emb.detach().cpu().float()
+    if emb.shape[0] == len(seq) + 2:
+        return emb[1:-1].detach().cpu().float()
+    if emb.shape[0] == len(seq) + 1:
+        return emb[:len(seq)].detach().cpu().float()
+    if emb.shape[0] > len(seq):
+        return emb[:len(seq)].detach().cpu().float()
 
-    for rec in records:
-        clean_seq, warnings = clean_sequence(rec["sequence"])
-        cleaned_records.append({"id": rec["id"], "sequence": clean_seq})
-        for w in warnings:
-            global_warnings.append(f"{rec['id']}: {w}")
+    raise ValueError(f"ESMC returned shape {tuple(emb.shape)} for sequence length {len(seq)}.")
 
-    MODEL_MANAGER.load(model_key, device)
-    spec = MODEL_SPECS[model_key]
 
-    embeddings_by_id: Dict[str, np.ndarray] = {}
-    summary_rows = []
-    first_preview = None
-    first_preview_name = None
+@torch.no_grad()
+def embed_prosst(seq: str) -> torch.Tensor:
+    # Sequence-only mode:
+    # 1) predict structure token sequence from amino-acid sequence
+    # 2) feed sequence + structure tokens into ProSST
+    structure_tokens = RUNNER.sst_predictor.predict(seq)
+
+    # Structure tokens may come back as list[int], np.ndarray, or space-separated string
+    if isinstance(structure_tokens, str):
+        sst_seq = structure_tokens
+    else:
+        sst_seq = " ".join([str(x) for x in structure_tokens])
 
-    for idx, rec in enumerate(cleaned_records, start=1):
-        seq_id = rec["id"]
-        seq = rec["sequence"]
-        progress((idx - 1) / max(len(cleaned_records), 1), desc=f"Embedding {seq_id}")
+    aa_spaced = protein_to_spaced(seq)
 
-        if spec.family == "hf_encoder":
-            emb = embed_one_hf_encoder(seq, MODEL_MANAGER.model, MODEL_MANAGER.tokenizer, MODEL_MANAGER.device)
-        elif spec.family == "t5_encoder":
-            emb = embed_one_t5_encoder(seq, MODEL_MANAGER.model, MODEL_MANAGER.tokenizer, MODEL_MANAGER.device)
-        elif spec.family == "esmc":
-            emb = embed_one_esmc(seq, MODEL_MANAGER.model, MODEL_MANAGER.device)
-        else:
-            raise ValueError(f"Unsupported family: {spec.family}")
+    enc = RUNNER.tokenizer(
+        aa_spaced,
+        return_tensors="pt",
+        add_special_tokens=True,
+        return_special_tokens_mask=True,
+        truncation=False,
+    )
+    enc = {k: v.to(RUNNER.device) for k, v in enc.items()}
+
+    # Different ProSST remote-code implementations may expect different kwarg names.
+    # Try the common names first.
+    tried = []
+
+    for kw in ("ss_input_ids", "structure_ids", "sst_input_ids", "struc_input_ids"):
+        try:
+            sst_enc = RUNNER.tokenizer(
+                sst_seq,
+                return_tensors="pt",
+                add_special_tokens=True,
+                truncation=False,
+            )
+            sst_ids = sst_enc["input_ids"].to(RUNNER.device)
 
-        if emb.shape[0] != len(seq):
-            raise ValueError(
-                f"Normalization failed for {seq_id}: got {emb.shape}, expected first dimension {len(seq)}."
+            out = RUNNER.model(
+                input_ids=enc["input_ids"],
+                attention_mask=enc.get("attention_mask", None),
+                output_hidden_states=True,
+                **{kw: sst_ids},
             )
 
-        embeddings_by_id[seq_id] = emb
-
-        summary_rows.append({
-            "id": seq_id,
-            "length_L": len(seq),
-            "embedding_dim_d": emb.shape[1],
-            "shape": f"{emb.shape[0]} x {emb.shape[1]}",
-            "model": model_key,
-        })
-
-        if first_preview is None:
-            preview_rows = min(20, emb.shape[0])
-            preview_cols = min(8, emb.shape[1])
-            df = pd.DataFrame(
-                emb[:preview_rows, :preview_cols],
-                index=[f"res_{i+1}" for i in range(preview_rows)],
-                columns=[f"dim_{j+1}" for j in range(preview_cols)],
+            hidden = out.hidden_states[-1][0]
+            emb = normalize_to_Ld(
+                hidden=hidden,
+                expected_len=len(seq),
+                special_tokens_mask=enc.get("special_tokens_mask", None)[0] if enc.get("special_tokens_mask", None) is not None else None,
+                attention_mask=enc.get("attention_mask", None)[0] if enc.get("attention_mask", None) is not None else None,
             )
-            first_preview = df
-            first_preview_name = seq_id
+            return emb.detach().cpu().float()
 
-    progress(1.0, desc="Packaging outputs")
+        except Exception as e:
+            tried.append(f"{kw}: {repr(e)}")
 
-    out_zip = package_outputs(
-        embeddings_by_id=embeddings_by_id,
-        sequences={x["id"]: x["sequence"] for x in cleaned_records},
-        model_key=model_key,
-        notes=global_warnings,
+    raise RuntimeError(
+        "Failed to run ProSST. The installed ProSST remote-code signature may differ. "
+        + " | ".join(tried)
     )
 
-    summary_df = pd.DataFrame(summary_rows)
-    log_text = []
-    log_text.append(f"Loaded model: {model_key}")
-    log_text.append(f"Resolved device: {MODEL_MANAGER.device}")
-    log_text.append(f"Processed sequences: {len(cleaned_records)}")
-    if global_warnings:
-        log_text.append("")
-        log_text.append("Warnings:")
-        log_text.extend(global_warnings)
-
-    preview_markdown = f"Preview shown for: `{first_preview_name}`"
-
-    return summary_df, first_preview, preview_markdown, out_zip, "\n".join(log_text)
-
-
-def package_outputs(
-    embeddings_by_id: Dict[str, np.ndarray],
-    sequences: Dict[str, str],
-    model_key: str,
-    notes: List[str],
-) -> str:
-    tmpdir = tempfile.mkdtemp(prefix="protein_embedding_hub_")
-    zip_path = os.path.join(tmpdir, f"{safe_filename(model_key)}_embeddings.zip")
-
-    summary_rows = []
-    for seq_id, emb in embeddings_by_id.items():
-        summary_rows.append({
-            "id": seq_id,
-            "length_L": sequences[seq_id].__len__(),
-            "embedding_dim_d": emb.shape[1],
-            "shape": f"{emb.shape[0]} x {emb.shape[1]}",
-            "npy_file": f"{safe_filename(seq_id)}.npy",
-        })
-
-    summary_df = pd.DataFrame(summary_rows)
-    sequences_df = pd.DataFrame(
-        [{"id": k, "sequence": v} for k, v in sequences.items()]
-    )
+
+def embed_one_sequence(seq: str) -> torch.Tensor:
+    if RUNNER.family == "hf_encoder":
+        return embed_hf_encoder(seq)
+    if RUNNER.family == "t5_encoder":
+        return embed_t5_encoder(seq)
+    if RUNNER.family == "esmc":
+        return embed_esmc(seq)
+    if RUNNER.family == "prosst":
+        return embed_prosst(seq)
+    raise ValueError(f"Unsupported family: {RUNNER.family}")
+
+
+def run_embedding(fasta_text: str, model_keys: List[str], device: str, progress=gr.Progress()):
+    if not model_keys:
+        raise ValueError("Please select at least one model.")
+
+    records = parse_fasta(fasta_text)
+    records = [{"id": r["id"], "sequence": clean_sequence(r["sequence"])} for r in records]
+
+    tmpdir = tempfile.mkdtemp(prefix="protein_embeddings_")
+    zip_path = os.path.join(tmpdir, "embeddings.zip")
+
+    total_steps = len(model_keys) * len(records)
+    step = 0
 
     with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
-        # summary.csv
-        with io.StringIO() as s:
-            summary_df.to_csv(s, index=False)
-            zf.writestr("summary.csv", s.getvalue())
+        for model_key in model_keys:
+            RUNNER.load(model_key, device)
 
-        # sequences.csv
-        with io.StringIO() as s:
-            sequences_df.to_csv(s, index=False)
-            zf.writestr("sequences.csv", s.getvalue())
+            for rec in records:
+                step += 1
+                progress(step / total_steps, desc=f"{model_key} | {rec['id']}")
+                emb = embed_one_sequence(rec["sequence"])
 
-        # notes.txt
-        note_text = "\n".join(notes) if notes else "No warnings."
-        zf.writestr("notes.txt", note_text)
+                if emb.ndim != 2 or emb.shape[0] != len(rec["sequence"]):
+                    raise ValueError(
+                        f"{model_key} failed on {rec['id']}: got shape {tuple(emb.shape)}, expected ({len(rec['sequence'])}, d)"
+                    )
 
-        # per-sequence npy
-        for seq_id, emb in embeddings_by_id.items():
-            npy_name = f"embeddings/{safe_filename(seq_id)}.npy"
-            buf = io.BytesIO()
-            np.save(buf, emb)
-            zf.writestr(npy_name, buf.getvalue())
+                inner_name = f"{safe_filename(model_key)}/{safe_filename(rec['id'])}.pt"
+                buffer = io.BytesIO()
+                torch.save(emb, buffer)
+                zf.writestr(inner_name, buffer.getvalue())
 
-    return zip_path
+    return zip_path, f"Done: {len(records)} sequence(s), {len(model_keys)} model(s)."
 
 
-def clear_loaded_model():
-    MODEL_MANAGER.unload()
-    return "Model cache cleared."
+def clear_cache():
+    RUNNER.unload()
+    return "Cache cleared."
 
 
-# =========================
-# Gradio UI
-# =========================
 EXAMPLE_FASTA = """>seq1
 MKWVTFISLLLLFSSAYSRGVFRRDTHKSEIAHRFKDLGE
 >seq2
@@ -554,70 +458,45 @@ GAVLILKKKGHHEAELKPLAQSHATKHKIPIKYLEFISEAIIHVLHSR
 
 with gr.Blocks(title=APP_TITLE) as demo:
     gr.Markdown(f"# {APP_TITLE}")
-    gr.Markdown(APP_DESC)
-
-    with gr.Row():
-        with gr.Column(scale=2):
-            fasta_input = gr.Textbox(
-                label="Protein FASTA input",
-                lines=16,
-                value=EXAMPLE_FASTA,
-                placeholder="Paste FASTA here..."
-            )
 
-            model_dropdown = gr.Dropdown(
-                choices=list(MODEL_SPECS.keys()),
-                value="ESM2-150M",
-                label="Model"
-            )
+    fasta_input = gr.Textbox(
+        label="FASTA",
+        lines=16,
+        value=EXAMPLE_FASTA,
+        placeholder="Paste FASTA here",
+    )
 
-            device_dropdown = gr.Dropdown(
-                choices=["auto", "cuda", "cpu"],
-                value="auto",
-                label="Device"
-            )
+    model_select = gr.CheckboxGroup(
+        choices=list(MODEL_SPECS.keys()),
+        value=["ESM2-150M"],
+        label="Models",
+    )
 
-            with gr.Row():
-                run_btn = gr.Button("Run embedding", variant="primary")
-                clear_btn = gr.Button("Clear loaded model")
-
-        with gr.Column(scale=1):
-            gr.Markdown("## Notes")
-            gr.Markdown(
-                "- Output is always normalized to residue-level `L*d`\n"
-                "- ZIP contains one `.npy` per sequence\n"
-                "- `summary.csv` records final shapes\n"
-                "- Large models need GPU"
-            )
-            model_note = gr.Markdown(
-                value="\n".join(
-                    [f"- **{k}**: {v.note}" for k, v in MODEL_SPECS.items()]
-                )
-            )
+    device_select = gr.Dropdown(
+        choices=["auto", "cuda", "cpu"],
+        value="auto",
+        label="Device",
+    )
 
     with gr.Row():
-        summary_output = gr.Dataframe(label="Summary", interactive=False)
-    with gr.Row():
-        preview_note = gr.Markdown()
-    with gr.Row():
-        preview_output = gr.Dataframe(label="Embedding preview (first sequence)", interactive=False)
+        run_btn = gr.Button("Run", variant="primary")
+        clear_btn = gr.Button("Clear cache")
 
-    with gr.Row():
-        download_output = gr.File(label="Download ZIP")
-    with gr.Row():
-        log_output = gr.Textbox(label="Log", lines=10)
+    output_file = gr.File(label="Download")
+    log_box = gr.Textbox(label="Log", lines=4)
 
     run_btn.click(
-        fn=embed_sequences,
-        inputs=[fasta_input, model_dropdown, device_dropdown],
-        outputs=[summary_output, preview_output, preview_note, download_output, log_output],
+        fn=run_embedding,
+        inputs=[fasta_input, model_select, device_select],
+        outputs=[output_file, log_box],
     )
 
     clear_btn.click(
-        fn=clear_loaded_model,
+        fn=clear_cache,
         inputs=[],
-        outputs=[log_output],
+        outputs=[log_box],
     )
 
-demo.queue(max_size=16)
-demo.launch()
+demo.queue(max_size=8)
+demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)
+