app.py

import gradio as gr
from huggingface_hub import HfApi, ModelFilter
import pandas as pd
from re import match
from tempfile import NamedTemporaryFile
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

    # fetch suitable ESM models from HuggingFace Hub
MODELS = [m.modelId for m in HfApi().list_models(filter=ModelFilter(author="facebook", model_name="esm", task="fill-mask"), sort="lastModified", direction=-1)]
if not any(MODELS):
    raise RuntimeError("Error while retrieving models from HuggingFace Hub")

    # scoring strategies
SCORING = ["masked-marginals (more accurate)", "wt-marginals (faster)"]

class Model:
    """Wrapper for ESM models"""
    def __init__(self, model_name:str=""):
        "load selected model and tokenizer"
        self.model_name = model_name
        if model_name:
            self.model = AutoModelForMaskedLM.from_pretrained(model_name)
            self.batch_converter = AutoTokenizer.from_pretrained(model_name)
            self.alphabet = self.batch_converter.get_vocab()
            if torch.cuda.is_available():
                self.model = self.model.cuda()

    def __rshift__(self, batch_tokens:torch.Tensor) -> torch.Tensor:
        "run model on batch of tokens"
        return self.model(batch_tokens)["logits"]
    
    def __lshift__(self, input:str) -> torch.Tensor:
        "convert input string to batch of tokens"
        return self.batch_converter(input, return_tensors="pt")["input_ids"]

    def __getitem__(self, key:str) -> int:
        "get token ID from character"
        return self.alphabet[key]
    
    def run_model(self, data):
        "run model on data"
        def label_row(row, token_probs):
            "label row with score"
            wt, idx, mt = row[0], int(row[1:-1])-1, row[-1]
            score = token_probs[0, 1+idx, self[mt]] - token_probs[0, 1+idx, self[wt]]
            return score.item()
        
        batch_tokens = self<<data.seq

            # run model with selected scoring strategy (info thereof available in the original ESM paper)
        if data.scoring_strategy.startswith("wt-marginals"):
            with torch.no_grad():
                token_probs = torch.log_softmax(self>>batch_tokens, dim=-1)
            data.out[self.model_name] = data.sub.apply(
                lambda row: label_row(
                    row['0'],
                    token_probs,
                ),
                axis=1,
            )
        elif data.scoring_strategy.startswith("masked-marginals"):
            all_token_probs = []
            for i in range(batch_tokens.size()[1]):
                batch_tokens_masked = batch_tokens.clone()
                batch_tokens_masked[0, i] = self['<mask>']
                with torch.no_grad():
                    token_probs = torch.log_softmax(
                        self>>batch_tokens_masked, dim=-1
                    )
                all_token_probs.append(token_probs[:, i])
            token_probs = torch.cat(all_token_probs, dim=0).unsqueeze(0)
            data.out[self.model_name] = data.sub.apply(
                lambda row: label_row(
                    row['0'],
                    token_probs,
                ),
                axis=1,
            )
    
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()
        if not all(x in self.model.alphabet for x in src):
            raise RuntimeError("Unrecognised characters in sequence")

    def parse_sub(self, trg:str):
        "parse input substitutions"
        self.mode = None
        self.sub = list()
        self.trg = trg.strip().upper()

            # identify running mode
        if len(self.trg.split()) == 1 and len(self.trg.split()[0]) == len(self.seq):    # if single string of same length as sequence, seq vs seq mode
            self.mode = 'SVS'
            for resi,(src,trg) in enumerate(zip(self.seq, self.trg), 1):
                if src != trg:
                    self.sub.append(f"{src}{resi}{trg}")
        else:
            self.trg = self.trg.split()
            if all(match(r'\d+', x) for x in self.trg):                                 # if all strings are numbers, deep mutational scanning mode
                self.mode = 'DMS'
                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}")
            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
            else:
                raise RuntimeError("Unrecognised running mode; wrong inputs?")
        
        self.sub = pd.DataFrame(self.sub, columns=['0'])

    def __init__(self, src:str, trg:str, model_name:str, scoring_strategy:str, out_file):
        "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.parse_seq(src)
        self.parse_sub(trg)
        self.scoring_strategy = scoring_strategy
        self.out = pd.DataFrame(self.sub, columns=['0', self.model_name])
        self.out_buffer = out_file.name

    def parse_output(self) -> str:
        "format output data for visualisation"
        if self.mode == 'MUT':      # if single substitution mode, sort by score
            self.out = self.out.sort_values(self.model_name, ascending=False)
        elif self.mode == 'DMS':    # if deep mutational scanning mode, sort by residue and score
            self.out = pd.concat([(self.out.assign(resi=self.out['0'].str.extract(r'(\d+)', expand=False).astype(int))      # FIX: this doesn't work if there's jolly characters in the input sequence
                                    .sort_values(['resi', self.model_name], ascending=[True,False])
                                    .groupby(['resi'])
                                    .head(19)
                                    .drop(['resi'], axis=1)).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')
            # save to temporary file to be downloaded
        self.out.round(2).to_csv(self.out_buffer, index=False)
        return (self.out.style
                        .format(lambda x: f'{x:.2f}' if isinstance(x, float) else x)
                        .hide(axis=0)
                        .hide(axis=1)
                        .background_gradient(cmap="RdYlGn", vmax=8, vmin=-8)
                        .to_html())
    
    def calculate(self):
        "run model and parse output"
        self.model.run_model(self)
        return self, self.parse_output()

def app(*argv):
    "run app"
    seq, trg, model_name, scoring_strategy, out_file, *_ = argv
    data, html = Data(seq, trg, model_name, scoring_strategy, out_file).calculate()
    return html, gr.File.update(value=out_file.name, visible=True)
    # df = pd.DataFrame((pd.np.random.random((10, 5))-0.5)*10, columns=list('ABCDE'))
    # df.to_csv(out_file.name, index=False)
    # return df.to_html(), gr.File.update(value=out_file.name, visible=True)

with gr.Blocks() as demo, NamedTemporaryFile(mode='w+', prefix='out_', suffix='.csv') as out_file, open("instructions.md", "r") as md:
    gr.Markdown(md.read())
    seq = gr.Textbox(lines=2, label="Sequence", placeholder="Sequence here...", value='MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ')
    trg = gr.Textbox(lines=1, label="Substitutions", placeholder="Substitutions here...", value="61 214 19 30 122 140")
    model_name = gr.Dropdown(MODELS, label="Model", value=MODELS[1])
    scoring_strategy = gr.Dropdown(SCORING, label="Scoring strategy", value=SCORING[1])
    btn = gr.Button(value="Submit")
    out = gr.HTML()
    bto = gr.File(value=out_file.name, visible=False, label="Download", file_count='single', interactive=False)
    btn.click(fn=app, inputs=[seq, trg, model_name, scoring_strategy, bto], outputs=[out, bto]) 

# demo.launch(share=True, server_name="0.0.0.0", server_port=7878)