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    # checking sequence inputs42    if not sequence.strip():43      raise gr.Error("Error: The sequence input is empty. Please enter a valid protein sequence.") 44      return None, None, None45    if any(char not in AAs_tokens for char in sequence):46      raise gr.Error("Error: The sequence input contains non-amino acid characters. Please enter a valid protein sequence.") 47      return None, None, None48 49    # checking domain bounds inputs50    try:51      start = int(domain_bounds['start'][0])52      end = int(domain_bounds['end'][0])53    except ValueError:54      raise gr.Error("Error: Start and end indices must be integers.")55      return None, None, None56    if start >= end:57      raise gr.Error("Start index must be smaller than end index.")58      return None, None, None59    if start == 0 and end != 0:60      raise gr.Error("Indexing starts at 1. Please enter valid domain bounds.")61      return None, None, None62    if start <= 0 or end <= 0:63      raise gr.Error("Domain bounds must be positive integers. Please enter valid domain bounds.")64      return None, None, None65    if start > len(sequence) or end > len(sequence):66      raise gr.Error("Domain bounds exceed sequence length.")67      return None, None, None68 69    # checking n inputs70    if n == None:71      raise gr.Error("Choose Top N Tokens from the dropdown menu.")72      return None, None, None73 74    start_index = int(domain_bounds['start'][0]) - 175    end_index = int(domain_bounds['end'][0])76 77    top_n_mutations = {}78    all_logits = []79 80    for i in range(len(sequence)):81          if start_index <= i <= (end_index - 1):82              masked_seq = sequence[:i] + '<mask>' + sequence[i+1:]83              inputs = tokenizer(masked_seq, return_tensors="pt", padding=True, truncation=True, max_length=2000)84              inputs = {k: v.to(device) for k, v in inputs.items()}85              with torch.no_grad():86                  logits = model(**inputs).logits87              mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]88              mask_token_logits = logits[0, mask_token_index, :]89 90              # Define amino acid tokens91              all_tokens_logits = mask_token_logits.squeeze(0)92              top_tokens_indices = torch.argsort(all_tokens_logits, dim=0, descending=True)93              top_tokens_logits = all_tokens_logits[top_tokens_indices]94              mutation = []95              # make sure we don't include non-AA tokens96              for token_index in top_tokens_indices:97                  decoded_token = tokenizer.decode([token_index.item()])98                  if decoded_token in AAs_tokens:99                      mutation.append(decoded_token)100                      if len(mutation) == n:101                          break102              top_n_mutations[(sequence[i], i)] = mutation103 104              # collecting logits for the heatmap105              logits_array = mask_token_logits.cpu().numpy()106              # filter out non-amino acid tokens107              filtered_indices = list(range(4, 23 + 1))108              filtered_logits = logits_array[:, filtered_indices]109              all_logits.append(filtered_logits)110 111 112    token_indices = torch.arange(logits.size(-1))113    tokens = [tokenizer.decode([idx]) for idx in token_indices]114    filtered_tokens = [tokens[i] for i in filtered_indices]115 116    all_logits_array = np.vstack(all_logits)117    normalized_logits_array = F.softmax(torch.tensor(all_logits_array), dim=-1).numpy()118    transposed_logits_array = normalized_logits_array.T119 120    # Plotting the heatmap121    domain_len = end - start122    if 500 > domain_len > 100:123      step_size = 49124    elif 500 <= domain_len:125      step_size = 99126    elif domain_len < 10:127      step_size = 1128    else:129      step_size = 9130    x_tick_positions = np.arange(start_index, end_index, step_size)131    x_tick_labels = [str(pos + 1) for pos in x_tick_positions]132 133    plt.figure(figsize=(15, 8))134    plt.rcParams.update({'font.size': 18})135 136    sns.heatmap(transposed_logits_array, cmap='plasma', xticklabels=x_tick_labels, yticklabels=filtered_tokens)137    plt.title('Token Probability Heatmap')138    plt.ylabel('Token')139    plt.xlabel('Residue Index')140    plt.yticks(rotation=0)141    plt.xticks(x_tick_positions - start_index + 0.5, x_tick_labels, rotation=0)142 143    # Save the figure to a BytesIO object144    buf = BytesIO()145    plt.savefig(buf, format='png', dpi = 300)146    buf.seek(0)147    plt.close()148 149    # Convert BytesIO object to an image150    img = Image.open(buf)151 152    original_residues = []153    mutations = []154    positions = []155 156    for key, value in top_n_mutations.items():157        original_residue, position = key158        original_residues.append(original_residue)159        mutations.append(value)160        positions.append(position + 1)161 162    df = pd.DataFrame({163        'Original Residue': original_residues,164        'Predicted Residues': mutations,165        'Position': positions166    })167    df.to_csv("predicted_tokens.csv", index=False)168    img.save("heatmap.png", dpi=(300, 300))169    zip_path = "outputs.zip"170    with zipfile.ZipFile(zip_path, 'w') as zipf:171        zipf.write("predicted_tokens.csv")172        zipf.write("heatmap.png")173 174 175    return df, img, zip_path176 177demo = gr.Interface(178    fn=process_sequence,179    inputs=[180        gr.Textbox(label="Sequence", placeholder="Enter the protein sequence here"),181        gr.Dataframe(182            headers=["start", "end"],183            datatype=["number", "number"],184            row_count=(1, "fixed"),185            col_count=(2, "fixed"),186            label="Domain Bounds"187        ),188        gr.Dropdown([i for i in range(1, 21)], label="Top N Tokens"),189    ],190     outputs=[191        gr.Dataframe(label="Predicted Tokens (in order of decreasing likelihood)"),192        gr.Image(type="pil", label="Heatmap"),193        gr.File(label="Download Outputs"),194    ],195)196if __name__ == "__main__":197    with suppress_output():198      demo.launch()