ChatterjeeLab/zero_shot_mutation_prediction

1import gradio as gr2import pandas as pd3import torch4from transformers import AutoTokenizer, AutoModelForMaskedLM5import torch.nn.functional as F6import logging7import numpy as np8import matplotlib.pyplot as plt9import seaborn as sns10from io import BytesIO11from PIL import Image12from contextlib import contextmanager13import warnings14import sys15import os16import zipfile17 18logging.getLogger("transformers.modeling_utils").setLevel(logging.ERROR)19device = torch.device("cuda" if torch.cuda.is_available() else "cpu")20print(f"Using device: {device}")21 22# Load the tokenizer and model23model_name = "ChatterjeeLab/FusOn-pLM"24tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)25model = AutoModelForMaskedLM.from_pretrained(model_name, trust_remote_code=True)26model.to(device)27model.eval()28 29@contextmanager30def suppress_output():31    with open(os.devnull, 'w') as devnull:32        old_stdout = sys.stdout33        sys.stdout = devnull34        try:35            yield36        finally:37            sys.stdout = old_stdout38 39def process_sequence(sequence, domain_bounds, n):40    AAs_tokens = ['L', 'A', 'G', 'V', 'S', 'E', 'R', 'T', 'I', 'D', 'P', 'K', 'Q', 'N', 'F', 'Y', 'M', 'H', 'W', 'C']41    AAs_tokens_indices = {'L' : 4, 'A' : 5, 'G' : 6, 'V': 7, 'S' : 8, 'E' : 9, 'R' : 10, 'T' : 11, 'I': 12, 'D' : 13, 'P' : 14,42                          'K' : 15, 'Q' : 16, 'N' : 17, 'F' : 18, 'Y' : 19, 'M' : 20, 'H' : 21, 'W' : 22, 'C' : 23}43    # checking sequence inputs44    if not sequence.strip():45      raise gr.Error("Error: The sequence input is empty. Please enter a valid protein sequence.")46      return None, None, None47    if any(char not in AAs_tokens for char in sequence):48      raise gr.Error("Error: The sequence input contains non-amino acid characters. Please enter a valid protein sequence.")49      return None, None, None50 51    # checking domain bounds inputs52    try:53      start = int(domain_bounds['start'][0])54      end = int(domain_bounds['end'][0])55    except ValueError:56      raise gr.Error("Error: Start and end indices must be integers.")57      return None, None, None58    if start >= end:59      raise gr.Error("Start index must be smaller than end index.")60      return None, None, None61    if start == 0 and end != 0:62      raise gr.Error("Indexing starts at 1. Please enter valid domain bounds.")63      return None, None, None64    if start <= 0 or end <= 0:65      raise gr.Error("Domain bounds must be positive integers. Please enter valid domain bounds.")66      return None, None, None67    if start > len(sequence) or end > len(sequence):68      raise gr.Error("Domain bounds exceed sequence length.")69      return None, None, None70 71    # checking top n tokens input72    if n == None:73      raise gr.Error("Choose Top N Tokens from the dropdown menu.")74      return None, None, None75 76    start_index = int(domain_bounds['start'][0]) - 177    end_index = int(domain_bounds['end'][0])78 79    top_n_mutations = {}80    all_logits = []81 82    # these 2 lists are for the 2nd heatmap83    originals_logits = []84    conservation_likelihoods = {}85 86    for i in range(len(sequence)):87      # only iterate through the residues inside the domain88          if start_index <= i <= (end_index - 1):89              original_residue = sequence[i]90              original_residue_index = AAs_tokens_indices[original_residue]91              masked_seq = sequence[:i] + '<mask>' + sequence[i+1:]92              inputs = tokenizer(masked_seq, return_tensors="pt", padding=True, truncation=True, max_length=2000)93              inputs = {k: v.to(device) for k, v in inputs.items()}94              with torch.no_grad():95                  logits = model(**inputs).logits96              mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]97              mask_token_logits = logits[0, mask_token_index, :]98 99              # Pick top N tokens100              all_tokens_logits = mask_token_logits.squeeze(0)101              top_tokens_indices = torch.argsort(all_tokens_logits, dim=0, descending=True)102              top_tokens_logits = all_tokens_logits[top_tokens_indices]103              mutation = []104              # make sure we don't include non-AA tokens105              for token_index in top_tokens_indices:106                  decoded_token = tokenizer.decode([token_index.item()])107                  # decoded all tokens, pick the top n amino acid ones108                  if decoded_token in AAs_tokens:109                      mutation.append(decoded_token)110                      if len(mutation) == n:111                          break112              top_n_mutations[(sequence[i], i)] = mutation113 114              # collecting logits for the heatmap115              logits_array = mask_token_logits.cpu().numpy()116              # filter out non-amino acid tokens117              filtered_indices = list(range(4, 23 + 1))118              filtered_logits = logits_array[:, filtered_indices]119              all_logits.append(filtered_logits)120 121              # code for the second heatmap122              normalized_mask_token_logits = F.softmax(torch.tensor(mask_token_logits).cpu(), dim=-1).numpy()123              normalized_mask_token_logits = np.squeeze(normalized_mask_token_logits)124              originals_logit = normalized_mask_token_logits[original_residue_index]125              originals_logits.append(originals_logit)126 127              if originals_logit > 0.7:128                  conservation_likelihoods[(original_residue, i)] = 1129              else:130                  conservation_likelihoods[(original_residue, i)] = 0131 132 133 134   # Plotting heatmap 2135    domain_len = end - start136    if 500 > domain_len > 100:137      step_size = 50138    elif 500 <= domain_len:139      step_size = 100140    elif domain_len < 10:141      step_size = 1142    else:143      step_size = 10144    x_tick_positions = np.arange(start_index, end_index, step_size)145    x_tick_labels = [str(pos + 1) for pos in x_tick_positions]146 147    all_logits_array = np.vstack(originals_logits)148    transposed_logits_array = all_logits_array.T149    conservation_likelihoods_array = np.array(list(conservation_likelihoods.values())).reshape(1, -1)150    # combine to make a 2D heatmap151    combined_array = np.vstack((transposed_logits_array, conservation_likelihoods_array))152 153    plt.figure(figsize=(15, 5))154    plt.rcParams.update({'font.size': 16.5})155    sns.heatmap(combined_array, cmap='viridis', xticklabels=x_tick_labels, yticklabels=['Residue \nLogits', 'Residue \nConservation'], cbar=True)156    plt.xticks(x_tick_positions - start_index + 0.5, x_tick_labels, rotation=0)157    plt.title('Original Residue Probability and Conservation')158    plt.xlabel('Residue Index')159    plt.show()160    buf = BytesIO()161    plt.savefig(buf, format='png', dpi=300)162    buf.seek(0)163    plt.close()164    img_2 = Image.open(buf)165 166 167# plotting heatmap 1168    token_indices = torch.arange(logits.size(-1))169    tokens = [tokenizer.decode([idx]) for idx in token_indices]170    filtered_tokens = [tokens[i] for i in filtered_indices]171    all_logits_array = np.vstack(all_logits)172    normalized_logits_array = F.softmax(torch.tensor(all_logits_array), dim=-1).numpy()173    transposed_logits_array = normalized_logits_array.T174 175 176    plt.figure(figsize=(15, 8))177    plt.rcParams.update({'font.size': 16.5})178    sns.heatmap(transposed_logits_array, cmap='plasma', xticklabels=x_tick_labels, yticklabels=filtered_tokens)179    plt.title('Token Probability')180    plt.ylabel('Amino Acid')181    plt.xlabel('Residue Index')182    plt.yticks(rotation=0)183    plt.xticks(x_tick_positions - start_index + 0.5, x_tick_labels, rotation=0)184 185    buf = BytesIO()186    plt.savefig(buf, format='png', dpi = 300)187    buf.seek(0)188    plt.close()189 190    img_1 = Image.open(buf)191 192# store the predicted mutations in a dataframe193    original_residues = []194    mutations = []195    positions = []196 197    for key, value in top_n_mutations.items():198        original_residue, position = key199        original_residues.append(original_residue)200        mutations.append(value)201        positions.append(position + 1)202 203    df = pd.DataFrame({204        'Original Residue': original_residues,205        'Predicted Residues': mutations,206        'Position': positions207    })208    df.to_csv("predicted_tokens.csv", index=False)209    img_1.save("heatmap.png", dpi=(300, 300))210    img_2.save("heatmap_2.png", dpi=(300, 300))211    zip_path = "outputs.zip"212    with zipfile.ZipFile(zip_path, 'w') as zipf:213        zipf.write("predicted_tokens.csv")214        zipf.write("heatmap.png")215        zipf.write("heatmap_2.png")216 217    return df, img_1, img_2, zip_path218 219# launch the demo220demo = gr.Interface(221    fn=process_sequence,222    inputs=[223        gr.Textbox(label="Sequence", placeholder="Enter the protein sequence here"),224        gr.Dataframe(225            value = [[1, 1]],226            headers=["start", "end"],227            datatype=["number", "number"],228            row_count=(1, "fixed"),229            col_count=(2, "fixed"),230            label="Domain Bounds"231        ),232        gr.Dropdown([i for i in range(1, 21)], label="Top N Tokens"),233    ],234     outputs=[235        gr.Dataframe(label="Predicted Tokens (in order of decreasing likelihood)"),236        gr.Image(type="pil", label="Probability Distribution for All Tokens"),237        gr.Image(type="pil", label="Residue Conservation"),238        gr.File(label="Download Outputs"),239    ],240)241if __name__ == "__main__":242    with suppress_output():243      demo.launch()