|
import gradio as gr |
|
from transformers import AutoTokenizer, EsmForMaskedLM |
|
import torch |
|
import matplotlib.pyplot as plt |
|
import numpy as np |
|
import os |
|
|
|
def generate_heatmap(protein_sequence, start_pos=1, end_pos=None): |
|
|
|
model_name = "facebook/esm2_t6_8M_UR50D" |
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
model = EsmForMaskedLM.from_pretrained(model_name) |
|
|
|
|
|
input_ids = tokenizer.encode(protein_sequence, return_tensors="pt") |
|
sequence_length = input_ids.shape[1] - 2 |
|
|
|
|
|
if end_pos is None: |
|
end_pos = sequence_length |
|
|
|
|
|
amino_acids = list("ACDEFGHIKLMNPQRSTVWY") |
|
|
|
|
|
heatmap = np.zeros((20, end_pos - start_pos + 1)) |
|
|
|
|
|
for position in range(start_pos, end_pos + 1): |
|
|
|
masked_input_ids = input_ids.clone() |
|
masked_input_ids[0, position] = tokenizer.mask_token_id |
|
|
|
|
|
with torch.no_grad(): |
|
logits = model(masked_input_ids).logits |
|
|
|
|
|
probabilities = torch.nn.functional.softmax(logits[0, position], dim=0) |
|
log_probabilities = torch.log(probabilities) |
|
|
|
|
|
wt_residue = input_ids[0, position].item() |
|
log_prob_wt = log_probabilities[wt_residue].item() |
|
|
|
|
|
for i, amino_acid in enumerate(amino_acids): |
|
log_prob_mt = log_probabilities[tokenizer.convert_tokens_to_ids(amino_acid)].item() |
|
heatmap[i, position - start_pos] = log_prob_mt - log_prob_wt |
|
|
|
|
|
plt.figure(figsize=(15, 5)) |
|
plt.imshow(heatmap, cmap="viridis_r", aspect="auto") |
|
plt.xticks(range(end_pos - start_pos + 1), list(protein_sequence[start_pos-1:end_pos])) |
|
plt.yticks(range(20), amino_acids) |
|
plt.xlabel("Position in Protein Sequence") |
|
plt.ylabel("Amino Acid Mutations") |
|
plt.title("Predicted Effects of Mutations on Protein Sequence (LLR)") |
|
plt.colorbar(label="Log Likelihood Ratio (LLR)") |
|
plt.show() |
|
|
|
|
|
temp_file = "temp_heatmap.png" |
|
plt.savefig(temp_file) |
|
plt.close() |
|
return temp_file |
|
|
|
def heatmap_interface(sequence, start, end=None): |
|
|
|
start = int(start) |
|
if end is not None: |
|
end = int(end) |
|
|
|
|
|
if end is None or end > len(sequence) or end <= 0: |
|
end = len(sequence) |
|
|
|
|
|
if start < 1 or start > len(sequence): |
|
return "Start position is out of bounds." |
|
|
|
|
|
heatmap_path = generate_heatmap(sequence, start, end) |
|
return heatmap_path |
|
|
|
|
|
iface = gr.Interface( |
|
fn=heatmap_interface, |
|
inputs=[ |
|
gr.Textbox(lines=2, placeholder="Enter Protein Sequence Here..."), |
|
gr.Number(label="Start Position", value=1), |
|
gr.Number(label="End Position") |
|
], |
|
outputs="image", |
|
live=True |
|
) |
|
|
|
|
|
iface.launch() |