Create app.py

app.py

@@ -0,0 +1,623 @@
+import gc
+import io
+import os
+import re
+import zipfile
+import tempfile
+from dataclasses import dataclass
+from typing import Dict, List, Tuple, Optional
+
+import gradio as gr
+import numpy as np
+import pandas as pd
+import torch
+
+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.
+"""
+
+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"
+    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."
+    ),
+}
+
+
+# =========================
+# 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
+
+
+MODEL_MANAGER = ModelManager()
+
+
+# =========================
+# FASTA and sequence utils
+# =========================
+def parse_fasta(text: str) -> List[Dict[str, str]]:
+    text = text.strip()
+    if not text:
+        raise ValueError("Empty FASTA input.")
+
+    records = []
+    current_id = None
+    current_seq = []
+
+    for raw_line in text.splitlines():
+        line = raw_line.strip()
+        if not line:
+            continue
+        if line.startswith(">"):
+            if current_id is not None:
+                seq = "".join(current_seq).strip()
+                if not seq:
+                    raise ValueError(f"Sequence for record '{current_id}' is empty.")
+                records.append({"id": current_id, "sequence": seq})
+            current_id = line[1:].strip() or f"seq_{len(records)+1}"
+            current_seq = []
+        else:
+            if current_id is None:
+                current_id = f"seq_{len(records)+1}"
+            current_seq.append(line)
+
+    if current_id is not None:
+        seq = "".join(current_seq).strip()
+        if not seq:
+            raise ValueError(f"Sequence for record '{current_id}' is empty.")
+        records.append({"id": current_id, "sequence": seq})
+
+    if not records:
+        raise ValueError("No FASTA records found.")
+    return records
+
+
+def clean_sequence(seq: str) -> Tuple[str, List[str]]:
+    seq = re.sub(r"\s+", "", seq).upper()
+    warnings = []
+
+    if not seq:
+        raise ValueError("Empty sequence after cleaning.")
+
+    bad = sorted({c for c in seq if c not in ALLOWED_AA})
+    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
+
+
+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(
+    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)}")
+
+    T, d = hidden.shape
+
+    if special_tokens_mask is not None:
+        mask = 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]
+
+    if T == expected_len:
+        return hidden
+    if T == expected_len + 2:
+        return hidden[1:-1]
+    if T == expected_len + 1:
+        return hidden[:expected_len]
+    if T > expected_len:
+        return hidden[:expected_len]
+
+    raise ValueError(
+        f"Could not normalize token length {T} to residue length {expected_len}."
+    )
+
+
+# =========================
+# Model-specific embedding
+# =========================
+@torch.no_grad()
+def embed_one_hf_encoder(seq: str, model, tokenizer, device: str) -> np.ndarray:
+    enc = 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)
+
+    residue_hidden = normalize_to_residue_level(
+        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,
+    )
+    return residue_hidden.detach().cpu().float().numpy()
+
+
+@torch.no_grad()
+def embed_one_t5_encoder(seq: str, model, tokenizer, device: str) -> np.ndarray:
+    # ProtT5 style preprocessing: uppercase residues separated by spaces.
+    spaced = protein_to_spaced(seq)
+    enc = 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"})
+    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(
+        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,
+    )
+    return residue_hidden.detach().cpu().float().numpy()
+
+
+@torch.no_grad()
+def embed_one_esmc(seq: str, model, device: str) -> np.ndarray:
+    from esm.sdk.api import ESMProtein, LogitsConfig
+
+    protein = ESMProtein(sequence=seq)
+    protein_tensor = model.encode(protein)
+    out = 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)}."
+    )
+
+
+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 = []
+
+    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}")
+
+    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
+
+    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}")
+
+        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}")
+
+        if emb.shape[0] != len(seq):
+            raise ValueError(
+                f"Normalization failed for {seq_id}: got {emb.shape}, expected first dimension {len(seq)}."
+            )
+
+        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)],
+            )
+            first_preview = df
+            first_preview_name = seq_id
+
+    progress(1.0, desc="Packaging outputs")
+
+    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,
+    )
+
+    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()]
+    )
+
+    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())
+
+        # sequences.csv
+        with io.StringIO() as s:
+            sequences_df.to_csv(s, index=False)
+            zf.writestr("sequences.csv", s.getvalue())
+
+        # notes.txt
+        note_text = "\n".join(notes) if notes else "No warnings."
+        zf.writestr("notes.txt", note_text)
+
+        # 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())
+
+    return zip_path
+
+
+def clear_loaded_model():
+    MODEL_MANAGER.unload()
+    return "Model cache cleared."
+
+
+# =========================
+# Gradio UI
+# =========================
+EXAMPLE_FASTA = """>seq1
+MKWVTFISLLLLFSSAYSRGVFRRDTHKSEIAHRFKDLGE
+>seq2
+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"
+            )
+
+            device_dropdown = gr.Dropdown(
+                choices=["auto", "cuda", "cpu"],
+                value="auto",
+                label="Device"
+            )
+
+            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()]
+                )
+            )
+
+    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)
+
+    with gr.Row():
+        download_output = gr.File(label="Download ZIP")
+    with gr.Row():
+        log_output = gr.Textbox(label="Log", lines=10)
+
+    run_btn.click(
+        fn=embed_sequences,
+        inputs=[fasta_input, model_dropdown, device_dropdown],
+        outputs=[summary_output, preview_output, preview_note, download_output, log_output],
+    )
+
+    clear_btn.click(
+        fn=clear_loaded_model,
+        inputs=[],
+        outputs=[log_output],
+    )
+
+demo.queue(max_size=16)
+demo.launch()