ESMC integration

AGENTS.md

@@ -0,0 +1,55 @@
+# ZSP — Agent Guide
+
+**What it is.** Gradio web app that scores protein sequence variants using ESM2, E1, and ESMC language models. Implements ESM-Scan methodology ([doi:10.1002/pro.5221](https://doi.org/10.1002/pro.5221)).
+
+## Structure
+
+```
+app.py        — Gradio UI, entry point
+model.py      — ModelFactory + wrappers for ESM/E1/ESMC models
+data.py       — Input parsing, scoring modes, output rendering (heatmap/table)
+requirements.txt — deps (gradio, torch, esm[Biohub], pandas, seaborn, E1)
+header.md     — UI header text
+instructions.md — user instructions tab content
+out.*         — generated outputs (gitignored)
+```
+
+## Run
+
+```bash
+source .venv/bin/activate   # or however you activate
+python app.py
+```
+
+Deploys as a HuggingFace Space (`sdk: gradio`, `app_file: app.py`). Python 3.12.
+
+## Scoring Modes
+
+Parsed from the substitutions input in `data.py`:
+
+| Mode | Trigger | What it does |
+|------|---------|-------------|
+| **TMS** (triple mutational scan) | No match to other patterns | All-positions × all-amino-acids sweep → heatmap image |
+| **DMS** (deep mutational scan) | All tokens are integers | Specified positions × all alternatives → sorted table |
+| **MUT** (single mutations) | Tokens match `[A-Z]\d+[A-Z]` | Explicit substitution list → scored table |
+| **MUT-seq-vs-seq** | Single token same length as sequence | Align two sequences, score differences → scored table |
+
+Scoring strategy flag: `"masked-marginals"` (higher accuracy, slower) vs `"wt-marginals"`.
+
+## Key Design Notes
+
+- Models loaded lazily via `ModelFactory`; ESM and ESMC use shared `ESMModel` wrapper, E1 uses its own package (`E1@git+...`)
+- ESMC requires `esm@git+Biohub/esm` — registers custom `esmc` architecture with transformers fork
+- GPU auto-detected; falls back to CPU
+- E1 model chunks batches at 60 to avoid OOM
+- TMS mode produces a seaborn heatmap PNG; other modes produce styled HTML tables
+- Output files always written to `out.csv` / `out.png` in CWD
+- Progress tracking wired into Gradio's `Progress` object
+
+## Edit Rules
+
+- Match existing style: one-line docstrings on classes/methods, minimal comments
+- No new dependencies unless absolutely necessary — stdlib or already-installed first
+- Torch/transformers signatures: grep before guessing versions
+- Output format changes affect both CSV and image/table paths — update both
+- ModelFactory registry is class-level `_models` dict; adding a model type means registering there

app.py

@@ -47,7 +47,7 @@ with Blocks() as demo:
                      , label="Substitutions"
                      , placeholder="Substitutions here…"
                      , value='' )
-        model_name  = Dropdown(MODELS, label="Model", value="facebook/esm2_t30_150M_UR50D")
+        model_name  = Dropdown(MODELS, label="Model", value="biohub/ESMC-300M")
         acc_box     = Checkbox(value=True, label="Use higher accuracy scoring", interactive=True)
         run_btn     = Button(value="Run", variant="primary")
         dl_btn      = DownloadButton(label="Download raw data", visible=False)

header.md

@@ -1,7 +1,7 @@
-Calculate the fitness of single amino acid substitutions on proteins, using a [zero-shot](https://doi.org/10.1101/2021.07.09.450648) protein language predictor of the [ESM model family](https://huggingface.co/facebook/esm2_t6_8M_UR50D).
+Calculate the fitness of single amino acid substitutions on proteins, using zero-shot protein language models from the [ESM](https://huggingface.co/facebook) and [ESMC](https://huggingface.co/collections/biohub/esmc-model-family) families.
 
-**UPDATE:** 
-[Profluent-Bio](https://huggingface.co/Profluent-Bio)'s [E1 model family](https://huggingface.co/Profluent-Bio/E1-150m) is now available for inference.
+**UPDATE:**
+[Biohub](https://biohub.ai)'s [ESMC model family](https://huggingface.co/collections/biohub/esmc-model-family) (trained on UniRef + MGnify + JGI) is now available alongside ESM2 models.
 
 **WARNING:** 
 Due to high server traffic, the tool might become slow or unresponsive.

instructions.md

@@ -16,8 +16,10 @@ Note: While jolly characters (e.g., `-X.B`) can be included, they currently cann
 - **Same-Length Sequence**: Analyze differing amino acid substitutions one by one within sequences of equal length.
 - **Different Inputs**: For any other input format, a deep mutational scan of the full sequence will be performed.
 
-**Model Selection**: Choose a model for calculations from those available on Hugging Face Model Hub.  
-The `esm2_t33_650M_UR50D` model offers an optimal balance between cost and accuracy [*](https://doi.org/10.1126/science.ade2574).
+**Model Selection**: Choose a model for calculations from those available on Hugging Face Model Hub. Two model families are supported:
+
+- **Facebook ESM2** — trained on UniRef50. `esm2_t33_650M_UR50D` offers an optimal balance between cost and accuracy [*](https://doi.org/10.1126/science.ade2574).
+- **Biohub ESMC** — trained on UniRef + MGnify + JGI (broader metagenomic coverage). Pre-LN architecture with rotary embeddings, max 2048 token context. Rankings correlate strongly with ESM2 (Spearman ρ ≥ 0.93) while absolute scores may differ due to training data differences.
 
 **Accuracy Option**: The **Use higher accuracy** option applies a masked-marginals scoring strategy, which considers sequence context during inference.  
 While this method is slower, it enhances accuracy. 

model.py

@@ -4,11 +4,9 @@ from typing import Any
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 from transformers.tokenization_utils_base import BatchEncoding
 from transformers.modeling_outputs import MaskedLMOutput
-from E1.modeling import E1ForMaskedLM
-from E1.scorer import E1Scorer, EncoderScoreMethod
 
 class ESMModel:
-    "Wrapper for ESM models"
+    "Wrapper for ESM/ESMC models"
     def __init__(self, model_name:str):
         "Load selected model and tokenizer"
         self.model_name = model_name
@@ -75,29 +73,6 @@ class ESMModel:
             lambda row: label_row(row['0'], token_probs)
           , axis=1 )
 
-class E1Model:
-    "Wrapper for E1 models"
-    def __init__(self, model_name:str):
-        self.model_name = model_name
-        self.scoring_strategy = EncoderScoreMethod.MASKED_MARGINAL
-        if model_name:            
-            self.model = E1ForMaskedLM.from_pretrained(model_name, dtype=torch.float)
-            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            self.model = self.model.to(self.device)
-            self.scorer = E1Scorer(self.model, self.scoring_strategy)
-            self.alphabet = self.scorer.vocab
-
-    def run_model(self, data):
-        "Run model on data"
-        if not data.scoring_strategy.startswith("masked-marginals"):
-            self.scorer = E1Scorer(self.model, EncoderScoreMethod.WILDTYPE_MARGINAL)
-        batch_size = 60 ## chunking to avoid OOM
-        out = []
-        for chunk in data.progress.tqdm([data.trg[i:i+batch_size] for i in range(0, len(data.trg), batch_size)], desc="Calculating"):
-            scores = self.scorer.score(parent_sequence=data.seq, sequences=chunk)
-            out.extend(s['score'] for s in scores)
-        data.out[self.model_name] = out
-
 class ModelFactory:
     "Factory for models"
     _models = { **{ m.id:ESMModel for m in HfApi().list_models(
@@ -106,11 +81,8 @@ class ModelFactory:
                     , filter="fill-mask" 
                     , sort="lastModified"
                     , direction=-1 ) }
-            ,  **{ m.id:E1Model for m in HfApi().list_models(
-                      author="Profluent-Bio"
-                    , model_name="E1"
-                    , sort="lastModified"
-                    , direction=-1 ) } }
+            ,  **{ m.id:ESMModel for m in HfApi().list_models(
+                      author="biohub", model_name="esmc", filter="fill-mask" ) if "-step" not in m.id and "-2024-12" not in m.id } }
     
     @classmethod
     def register(cls, model_name, model_cls):

requirements.txt

@@ -2,5 +2,4 @@ gradio
 pandas
 seaborn
 torch
-transformers
-E1@git+https://github.com/Profluent-AI/E1.git@main#egg=E1
+esm@git+https://github.com/Biohub/esm.git@main