add application file

app.py

@@ -0,0 +1,218 @@
+from huggingface_hub import HfApi, ModelFilter
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+import pandas as pd
+import re
+from tqdm import tqdm
+import torch
+import gradio as gr
+import warnings
+warnings.filterwarnings('ignore')
+
+MODEL, MODEL_NAME, BATCH_CONVERTER, ALPHABET = None, None, None, None
+OFFSET = 1
+MODELS = [m.modelId for m in HfApi().list_models(filter=ModelFilter(author="facebook", model_name="esm", task="fill-mask"), sort="lastModified", direction=-1)]
+SCORING = ["masked-marginals (more accurate)", "wt-marginals (faster)"]
+
+def label_row(row, sequence, token_probs):
+    wt, idx, mt = row[0], int(row[1:-1]) - OFFSET, row[-1]
+    assert sequence[idx] == wt, "The listed wildtype does not match the provided sequence"
+
+    wt_encoded, mt_encoded = ALPHABET[wt], ALPHABET[mt]
+
+    score = token_probs[0, 1 + idx, mt_encoded] - token_probs[0, 1 + idx, wt_encoded]
+    return score.item()
+
+def initialise_model(model_name):
+    global MODEL, MODEL_NAME, BATCH_CONVERTER, ALPHABET
+    MODEL_NAME = model_name
+    MODEL = AutoModelForMaskedLM.from_pretrained(model_name)
+    BATCH_CONVERTER = AutoTokenizer.from_pretrained(model_name)
+    ALPHABET = BATCH_CONVERTER.get_vocab()
+    if torch.cuda.is_available():
+        MODEL = MODEL.cuda()
+
+def parse_input(seq, sub):
+    assert seq.isalpha(), "Sequence must be alphabetic"
+    substitutions, mode = list(), None
+
+    if len(sub.split()) == 1 and len(sub.split()[0]) == len(seq):
+        mode = 'seq vs seq'
+        for resi,(src,trg) in enumerate(zip(seq,sub), OFFSET):
+            if src != trg:
+                substitutions.append(f"{src}{resi}{trg}")
+    elif len(targets := sub.split()) > 1:
+        if all(re.match(r'\d+', x) for x in targets):
+            mode = 'deep mutational scan'
+            for resi in map(int, sub.split()):
+                src = seq[resi-OFFSET]
+                for trg in "ACDEFGHIKLMNPQRSTVWY".replace(src,''):
+                    substitutions.append(f"{src}{resi}{trg}")
+        elif all(re.match(r'[A-Z]\d+[A-Z]', x) for x in targets):
+            mode = 'aa substitutions'
+            substitutions = targets
+    
+    if not mode:
+        raise RuntimeError("Unrecognised running mode")
+    
+    return mode, pd.DataFrame(substitutions, columns=['0'])
+
+def run_model(sequence, substitutions, batch_tokens, scoring_strategy):
+    if scoring_strategy.startswith("wt-marginals"):
+        with torch.no_grad():
+            token_probs = torch.log_softmax(MODEL(batch_tokens)["logits"], dim=-1)
+        substitutions[MODEL_NAME] = substitutions.apply(
+            lambda row: label_row(
+                row['0'],
+                sequence,
+                token_probs,
+            ),
+            axis=1,
+        )
+    elif scoring_strategy.startswith("masked-marginals"):
+        all_token_probs = []
+        for i in tqdm(range(batch_tokens.size()[1])):
+            batch_tokens_masked = batch_tokens.clone()
+            batch_tokens_masked[0, i] = ALPHABET['<mask>']
+            with torch.no_grad():
+                token_probs = torch.log_softmax(
+                    MODEL(batch_tokens_masked)["logits"], dim=-1
+                )
+            all_token_probs.append(token_probs[:, i])
+        token_probs = torch.cat(all_token_probs, dim=0).unsqueeze(0)
+        substitutions[MODEL_NAME] = substitutions.apply(
+            lambda row: label_row(
+                row['0'],
+                sequence,
+                token_probs,
+            ),
+            axis=1,
+        )
+
+    return substitutions
+
+def parse_output(output, mode):
+    if mode == 'aa substitutions':
+        output = output.sort_values(MODEL_NAME, ascending=False)
+    elif mode == 'deep mutational scan':
+        output = pd.concat([(output.assign(resi=output['0'].str.extract(r'(\d+)', expand=False).astype(int))
+                .sort_values(['resi', 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(output.shape[0]//19)]
+                , axis=1).set_axis(range(output.shape[0]//19*2), axis='columns')
+
+    return output.style.format(lambda x: f'{x:.2f}' if isinstance(x, float) else x).hide_index().hide_columns().background_gradient(cmap="RdYlGn", vmax=8, vmin=-8).to_html()
+
+    
+# mode = 'deep mutational scan'  #@param ['seq vs seq', 'deep mutational scan', 'aa substitutions']
+# sequence = "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ" #@param {type:"string"}
+# target = "61 214 19 30 122 140" #@param {type:"string"}
+# substitutions = list()
+# scoring_strategy = "masked-marginals"
+
+# if mode == 'seq vs seq':
+#   for resi,(seq,trg) in enumerate(zip(sequence,target), OFFSET):
+#     if seq != trg:
+#       substitutions.append(f"{seq}{resi}{trg}")
+# elif mode == 'deep mutational scan':
+#   for resi in map(int, target.split()):
+#     seq = sequence[resi-OFFSET]
+#     for trg in "ACDEFGHIKLMNPQRSTVWY".replace(seq,''):
+#       substitutions.append(f"{seq}{resi}{trg}")
+# elif mode == 'aa substitutions':
+#   substitutions = target.split()
+# else:
+#   raise RuntimeError("Unrecognised running mode")
+
+# df = pd.DataFrame(substitutions, columns=['0'])
+# mutation_col = df.columns[0]
+
+# batch_tokens = batch_converter(sequence, return_tensors='pt')['input_ids']
+
+# if scoring_strategy == "wt-marginals":
+#     with torch.no_grad():
+#         token_probs = torch.log_softmax(model(batch_tokens)["logits"], dim=-1)
+#     df[model_name] = df.apply(
+#         lambda row: label_row(
+#             row[mutation_col],
+#             sequence,
+#             token_probs,
+#             alphabet,
+#             OFFSET,
+#         ),
+#         axis=1,
+#     )
+# elif scoring_strategy == "masked-marginals":
+#     all_token_probs = []
+#     for i in tqdm(range(batch_tokens.size()[1])):
+#         batch_tokens_masked = batch_tokens.clone()
+#         batch_tokens_masked[0, i] = alphabet['<mask>']
+#         with torch.no_grad():
+#             token_probs = torch.log_softmax(
+#                 model(batch_tokens_masked)["logits"], dim=-1
+#             )
+#         all_token_probs.append(token_probs[:, i])  # vocab size
+#     token_probs = torch.cat(all_token_probs, dim=0).unsqueeze(0)
+#     df[model_name] = df.apply(
+#         lambda row: label_row(
+#             row[mutation_col],
+#             sequence,
+#             token_probs,
+#             alphabet,
+#             OFFSET,
+#         ),
+#         axis=1,
+#     )
+
+# if mode == 'aa substitutions':
+#   df = df.sort_values(model_name, ascending=False)
+# elif mode == 'deep mutational scan':
+#   df = pd.concat([(df.assign(resi=df['0'].str.extract(f'(\d+)', expand=False).astype(int))
+#           .sort_values(['resi', 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(df.shape[0]//19)]
+#         , axis=1).set_axis(range(df.shape[0]//19*2), axis='columns')
+
+# df.style.hide_index().hide_columns().background_gradient(cmap="RdYlGn", vmax=8, vmin=-8)
+
+def app(*argv):
+    seq, trg, model_name, scoring_strategy, *_ = argv
+
+    mode, substitutions = parse_input(seq, trg)
+
+    if model_name != MODEL_NAME:
+        initialise_model(model_name)
+
+    batch_tokens = BATCH_CONVERTER(seq, return_tensors='pt')['input_ids']
+
+    df = run_model(seq, substitutions, batch_tokens, scoring_strategy)
+
+    return parse_output(df, mode)
+
+# demo = gr.Interface(
+#     theme=gr.themes.Base(),
+#     title="Protein Sequence Mutagenesis",
+#     description="Predict the effect of mutations on protein stability",
+#     fn=app,
+#     inputs=[gr.Textbox(lines=2, label="Sequence", placeholder="Sequence here...", required=True, value='MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ'),
+#             gr.Textbox(lines=2, label="Substitutions", placeholder="Substitutions here...", required=True, value="61 214 19 30 122 140"),
+#             gr.Dropdown(MODELS, label="Model", value=MODELS[1]),
+#             gr.Dropdown(["masked-marginals (more accurate)", "wt-marginals (faster)"], label="Scoring strategy", value="wt-marginals (faster)"),
+#             ],
+#     outputs=gr.HTML(formatter="html", label="Output"),
+# )
+
+with gr.Blocks() as demo:
+    gr.Markdown("""Protein Sequence Mutagenesis""", name="title")
+    gr.Markdown("""Predict the effect of mutations on protein stability""", name="description")
+    seq = gr.Textbox(lines=2, label="Sequence", placeholder="Sequence here...", required=True, value='MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ')
+    trg = gr.Textbox(lines=1, label="Substitutions", placeholder="Substitutions here...", required=True, 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(label="Submit", type="submit")
+    btn.click(fn=app, inputs=[seq, trg, model_name, scoring_strategy], outputs=[gr.HTML()])
+
+if __name__ == '__main__':
+    demo.launch()
+    # demo.launch(share=True, server_name="0.0.0.0", server_port=7878)