Initial sync of ESM-2 model-serving app from cancer-mutation-predictor repo

Dockerfile

@@ -0,0 +1,11 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+EXPOSE 7860
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,10 +1,21 @@
 ---
-title: Cancer Mutation Esm2 Serving
-emoji: 👁
-colorFrom: pink
-colorTo: green
+title: Cancer Mutation ESM2 Serving
+emoji: 🧬
+colorFrom: blue
+colorTo: gray
 sdk: docker
+app_port: 7860
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# ESM-2 model serving
+
+Internal-only inference API for the Cancer Mutation Predictor project. Loads
+a fine-tuned `facebook/esm2_t12_35M_UR50D` checkpoint from
+`Arko007/esm2-cancer-nlr-35M` at startup and exposes `/score` and
+`/score/batch`, gated by a shared-secret `X-Internal-Key` header so only the
+orchestrator backend can call it.
+
+This folder is synced here automatically from
+[Anamitra-Sarkar/cancer-mutation-predictor](https://github.com/Anamitra-Sarkar/cancer-mutation-predictor)
+via GitHub Actions — do not edit directly in the Space.

app.py

@@ -0,0 +1,113 @@
+"""
+HF Space (Arko006 account) — the ONLY component that loads the ESM-2 model.
+Downloads the fine-tuned checkpoint from the Arko007 HF *model* repo at
+container startup (never bundled into this Space's own git history) and
+exposes a small, shared-secret-gated scoring API for the Render orchestrator.
+"""
+
+import os
+import threading
+from contextlib import asynccontextmanager
+from typing import Optional
+
+import torch
+from fastapi import FastAPI, Header, HTTPException
+from pydantic import BaseModel
+from huggingface_hub import hf_hub_download
+
+from model import TransformerDMSRegressor, DEFAULT_MODEL_NAME
+
+CHECKPOINT_REPO = os.getenv("CHECKPOINT_REPO", "Arko007/esm2-cancer-nlr-35M")
+CHECKPOINT_FILE = os.getenv("CHECKPOINT_FILE", "best.pt")
+SHARED_SECRET = os.getenv("MODEL_SERVING_SHARED_SECRET", "")
+
+_state = {"model": None, "loading": True, "revision": None, "error": None}
+_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def _load_model():
+    try:
+        model = TransformerDMSRegressor(DEFAULT_MODEL_NAME)
+        checkpoint_path = hf_hub_download(repo_id=CHECKPOINT_REPO, filename=CHECKPOINT_FILE)
+        state_dict = torch.load(checkpoint_path, map_location=_device)
+        model.load_state_dict(state_dict["model_state_dict"])
+        model.to(_device)
+        model.eval()
+        _state["model"] = model
+        _state["revision"] = f"{CHECKPOINT_REPO}/{CHECKPOINT_FILE} (epoch {state_dict.get('epoch', '?')})"
+    except Exception as e:
+        # No fine-tuned checkpoint yet is a legitimate startup state (e.g.
+        # before the first Kaggle training run has completed) — fall back to
+        # the base pretrained ESM-2 model so /score still returns real
+        # zero-shot LLR scores (nlr_fitness will be None until fine-tuned).
+        print(f"Could not load fine-tuned checkpoint ({e}); falling back to base ESM-2 zero-shot scoring.")
+        try:
+            model = TransformerDMSRegressor(DEFAULT_MODEL_NAME)
+            model.to(_device)
+            model.eval()
+            _state["model"] = model
+            _state["revision"] = f"{DEFAULT_MODEL_NAME} (base, not fine-tuned)"
+        except Exception as inner_e:
+            _state["error"] = str(inner_e)
+    finally:
+        _state["loading"] = False
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    threading.Thread(target=_load_model, daemon=True).start()
+    yield
+
+
+app = FastAPI(title="Cancer Mutation Predictor — ESM-2 model serving", lifespan=lifespan)
+
+
+class ScoreRequest(BaseModel):
+    sequence: str
+    position: int
+    ref_aa: str
+    alt_aa: str
+
+
+def _check_secret(x_internal_key: Optional[str]):
+    if SHARED_SECRET and x_internal_key != SHARED_SECRET:
+        raise HTTPException(status_code=401, detail="Invalid or missing X-Internal-Key")
+
+
+@app.get("/health")
+def health():
+    return {
+        "status": "ok" if not _state["loading"] else "warming",
+        "model_loaded": _state["model"] is not None,
+        "checkpoint": _state["revision"],
+        "error": _state["error"],
+    }
+
+
+@app.post("/score")
+def score(req: ScoreRequest, x_internal_key: Optional[str] = Header(None)):
+    _check_secret(x_internal_key)
+    if _state["model"] is None:
+        raise HTTPException(status_code=503, detail="Model still warming up, try again shortly")
+
+    if not (1 <= req.position <= len(req.sequence)):
+        raise HTTPException(status_code=400, detail="position out of range for the given sequence")
+
+    import time
+    t0 = time.time()
+    llr, fitness = _state["model"].score(req.sequence, req.position, req.ref_aa, req.alt_aa, _device)
+    return {
+        "raw_llr": llr,
+        "nlr_fitness": fitness if "not fine-tuned" not in (_state["revision"] or "") else None,
+        "model_id": DEFAULT_MODEL_NAME,
+        "checkpoint_revision": _state["revision"],
+        "inference_ms": round((time.time() - t0) * 1000, 1),
+    }
+
+
+@app.post("/score/batch")
+def score_batch(reqs: list[ScoreRequest], x_internal_key: Optional[str] = Header(None)):
+    _check_secret(x_internal_key)
+    if len(reqs) > 50:
+        raise HTTPException(status_code=400, detail="Batch limited to 50 variants per call")
+    return [score(r, x_internal_key) for r in reqs]

model.py

@@ -0,0 +1,50 @@
+"""
+Inference-mode mirror of training/model.py's TransformerDMSRegressor.
+Kept as a separate copy (not an import) because this file is the ONLY
+directory git-synced to the HF Space — it must be fully self-contained.
+"""
+
+import torch
+import torch.nn as nn
+from transformers import EsmForMaskedLM, EsmTokenizer
+
+DEFAULT_MODEL_NAME = "facebook/esm2_t12_35M_UR50D"
+
+
+class TransformerDMSRegressor(nn.Module):
+    def __init__(self, model_name: str = DEFAULT_MODEL_NAME):
+        super().__init__()
+        self.tokenizer = EsmTokenizer.from_pretrained(model_name)
+        self.backbone = EsmForMaskedLM.from_pretrained(model_name)
+        for param in self.backbone.parameters():
+            param.requires_grad = False
+        for param in self.backbone.esm.encoder.layer[-2:].parameters():
+            param.requires_grad = True
+        self.regression_head = nn.Linear(1, 1)
+
+    @torch.no_grad()
+    def score(self, sequence: str, position: int, ref_aa: str, alt_aa: str, device: torch.device):
+        max_len = 1022  # ESM tokenizer budget minus special tokens
+        start = max(0, position - 1 - max_len // 2)
+        end = min(len(sequence), start + max_len)
+        start = max(0, end - max_len)
+        window = sequence[start:end]
+        local_pos = position - start  # 1-based within window, aligns with <cls> offset
+
+        encoding = self.tokenizer(window, return_tensors="pt", truncation=True, max_length=max_len + 2)
+        input_ids = encoding["input_ids"].to(device)
+        attention_mask = encoding["attention_mask"].to(device)
+
+        seq_len = input_ids.size(1)
+        mutation_idx = torch.tensor([min(local_pos, seq_len - 1)], device=device)
+        ref_id = torch.tensor([self.tokenizer.convert_tokens_to_ids(ref_aa)], device=device)
+        alt_id = torch.tensor([self.tokenizer.convert_tokens_to_ids(alt_aa)], device=device)
+
+        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
+        logits = outputs.logits
+        logits_at_mut = logits[0, mutation_idx[0], :]
+        log_probs = torch.log_softmax(logits_at_mut, dim=-1)
+        llr = (log_probs[alt_id[0]] - log_probs[ref_id[0]]).item()
+
+        fitness = self.regression_head(torch.tensor([[llr]], device=device)).item()
+        return llr, fitness

requirements.txt

@@ -0,0 +1,7 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+fastapi==0.115.6
+uvicorn[standard]==0.34.0
+pydantic==2.10.4
+torch
+transformers==4.47.1
+huggingface_hub==0.27.0