Spaces:
Running
Running
| 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) |