E1 model addition

app.py

@@ -1,13 +1,13 @@
-from gradio import Blocks, Button, Checkbox, DataFrame, DownloadButton, Dropdown, Examples, Image, Markdown, Tab, Textbox
+from gradio import Blocks, Button, Checkbox, DataFrame, DownloadButton, Dropdown, Error, Examples, Image, Markdown, Tab, Textbox
 
-from model import get_models
+from model import ModelFactory
 from data import Data
 
 # Define scoring strategies
 SCORING = ["wt-marginals", "masked-marginals"]
 
 # Get available models
-MODELS = get_models()
+MODELS = ModelFactory.models()
 
 def app(*argv):
     """
@@ -17,12 +17,15 @@ def app(*argv):
     seq, trg, model_name, *_ = argv
     scoring = SCORING[scoring_strategy.value]
     # Calculate the data based on the input parameters
-    data = Data(seq, trg, model_name, scoring).calculate()
-
-    if isinstance(data.image(), str):
-        out = Image(value=data.image(), type='filepath', visible=True), DataFrame(visible=False)
-    else:
-        out = Image(visible=False), DataFrame(value=data.image(), visible=True)
+    try:
+        data = Data(seq, trg, model_name, scoring).calculate()
+        if isinstance(data.image(), str):
+            out = Image(value=data.image(), type='filepath', visible=True), DataFrame(visible=False)
+        else:
+            out = Image(visible=False), DataFrame(value=data.image(), visible=True)
+    except Exception as e:
+        out = Image(visible=False), DataFrame(visible=False)
+        raise Error(str(e))
         
     return *out, DownloadButton(value=data.csv(), visible=True)
 
@@ -32,58 +35,52 @@ with Blocks() as esm_scan:
     # Define the interface components
     with Tab("App"):
         Markdown(open("header.md", "r", encoding="utf-8").read())
-        seq = Textbox(
-            lines=2,
-            label="Sequence",
-            placeholder="FASTA sequence here...",
-            value=''
-        )
-        trg = Textbox(
-            lines=1,
-            label="Substitutions",
-            placeholder="Substitutions here...",
-            value=""
-        )
+        seq = Textbox( lines=2
+                     , label="Sequence"
+                     , placeholder="FASTA sequence here..."
+                     , value=''
+                     )
+        trg = Textbox( lines=1
+                     , label="Substitutions"
+                     , placeholder="Substitutions here..."
+                     , value=""
+                     )
         model_name = Dropdown(MODELS, label="Model", value="facebook/esm2_t30_150M_UR50D")
         scoring_strategy = Checkbox(value=True, label="Use higher accuracy scoring", interactive=True)
         btn = Button(value="Run", variant="primary")
         dlb = DownloadButton(label="Download raw data", visible=False)
         out = Image(visible=False)
         ouu = DataFrame(visible=False)
-        btn.click(
-            fn=app,
-            inputs=[seq, trg, model_name],
-            outputs=[out, ouu, dlb]
-        )
+        btn.click( fn=app
+                 , inputs=[seq, trg, model_name]
+                 , outputs=[out, ouu, dlb]
+                 )
         ex = Examples(
             examples=[
-                [
-                    "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ", 
-                    "deep mutational scanning",
-                    "facebook/esm2_t6_8M_UR50D"
+                [   "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ"
+                ,   "deep mutational scanning"
+                ,   "facebook/esm2_t6_8M_UR50D"
                 ],
-                [
-                    "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ", 
-                    "217 218 219",
-                    "facebook/esm2_t12_35M_UR50D"
+                [   "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ"
+                ,   "217 218 219"
+                ,   "facebook/esm2_t12_35M_UR50D"
                 ],
-                [
-                    "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ", 
-                    "R218K R218S R218N R218A R218V R218D",
-                    "facebook/esm2_t30_150M_UR50D",
+                [   "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ"
+                ,   "R218K R218S R218N R218A R218V R218D"
+                ,   "facebook/esm2_t30_150M_UR50D"
                 ],
-                [
-                    "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ", 
-                    "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMWGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ", 
-                    "facebook/esm2_t33_650M_UR50D",
+                [   "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ"
+                ,   "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMWGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ"
+                ,   "facebook/esm2_t33_650M_UR50D"
                 ],
-            ],
-            inputs=[seq,
-                    trg,
-                    model_name],
-            outputs=[out],
-            fn=app,
-            cache_examples=False
+            ]
+          , inputs=[ seq
+                   , trg
+                   , model_name
+                   ]
+          , outputs=[out]
+          , fn=app
+          , cache_examples=False
         )
     with Tab("Instructions"):
         Markdown(open("instructions.md", "r", encoding="utf-8").read())

data.py

@@ -4,18 +4,15 @@ import pandas as pd
 from re import match
 import seaborn as sns
 
-from model import Model
+from model import ModelFactory
 
 class Data:
     """Container for input and output data"""
-    # Initialise empty model as static class member for efficiency
-    model = Model()
-
     def parse_seq(self, src: str):
         """Parse input sequence"""
         self.seq = src.strip().upper().replace('\n', '')
         if not all(x in self.model.alphabet for x in self.seq):
-            raise RuntimeError("Unrecognised characters in sequence")
+            raise RuntimeError(f"Unsupported characters in sequence: {''.join(x for x in self.seq if x not in self.model.alphabet)}")
 
     def parse_sub(self, trg: str):
         """Parse input substitutions"""
@@ -36,34 +33,42 @@ class Data:
             if all(match(r'\d+', x) for x in self.trg):                         
                 # If all strings are numbers, deep mutational scanning mode
                 self.mode = 'DMS'
+                trh = list()
                 for resi in map(int, self.trg):
                     src = self.seq[resi-1]
                     for trg in "ACDEFGHIKLMNPQRSTVWY".replace(src, ''):
                         self.sub.append(f"{src}{resi}{trg}")
+                        trh.append(self.seq[:resi-1]+trg+self.seq[resi:])
                     self.resi.append(resi)
+                self.trg = trh
             elif all(match(r'[A-Z]\d+[A-Z]', x) for x in self.trg):
                 # If all strings are of the form X#Y, single substitution mode
                 self.mode = 'MUT'
                 self.sub = self.trg
-                self.resi = [int(x[1:-1]) for x in self.trg]
+                trh = list()
+                for x in self.trg:
+                    idx = int(x[1:-1])
+                    self.resi.append(idx)
+                    trh.append(self.seq[:idx-1]+x[-1]+self.seq[idx:])
                 for s, *resi, _ in self.trg:
                     if self.seq[int(''.join(resi))-1] != s:
-                        raise RuntimeError(f"Unrecognised input substitution {self.seq[int(''.join(resi))]}{int(''.join(resi))} /= {s}{int(''.join(resi))}")
+                        raise RuntimeError(f"Unrecognised input substitution: {self.seq[int(''.join(resi))]}{int(''.join(resi))} /= {s}{int(''.join(resi))}")
+                self.trg = trh
             else:
                 self.mode = 'TMS'
+                self.trg = list()
                 for resi, src in enumerate(self.seq, 1):
                     for trg in "ACDEFGHIKLMNPQRSTVWY".replace(src, ''):
                         self.sub.append(f"{src}{resi}{trg}")
+                        self.trg.append(self.seq[:resi-1]+trg+self.seq[resi:])
                     self.resi.append(resi)
 
         self.sub = pd.DataFrame(self.sub, columns=['0'])
 
     def __init__(self, src:str, trg:str, model_name:str='facebook/esm2_t33_650M_UR50D', scoring_strategy:str='masked-marginals', out_file='out'):
         "initialise data"
-            # if model has changed, load new model
-        if self.model.model_name != model_name:
-            self.model_name = model_name
-            self.model = Model(model_name)
+        self.model_name = model_name
+        self.model = ModelFactory(model_name)
         self.parse_seq(src)
         self.offset = 0
         self.parse_sub(trg)
@@ -101,8 +106,8 @@ class Data:
                             .groupby(['resi'])
                             .head(19)
                             .drop(['resi'], axis=1))
-        self.out = pd.concat([self.out.iloc[19*x:19*(x+1)].reset_index(drop=True) for x in range(self.out.shape[0]//19)]
-                            , axis=1).set_axis(range(self.out.shape[0]//19*2), axis='columns')
+        self.out = pd.concat([ self.out.iloc[19*x:19*(x+1)].reset_index(drop=True) for x in range(self.out.shape[0]//19)]
+                             , axis=1).set_axis(range(self.out.shape[0]//19*2), axis='columns')
 
     def process_tms_mode(self):
         self.out = self.assign_resi_and_group()
@@ -124,8 +129,8 @@ class Data:
                             .pipe(self.create_dataframe)
                             .sort_values(['0'], ascending=[True])
                             .drop(['resi', '0'], axis=1)
-                            .set_axis(['A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L',
-                                       'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y'])
+                            .set_axis([ 'A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L'
+                                      , 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y'])
                             .astype(float)
                             ) for x in range(self.out.shape[0]//19)]
                         , axis=1)
@@ -152,16 +157,16 @@ class Data:
 
     def plot_single_heatmap(self):
         fig = plt.figure(figsize=(12, 6))
-        sns.heatmap(self.out
-                  , cmap='RdBu'
-                  , cbar=False
-                  , square=True
-                  , xticklabels=1
-                  , yticklabels=1
-                  , center=0
-                  , annot=self.out.map(lambda x: ' ' if x != 0 else '·')
-                  , fmt='s'
-                  , annot_kws={'size': 'xx-large'})
+        sns.heatmap( self.out
+                   , cmap='RdBu'
+                   , cbar=False
+                   , square=True
+                   , xticklabels=1
+                   , yticklabels=1
+                   , center=0
+                   , annot=self.out.map(lambda x: ' ' if x != 0 else '·')
+                   , fmt='s'
+                   , annot_kws={'size': 'xx-large'})
         fig.tight_layout()
 
     def plot_multiple_heatmaps(self, ncols, nrows):
@@ -169,17 +174,17 @@ class Data:
         for i in range(nrows):
             tmp = self.out.iloc[:,i*ncols:(i+1)*ncols]
             label = tmp.map(lambda x: ' ' if x != 0 else '·')
-            sns.heatmap(tmp
-                      , ax=ax[i]
-                      , cmap='RdBu'
-                      , cbar=False
-                      , square=True
-                      , xticklabels=1
-                      , yticklabels=1
-                      , center=0
-                      , annot=label
-                      , fmt='s'
-                      , annot_kws={'size': 'xx-large'})
+            sns.heatmap( tmp
+                       , ax=ax[i]
+                       , cmap='RdBu'
+                       , cbar=False
+                       , square=True
+                       , xticklabels=1
+                       , yticklabels=1
+                       , center=0
+                       , annot=label
+                       , fmt='s'
+                       , annot_kws={'size': 'xx-large'})
             ax[i].set_yticklabels(ax[i].get_yticklabels(), rotation=0)
             ax[i].set_xticklabels(ax[i].get_xticklabels(), rotation=90)
         fig.tight_layout()

model.py

@@ -1,31 +1,17 @@
 from huggingface_hub import HfApi
 import torch
 from tqdm import tqdm
+from typing import Any
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 from transformers.tokenization_utils_base import BatchEncoding
 from transformers.modeling_outputs import MaskedLMOutput
-
-# Function to fetch suitable ESM models from HuggingFace Hub
-def get_models() -> list[None|str]:
-    """Fetch suitable ESM models from HuggingFace Hub."""
-    if not any(
-        out := [
-            m.modelId for m in HfApi().list_models(
-                author="facebook", 
-                model_name="esm", 
-                task="fill-mask", 
-                sort="lastModified", 
-                direction=-1
-            )
-        ]
-    ):
-        raise RuntimeError("Error while retrieving models from HuggingFace Hub")
-    return out
+from E1.modeling import E1ForMaskedLM
+from E1.scorer import E1Scorer, EncoderScoreMethod
 
 # Class to wrap ESM models
-class Model:
+class ESMModel:
     """Wrapper for ESM models."""
-    def __init__(self, model_name: str = ""):
+    def __init__(self, model_name:str):
         """Load selected model and tokenizer."""
         self.model_name = model_name
         if model_name:
@@ -95,9 +81,63 @@ class Model:
 
         # Apply the label_row function to each row of the substitutions dataframe
         data.out[self.model_name] = data.sub.apply(
-            lambda row: label_row(
-                row['0'],
-                token_probs,
-            ),
-            axis=1,
-        )
+            lambda row: label_row( 
+                row['0']
+              , token_probs
+              )
+          , axis=1
+        )
+
+# Class to wrap E1 models
+class E1Model:
+    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)
+            if torch.cuda.is_available():
+                self.model = self.model.cuda()
+                self.device = torch.device("cuda")
+            else:
+                self.device = torch.device("cpu")
+            self.scorer = E1Scorer(self.model, method=self.scoring_strategy)
+            self.alphabet = self.scorer.vocab
+
+    def run_model(self, data):
+        if not data.scoring_strategy.startswith("masked-marginals"):
+            self.scorer = E1Scorer(self.model, method=EncoderScoreMethod.WILDTYPE_MARGINAL)
+        scores = self.scorer.score(parent_sequence=data.seq, sequences=data.trg)
+        data.out[self.model_name] = [s['score'] for s in scores]
+
+class ModelFactory:
+    _models = {
+                **{ m.modelId:ESMModel for m in HfApi().list_models(
+                    author="facebook"
+                    , model_name="esm"
+                    , filter="fill-mask" 
+                    , sort="lastModified"
+                    , direction=-1
+                    )
+                }
+            , **{ m.modelId:E1Model for m in HfApi().list_models(
+                    author="Profluent-Bio"
+                    , model_name="E1"
+                    , sort="lastModified"
+                    , direction=-1
+                    )
+                }
+              }
+    
+    @classmethod
+    def register(cls, model_name, model_cls):
+        cls._models[model_name] = model_cls
+
+    @classmethod
+    def models(cls):
+        return [m for m in cls._models.keys()]
+
+    def __new__(cls, model_name:str) -> Any:
+        return cls._models[model_name](model_name)
+
+for m,c in ModelFactory._models.items():
+    ModelFactory.register(m, c)

requirements.txt

@@ -3,3 +3,4 @@ pandas
 seaborn
 torch
 transformers
+E1@git+https://github.com/Profluent-AI/E1.git@main#egg=E1