Add /downstream endpoint serving 8 fine-tuned task heads per backbone

app.py

@@ -7,6 +7,7 @@ Endpoints:
     POST /predict         — per-position 20-AA softmax over the sequence
     POST /embed           — mean-pooled last-layer embedding
     POST /mutation        — wildtype-marginal Δlog-likelihood matrix [L, 20]
+    POST /downstream      — all fine-tuned task heads (classification + regression)
 """
 
 import os
@@ -17,7 +18,7 @@ import torch.nn.functional as F
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
-from transformers import EsmForMaskedLM, EsmTokenizer
+from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmTokenizer
 
 
 MODEL_REGISTRY: Dict[str, str] = {
@@ -26,11 +27,46 @@ MODEL_REGISTRY: Dict[str, str] = {
     "miplm-msa":    "HUBioDataLab/esm2-8m-msa",
 }
 
+
+# Downstream heads — one HF repo per backbone, with 8 task subfolders each.
+# Map frontend backbone name → HF repo.
+DOWNSTREAM_REGISTRY: Dict[str, str] = {
+    "miplm-ce":     "HUBioDataLab/miplm-ce-tasks",
+    "miplm-blosum": "HUBioDataLab/miplm-blosum-tasks",
+    "miplm-msa":    "HUBioDataLab/miplm-msa-tasks",
+}
+
+# Task metadata. `kind` drives post-processing:
+#   single_label  → softmax → top-K classes with probabilities
+#   multi_label   → sigmoid → top-K classes with independent probabilities
+#   regression    → raw scalar
+# `labels` is the human-readable class vocabulary (when known). When unset the
+# response falls back to numeric class indices.
+DEEPLOC10_LABELS = [
+    "Cytoplasm", "Nucleus", "Extracellular", "Cell membrane",
+    "Endoplasmic reticulum", "Plastid", "Golgi apparatus",
+    "Lysosome/Vacuole", "Mitochondrion", "Peroxisome",
+]
+DEEPLOC2_LABELS = ["Soluble", "Membrane"]
+METAL_LABELS    = ["Non-binder", "Metal-ion binder"]
+
+DOWNSTREAM_TASKS: Dict[str, Dict] = {
+    "deeploc-cls10":   {"kind": "single_label", "labels": DEEPLOC10_LABELS, "title": "Subcellular localization (10-way)"},
+    "deeploc-cls2":    {"kind": "single_label", "labels": DEEPLOC2_LABELS,  "title": "Soluble vs membrane"},
+    "metalionbinding": {"kind": "single_label", "labels": METAL_LABELS,     "title": "Metal-ion binding"},
+    "thermostability": {"kind": "regression",   "labels": None,             "title": "Thermostability", "unit": "normalised score"},
+    "ec":              {"kind": "multi_label",  "labels": None,             "title": "EC enzyme class",        "num_classes": 585},
+    "go-bp":           {"kind": "multi_label",  "labels": None,             "title": "GO biological process",  "num_classes": 1943},
+    "go-cc":           {"kind": "multi_label",  "labels": None,             "title": "GO cellular component",  "num_classes": 320},
+    "go-mf":           {"kind": "multi_label",  "labels": None,             "title": "GO molecular function",  "num_classes": 489},
+}
+
 STANDARD_AAS = "ACDEFGHIKLMNPQRSTVWY"
 MAX_LEN = 1024
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 _models: Dict[str, Tuple[EsmForMaskedLM, EsmTokenizer, List[int]]] = {}
+_downstream: Dict[Tuple[str, str], EsmForSequenceClassification] = {}
 
 
 def get_model(name: str):
@@ -100,6 +136,10 @@ def _warmup():
 def root():
     return {
         "models": list(MODEL_REGISTRY.keys()),
+        "downstream_tasks": [
+            {"id": t, "title": cfg["title"], "kind": cfg["kind"]}
+            for t, cfg in DOWNSTREAM_TASKS.items()
+        ],
         "device": DEVICE,
         "max_length": MAX_LEN,
         "aa_order": STANDARD_AAS,
@@ -157,3 +197,95 @@ def mutation(req: _SeqRequest):
     wt_logp = log_probs.gather(1, wt_idx[:, None])                    # [L, 1]
     delta = (log_probs - wt_logp).tolist()
     return MutationResponse(model=req.model, sequence=seq, aa_order=STANDARD_AAS, scores=delta)
+
+
+# ─── Downstream task heads ──────────────────────────────────────────────────
+
+def get_downstream(backbone: str, task: str) -> EsmForSequenceClassification:
+    """Lazy-load and cache fine-tuned classification/regression heads."""
+    if backbone not in DOWNSTREAM_REGISTRY:
+        raise HTTPException(400, f"unknown backbone {backbone!r}; available: {list(DOWNSTREAM_REGISTRY)}")
+    if task not in DOWNSTREAM_TASKS:
+        raise HTTPException(400, f"unknown task {task!r}; available: {list(DOWNSTREAM_TASKS)}")
+    key = (backbone, task)
+    if key not in _downstream:
+        repo = DOWNSTREAM_REGISTRY[backbone]
+        print(f"[downstream] loading {repo} :: {task}")
+        m = EsmForSequenceClassification.from_pretrained(repo, subfolder=task).to(DEVICE).eval()
+        _downstream[key] = m
+    return _downstream[key]
+
+
+class DownstreamRequest(BaseModel):
+    sequence: str = Field(..., min_length=1, max_length=MAX_LEN, pattern=r"^[A-Za-z]+$")
+    backbone: str = "miplm-blosum"
+    top_k: int = 3
+
+
+class TaskPrediction(BaseModel):
+    task: str
+    title: str
+    kind: str                                     # single_label | multi_label | regression
+    value: float | None = None                    # set for regression
+    unit: str | None = None
+    top: List[Dict] | None = None                 # set for classification — [{label/index, prob}]
+    num_classes: int | None = None
+
+
+class DownstreamResponse(BaseModel):
+    backbone: str
+    sequence: str
+    predictions: List[TaskPrediction]
+
+
+@app.post("/downstream", response_model=DownstreamResponse)
+@torch.inference_mode()
+def downstream(req: DownstreamRequest):
+    seq = req.sequence.upper()
+    if req.backbone not in DOWNSTREAM_REGISTRY:
+        raise HTTPException(400, f"unknown backbone {req.backbone!r}")
+
+    # All downstream heads share the ESM-2 tokenizer with the base backbone.
+    _, tokenizer, _ = get_model(req.backbone)
+    batch = tokenizer(seq, return_tensors="pt").to(DEVICE)
+
+    predictions: List[TaskPrediction] = []
+    for task, cfg in DOWNSTREAM_TASKS.items():
+        model = get_downstream(req.backbone, task)
+        logits = model(**batch).logits[0]   # [num_labels] for sequence-level head
+
+        kind = cfg["kind"]
+        if kind == "regression":
+            predictions.append(TaskPrediction(
+                task=task, title=cfg["title"], kind=kind,
+                value=float(logits.item()),
+                unit=cfg.get("unit"),
+            ))
+        elif kind == "single_label":
+            probs = F.softmax(logits, dim=-1)
+            top_idx = torch.topk(probs, k=min(req.top_k, probs.numel())).indices.tolist()
+            top = [
+                {"label": cfg["labels"][i] if cfg["labels"] else f"class_{i}",
+                 "index": int(i),
+                 "prob": float(probs[i].item())}
+                for i in top_idx
+            ]
+            predictions.append(TaskPrediction(
+                task=task, title=cfg["title"], kind=kind,
+                top=top, num_classes=int(probs.numel()),
+            ))
+        else:  # multi_label
+            probs = torch.sigmoid(logits)
+            top_idx = torch.topk(probs, k=min(req.top_k, probs.numel())).indices.tolist()
+            top = [
+                {"label": (cfg["labels"][i] if cfg["labels"] else f"class_{i}"),
+                 "index": int(i),
+                 "prob": float(probs[i].item())}
+                for i in top_idx
+            ]
+            predictions.append(TaskPrediction(
+                task=task, title=cfg["title"], kind=kind,
+                top=top, num_classes=int(probs.numel()),
+            ))
+
+    return DownstreamResponse(backbone=req.backbone, sequence=seq, predictions=predictions)