app.py

from gradio import Blocks, Button, Checkbox, DownloadButton, Dropdown, Error, Examples, Image, HTML, Markdown, Progress, Tab, Textbox

from model import ModelFactory
from data import Data

MODELS = ModelFactory.models()
HEADER_MD = open("header.md", encoding="utf-8").read()
INSTRUCTIONS_MD = open("instructions.md", encoding="utf-8").read()

def app(seq, sub, model_name, acc):
    "Main application function"
    global progress
    scoring = "masked-marginals" if acc else "wt-marginals"

    # Validate the input
    if 1 > len(seq):
        raise Error("Sequence cannot be empty")
    if 1 > len(sub):
        raise Error("Substitutions cannot be empty")

    # Calculate the data based on the input parameters
    try:
        data = Data(seq, sub, model_name, scoring).calculate(progress)
        if isinstance(data.image, str):
            return ( Image(value=data.image, type='filepath', height=None, visible=True)
                   , HTML(value='')
                   , DownloadButton(value=data.csv, visible=True) )
        else:
            return ( Image(visible=False)
                   , HTML(value=data.image.to_html(), visible=True)
                   , DownloadButton(value=data.csv, visible=True) )
    except Exception as e:
        raise Error(str(e))
    finally:
        progress = Progress()
    
# Create the Gradio interface
with Blocks() as demo:
    Markdown("# [ESM-Scan](https://doi.org/10.1002/pro.5221)")
    # Define the interface components
    with Tab("App"):
        Markdown(HEADER_MD)
        seq = Textbox( lines=2
                     , label="Sequence"
                     , placeholder="FASTA sequence here…"
                     , value='' )
        sub = Textbox( lines=1
                     , label="Substitutions"
                     , placeholder="Substitutions here…"
                     , value='' )
        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)
        out_html    = HTML(visible=False)
        progress    = Progress()
        out_img     = Image(visible=False)
        run_btn.click( fn=lambda : (HTML(value=None, visible=False), Image(value=None, visible=True, height=128))
                     , outputs=[out_html, out_img]
                     ).then( fn=app
                           , inputs=[seq, sub, model_name, acc_box]
                           , outputs=[out_img, out_html, dl_btn] )
        ex = Examples(
            examples=[
                [   "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ"
                ,   "deep mutational scanning"
                ,   "facebook/esm2_t6_8M_UR50D" ],
                [   "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ"
                ,   "217 218 219"
                ,   "facebook/esm2_t12_35M_UR50D" ],
                [   "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ"
                ,   "R218K R218S R218N R218A R218V R218D"
                ,   "facebook/esm2_t30_150M_UR50D" ],
                [   "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ"
                ,   "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMWGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ"
                ,   "facebook/esm2_t33_650M_UR50D" ],
            ]
          , inputs=[seq, sub, model_name]
          , outputs=[out_img]
          , fn=app
          , cache_examples=False )
    with Tab("Instructions"):
        Markdown(INSTRUCTIONS_MD)

# Launch the Gradio interface
if __name__ == "__main__":
    demo.launch()