Spaces:
Sleeping
Sleeping
import pandas as pd | |
from IPython.display import clear_output | |
import torch | |
from transformers import EsmForSequenceClassification, AdamW, AutoTokenizer | |
from torch.utils.data import DataLoader, TensorDataset, random_split | |
from sklearn.preprocessing import LabelEncoder | |
from tqdm import tqdm | |
import numpy as np | |
import seaborn as sns | |
from sklearn.model_selection import train_test_split | |
import matplotlib.pyplot as plt | |
import pickle | |
import torch.nn.functional as F | |
import gradio as gr | |
import io | |
from PIL import Image | |
import Bio | |
from Bio import SeqIO | |
import zipfile | |
import os | |
# Load the model from the file | |
with open('family_labels.pkl', 'rb') as filefam: | |
yfam = pickle.load(filefam) | |
tokenizerfam = AutoTokenizer.from_pretrained("facebook/esm2_t12_35M_UR50D") #facebook/esm2_t33_650M_UR50D | |
device = 'cpu' | |
device | |
modelfam = EsmForSequenceClassification.from_pretrained("facebook/esm2_t12_35M_UR50D", num_labels=len(yfam.classes_)) | |
modelfam = modelfam.to('cpu') | |
modelfam.load_state_dict(torch.load("family.pth", map_location=torch.device('cpu'))) | |
modelfam.eval() | |
x_testfam = ["""MAEVLRTLAGKPKCHALRPMILFLIMLVLVLFGYGVLSPRSLMPGSLERGFCMAVREPDH | |
LQRVSLPRMVYPQPKVLTPCRKDVLVVTPWLAPIVWEGTFNIDILNEQFRLQNTTIGLTV | |
FAIKKYVAFLKLFLETAEKHFMVGHRVHYYVFTDQPAAVPRVTLGTGRQLSVLEVRAYKR | |
WQDVSMRRMEMISDFCERRFLSEVDYLVCVDVDMEFRDHVGVEILTPLFGTLHPGFYGSS | |
REAFTYERRPQSQAYIPKDEGDFYYLGGFFGGSVQEVQRLTRACHQAMMVDQANGIEAVW | |
HDESHLNKYLLRHKPTKVLSPEYLWDQQLLGWPAVLRKLRFTAVPKNHQAVRNP | |
"""] | |
encoded_inputfam = tokenizerfam(x_testfam, padding=True, truncation=True, max_length=512, return_tensors="pt") | |
input_idsfam = encoded_inputfam["input_ids"] | |
attention_maskfam = encoded_inputfam["attention_mask"] | |
with torch.no_grad(): | |
outputfam = modelfam(input_idsfam, attention_mask=attention_maskfam) | |
logitsfam = outputfam.logits | |
probabilitiesfam = F.softmax(logitsfam, dim=1) | |
_, predicted_labelsfam = torch.max(logitsfam, dim=1) | |
probabilitiesfam[0] | |
decoded_labelsfam = yfam.inverse_transform(predicted_labelsfam.tolist()) | |
decoded_labelsfam | |
#Load donor model from file | |
tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t12_35M_UR50D") | |
with open('donor_labels.pkl', 'rb') as file: | |
label_encoder = pickle.load(file) | |
# encoded_labels = label_encoder.fit(y) | |
# labels = torch.tensor(encoded_labels) | |
model = EsmForSequenceClassification.from_pretrained("facebook/esm2_t12_35M_UR50D", num_labels=len(label_encoder.classes_)) | |
model = model.to('cpu') | |
model.load_state_dict(torch.load("best_model_35M_t12_5v5.pth", map_location=torch.device('cpu'))) #model_best_35v2M.pth | |
model.eval() | |
x_test = ["""MAEVLRTLAGKPKCHALRPMILFLIMLVLVLFGYGVLSPRSLMPGSLERGFCMAVREPDH | |
LQRVSLPRMVYPQPKVLTPCRKDVLVVTPWLAPIVWEGTFNIDILNEQFRLQNTTIGLTV | |
FAIKKYVAFLKLFLETAEKHFMVGHRVHYYVFTDQPAAVPRVTLGTGRQLSVLEVRAYKR | |
WQDVSMRRMEMISDFCERRFLSEVDYLVCVDVDMEFRDHVGVEILTPLFGTLHPGFYGSS | |
REAFTYERRPQSQAYIPKDEGDFYYLGGFFGGSVQEVQRLTRACHQAMMVDQANGIEAVW | |
HDESHLNKYLLRHKPTKVLSPEYLWDQQLLGWPAVLRKLRFTAVPKNHQAVRNP | |
"""] | |
encoded_input = tokenizer(x_test, padding=True, truncation=True, max_length=512, return_tensors="pt") | |
input_ids = encoded_input["input_ids"] | |
attention_mask = encoded_input["attention_mask"] | |
with torch.no_grad(): | |
output = model(input_ids, attention_mask=attention_mask) | |
logits = output.logits | |
probabilities = F.softmax(logits, dim=1) | |
_, predicted_labels = torch.max(logits, dim=1) | |
probabilities[0] | |
decoded_labels = label_encoder.inverse_transform(predicted_labels.tolist()) | |
decoded_labels | |
glycosyltransferase_db = { | |
"GT31-chsy" : {'CAZy Name': 'GT31', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '8 ', 'More Info': 'http://www.cazy.org/GT31.html'}, | |
"GT2-CesA2" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '1 ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT43-arath" : {'CAZy Name': 'GT43', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT43.html'}, | |
"GT8-Met1" : {'CAZy Name': 'GT8 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': '9 ', 'More Info': 'http://www.cazy.org/GT8.html' }, | |
"GT32-higher" : {'CAZy Name': 'GT32', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT32.html'}, | |
"GT40" : {'CAZy Name': 'GT40', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT40.html'}, | |
"GT16" : {'CAZy Name': 'GT16', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '6 ', 'More Info': 'http://www.cazy.org/GT16.html'}, | |
"GT27" : {'CAZy Name': 'GT27', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': '5 ', 'More Info': 'http://www.cazy.org/GT27.html'}, | |
"GT55" : {'CAZy Name': 'GT55', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': '2 ', 'More Info': 'http://www.cazy.org/GT55.html'}, | |
"GT8-Glycogenin" : {'CAZy Name': 'GT8 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': '9 ', 'More Info': 'http://www.cazy.org/GT8.html' }, | |
"GT8-1" : {'CAZy Name': 'GT8 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': '9 ', 'More Info': 'http://www.cazy.org/GT8.html' }, | |
"GT25" : {'CAZy Name': 'GT25', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '6 ', 'More Info': 'http://www.cazy.org/GT25.html'}, | |
"GT2-DPM_like" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '2 ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT31-fringe" : {'CAZy Name': 'GT31', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '8 ', 'More Info': 'http://www.cazy.org/GT31.html'}, | |
"GT2-Bact_puta" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT84" : {'CAZy Name': 'GT84', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': '1 ', 'More Info': 'http://www.cazy.org/GT84.html'}, | |
"GT13" : {'CAZy Name': 'GT13', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '6 ', 'More Info': 'http://www.cazy.org/GT13.html'}, | |
"GT43-cele" : {'CAZy Name': 'GT43', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT43.html'}, | |
"GT2-Bact_LPS1" : {'CAZy Name': 'GT92', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT2-Bact_Oant" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': ' ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT67" : {'CAZy Name': 'GT67', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '8 ', 'More Info': 'http://www.cazy.org/GT67.html'}, | |
"GT2-HAS" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '1 ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT82" : {'CAZy Name': 'GT82', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '7 ', 'More Info': 'http://www.cazy.org/GT82.html'}, | |
"GT24" : {'CAZy Name': 'GT24', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': '9 ', 'More Info': 'http://www.cazy.org/GT24.html'}, | |
"GT31-plant" : {'CAZy Name': 'GT31', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '8 ', 'More Info': 'http://www.cazy.org/GT31.html'}, | |
"GT81-Bact" : {'CAZy Name': 'GT81', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': '2 ', 'More Info': 'http://www.cazy.org/GT81.html'}, | |
"GT2-Bact_gt25Me": {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': ' ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT2-B3GntL" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '4 ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT49" : {'CAZy Name': 'GT49', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT49.html'}, | |
"GT34" : {'CAZy Name': 'GT34', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT34.html'}, | |
"GT45" : {'CAZy Name': 'GT45', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT45.html'}, | |
"GT32-lower" : {'CAZy Name': 'GT32', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT32.html'}, | |
"GT88" : {'CAZy Name': 'GT88', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': '9 ', 'More Info': 'http://www.cazy.org/GT88.html'}, | |
"GT21" : {'CAZy Name': 'GT21', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '1 ', 'More Info': 'http://www.cazy.org/GT21.html'}, | |
"GT2-DPG_synt" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '2 ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT43-b3gat2" : {'CAZy Name': 'GT43', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT43.html'}, | |
"GT2-Chitin_synt": {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '5 ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT8-Bact" : {'CAZy Name': 'GT8 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT8.html' }, | |
"GT8-Met2" : {'CAZy Name': 'GT8 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT8.html' }, | |
"GT2-Bact_Chlor1": {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': ' ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT54" : {'CAZy Name': 'GT54', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '6 ', 'More Info': 'http://www.cazy.org/GT54.html'}, | |
"GT2-Cel_bre3" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '1 ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT2-Bact_Rham" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT6" : {'CAZy Name': 'GT6 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT6.html' }, | |
"GT2-Bact_puta2" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': ' ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT7-1" : {'CAZy Name': 'GT7 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '5 ', 'More Info': 'http://www.cazy.org/GT7.html' }, | |
"GT2-Csl" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '4 ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT2-ExoU" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': ' ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT2-Csl2" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '4 ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT64" : {'CAZy Name': 'GT64', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT64.html'}, | |
"GT2-Bact_Chlor2": {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': ' ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT78" : {'CAZy Name': 'GT78', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': '2 ', 'More Info': 'http://www.cazy.org/GT78.html'}, | |
"GT12" : {'CAZy Name': 'GT12', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT12.html'}, | |
"GT31-gnt" : {'CAZy Name': 'GT31', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '8 ', 'More Info': 'http://www.cazy.org/GT31.html'}, | |
"GT2-Bact_CHS" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '5 ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT62" : {'CAZy Name': 'GT62', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': '3 ', 'More Info': 'http://www.cazy.org/GT62.html'}, | |
"GT8-Met_Pla" : {'CAZy Name': 'GT8 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT8.html' }, | |
"GT15" : {'CAZy Name': 'GT15', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': '8 ', 'More Info': 'http://www.cazy.org/GT15.html'}, | |
"GT43-b3gat1" : {'CAZy Name': 'GT43', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT43.html'}, | |
"GT31-b3glt" : {'CAZy Name': 'GT31', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '8 ', 'More Info': 'http://www.cazy.org/GT31.html'}, | |
"GT2-CesA1" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '1 ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT60" : {'CAZy Name': 'GT60', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': '5 ', 'More Info': 'http://www.cazy.org/GT60.html'}, | |
"GT14" : {'CAZy Name': 'GT14', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '7 ', 'More Info': 'http://www.cazy.org/GT14.html'}, | |
"GT2-Bact_DPM_sy": {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '2 ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT17" : {'CAZy Name': 'GT17', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '7 ', 'More Info': 'http://www.cazy.org/GT17.html'}, | |
"GT2-Bact_LPS2" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': '3 ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT77" : {'CAZy Name': 'GT77', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': '9 ', 'More Info': 'http://www.cazy.org/GT77.html'}, | |
"GT2-Bact_EpsO" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': ' ', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
"GT43-b3gat3" : {'CAZy Name': 'GT43', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT43.html'}, | |
"GT8-Fun" : {'CAZy Name': 'GT8 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Retaining', 'Clade': '9 ', 'More Info': 'http://www.cazy.org/GT8.html' }, | |
"GT75" : {'CAZy Name': 'GT75', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT75.html'}, | |
"GT2-Bact_GlfT" : {'CAZy Name': 'GT2 ', 'Alternative Name': '', 'Fold': 'A', 'Mechanism': 'Inverting', 'Clade': 'N/A', 'More Info': 'http://www.cazy.org/GT2.html' }, | |
} | |
def get_family_info(family_name): | |
family_info = glycosyltransferase_db.get(family_name, {}) | |
output = "" | |
for key, value in family_info.items(): | |
if key == "more_info": | |
output += "**{}:**".format(key.title().replace("_", " ")) + "\n" | |
for link in value: | |
output += "[{}]({}) ".format(link, link) | |
else: | |
output += "**{}:** {} ".format(key.title().replace("_", " "), value) | |
return output | |
def fig_to_img(fig): | |
"""Converts a matplotlib figure to a PIL Image and returns it""" | |
buf = io.BytesIO() | |
fig.savefig(buf, format='png', bbox_inches='tight') | |
buf.seek(0) | |
img = Image.open(buf) | |
return img | |
def preprocess_protein_sequence(protein_fasta): | |
lines = protein_fasta.split('\n') | |
headers = [line for line in lines if line.startswith('>')] | |
if len(headers) > 1: | |
return None, "Multiple fasta sequences detected. Please upload a fasta file with only one sequence." | |
protein_sequence = ''.join(line for line in lines if not line.startswith('>')) | |
# Check for invalid characters | |
valid_characters = set("ACDEFGHIKLMNPQRSTVWYacdefghiklmnpqrstvwy") # the 20 standard amino acids | |
if not set(protein_sequence).issubset(valid_characters): | |
return None, "Invalid protein sequence. It contains characters that are not one of the 20 standard amino acids. Does your sequence contain gaps?" | |
return protein_sequence, None | |
def process_family_sequence(protein_fasta): | |
protein_sequence, error_msg = preprocess_protein_sequence(protein_fasta) | |
if error_msg: | |
return None, None, None, error_msg | |
encoded_input = tokenizer([protein_sequence], padding=True, truncation=True, max_length=512, return_tensors="pt") | |
input_idsfam = encoded_input["input_ids"] | |
attention_maskfam = encoded_input["attention_mask"] | |
with torch.no_grad(): | |
outputfam = modelfam(input_idsfam, attention_mask=attention_maskfam) | |
logitsfam = outputfam.logits | |
probabilitiesfam = F.softmax(logitsfam, dim=1) | |
_, predicted_labelsfam = torch.max(logitsfam, dim=1) | |
decoded_labelsfam = yfam.inverse_transform(predicted_labelsfam.tolist()) | |
family_info = get_family_info(decoded_labelsfam[0]) | |
figfam = plt.figure(figsize=(10, 5)) | |
labelsfam = yfam.classes_ | |
probabilitiesfam = probabilitiesfam.tolist() | |
# Convert the nested list to a flat list of probabilities | |
probabilitiesfam_flat = probabilitiesfam[0] if probabilitiesfam else [] | |
# Sort labels and probabilities by probability | |
labels_probsfam = list(zip(labelsfam, probabilitiesfam_flat)) | |
labels_probsfam.sort(key=lambda x: x[1], reverse=True) | |
# Select the top 5 fams | |
labels_probs_top5fam = labels_probsfam[:5] | |
labels_top5, probabilities_top5 = zip(*labels_probs_top5fam) | |
y_posfam = np.arange(len(labels_top5)) | |
plt.barh(y_posfam, [prob*100 for prob in probabilities_top5], align='center', alpha=0.5) | |
plt.yticks(y_posfam, labels_top5) | |
plt.xlabel('Probability (%)') | |
plt.title('Top 5 Family Class Probabilities') | |
plt.xlim(0, 100) | |
plt.close(figfam) | |
img = fig_to_img(figfam) | |
if len(protein_sequence) < 100: | |
return decoded_labelsfam[0], img, None, f"**Warning:** The sequence is relatively short. Fragmentary and partial sequences may result in incorrect predictions. \n\n {family_info}" | |
return decoded_labelsfam[0], img, None, family_info | |
def process_single_sequence(protein_fasta): #, protein_file | |
protein_sequence, error_msg = preprocess_protein_sequence(protein_fasta) | |
if error_msg: | |
return None, None, None, error_msg | |
encoded_input = tokenizer([protein_sequence], padding=True, truncation=True, max_length=512, return_tensors="pt") | |
input_ids = encoded_input["input_ids"] | |
attention_mask = encoded_input["attention_mask"] | |
with torch.no_grad(): | |
output = model(input_ids, attention_mask=attention_mask) | |
logits = output.logits | |
dprobabilities = F.softmax(logits, dim=1)[0] | |
_, predicted_labels = torch.max(logits, dim=1) | |
decoded_labels = label_encoder.inverse_transform(predicted_labels.tolist()) | |
family_info = get_family_info(decoded_labels[0]) | |
fig = plt.figure(figsize=(10, 5)) | |
labels = label_encoder.classes_ | |
dprobabilities = dprobabilities.tolist() | |
# Sort labels and probabilities by probability | |
labels_probs = list(zip(labels, dprobabilities)) | |
labels_probs.sort(key=lambda x: x[1], reverse=True) | |
# Select the top 3 donors | |
labels_probs_top3 = labels_probs[:3] | |
labels_top3, probabilities_top3 = zip(*labels_probs_top3) | |
y_pos = np.arange(len(labels_top3)) | |
plt.barh(y_pos, [prob*100 for prob in probabilities_top3], align='center', alpha=0.5) | |
plt.yticks(y_pos, labels_top3) | |
plt.xlabel('Probability (%)') | |
plt.title('Top 3 Donor Class Probabilities') | |
plt.xlim(0, 100) | |
plt.close(fig) | |
img = fig_to_img(fig) | |
if len(protein_sequence) < 100: | |
return decoded_labels[0], img, None, f"**Warning:** The sequence is relatively short. Fragmentary and partial sequences may result in incorrect predictions. \n\n {family_info}" | |
return decoded_labels[0], img, None, None | |
def process_sequence_file(protein_file): # added progress parameter that is displayed in gradio #, progress=gr.Progress() | |
try: | |
records = list(SeqIO.parse(protein_file.name, "fasta")) | |
except Exception as e: | |
return str(e) | |
if not os.path.exists('results'): | |
os.makedirs('results') | |
total = len(records) | |
for idx, record in enumerate(records): | |
protein_sequence = str(record.seq) | |
valid_characters = set("ACDEFGHIKLMNPQRSTVWYacdefghiklmnpqrstvwy") | |
if not set(protein_sequence).issubset(valid_characters): | |
with open(f'results/result_{idx+1}.txt', 'w') as file: | |
file.write("Invalid protein sequence. It contains characters that are not one of the 20 standard amino acids. Does your sequence contain gaps?") | |
continue | |
label, img, _, info = process_single_sequence(protein_sequence) | |
img.save(f'results/result_{idx+1}.png') | |
with open(f'results/result_{idx+1}.txt', 'w') as file: | |
file.write(f'Predicted Donor: {label}\n\n{info}') | |
# progress(idx/total) # Update the progress bar | |
# Create a zip file w/ results -- To Do: Figure out how to improve compression for large files | |
with zipfile.ZipFile('predicted_results.zip', 'w', zipfile.ZIP_DEFLATED) as zipf: | |
for root, dirs, files in os.walk('results/'): | |
for file in files: | |
zipf.write(os.path.join(root, file)) | |
return 'predicted_results.zip' #Provide indication of how to interpret downloaded zip file? f"**Warning:** The sequence is relatively short. Fragmentary and partial sequences may result in incorrect predictions. | |
# Function to mask a residue at a particular position | |
def mask_residue(sequence, position): | |
return sequence[:position] + 'X' + sequence[position+1:] | |
def generate_heatmap(protein_fasta): | |
protein_sequence, error_msg = preprocess_protein_sequence(protein_fasta) | |
# Tokenize and predict for original sequence | |
encoded_input = tokenizer([protein_sequence], padding=True, truncation=True, max_length=512, return_tensors="pt") | |
with torch.no_grad(): | |
original_output = model(encoded_input["input_ids"], attention_mask=encoded_input["attention_mask"]) | |
original_probabilities = F.softmax(original_output.logits, dim=1).cpu().numpy()[0] | |
# Define the size of each group | |
group_size = 10 # allow user to change this | |
# Calculate the number of groups | |
num_groups = len(protein_sequence) // group_size + (len(protein_sequence) % group_size > 0) | |
# Initialize an array to hold the importance scores | |
importance_scores = np.zeros((num_groups, len(original_probabilities))) | |
# Initialize tqdm progress bar | |
# with tqdm(total=num_groups, desc="Processing groups", position=0, leave=True) as pbar: | |
# # Loop through each group of residues in the sequence | |
for i in range(0, len(protein_sequence), group_size): | |
# Mask the residues in the group at positions [i, i + group_size) | |
masked_sequence = protein_sequence[:i] + 'X' * min(group_size, len(protein_sequence) - i) + protein_sequence[i + group_size:] | |
# Tokenize and predict for the masked sequence | |
encoded_input = tokenizer([masked_sequence], padding=True, truncation=True, max_length=512, return_tensors="pt") | |
with torch.no_grad(): | |
masked_output = model(encoded_input["input_ids"], attention_mask=encoded_input["attention_mask"]) | |
masked_probabilities = F.softmax(masked_output.logits, dim=1).cpu().numpy()[0] | |
# Calculate the change in probabilities and store it as the importance score | |
group_index = i // group_size | |
importance_scores[group_index, :] = np.abs(original_probabilities - masked_probabilities) | |
progress = (i // group_size + 1) / num_groups * 100 | |
print(f"Progress: {progress:.2f}%") | |
figmap, ax = plt.subplots(figsize=(20, 20)) | |
sns.heatmap(importance_scores, annot=True, cmap="coolwarm", xticklabels=label_encoder.classes_, yticklabels=[f"{i}-{i+group_size-1}" for i in range(0, len(protein_sequence), group_size)], ax=ax) | |
ax.set_xlabel("Predicted Labels") | |
ax.set_ylabel("Residue Position Groups") | |
img = fig_to_img(figmap) | |
return img | |
def main_function_single(sequence, show_explanation): | |
# Process seq, and return outputs for both fam and don | |
family_label, family_img, _, family_info = process_family_sequence(sequence) | |
donor_label, donor_img, *_ = process_single_sequence(sequence) | |
figmap = None | |
if show_explanation: | |
figmap = generate_heatmap(sequence) | |
return family_label, family_img, family_info, donor_label, donor_img, figmap | |
def main_function_upload(protein_file): #, progress=gr.Progress() | |
return process_sequence_file(protein_file) #, progress | |
prediction_imagefam = gr.outputs.Image(type='pil', label="Family prediction graph") | |
prediction_imagedonor = gr.outputs.Image(type='pil', label="Donor prediction graph") | |
prediction_explain = gr.outputs.Image(type='pil', label="Donor prediction explanation") | |
with gr.Blocks() as app: | |
gr.Markdown("# Glydentify") | |
with gr.Tab("Single Sequence Prediction"): | |
with gr.Row().style(equal_height=True): | |
with gr.Column(): | |
sequence = gr.inputs.Textbox(lines=16, placeholder='Enter Protein Sequence Here...', label="Protein Sequence") | |
explanation_checkbox = gr.inputs.Checkbox(label="Show Explanation", default=False) | |
with gr.Column(): | |
with gr.Accordion("Example:"): | |
gr.Markdown(""" | |
\>sp|Q9Y5Z6|B3GT1_HUMAN Beta-1,3-galactosyltransferase 1 OS=Homo sapiens OX=9606 GN=B3GALT1 PE=1 SV=1 | |
MASKVSCLYVLTVVCWASALWYLSITRPTSSYTGSKPFSHLTVARKNFTFGNIRTRPINPHSFEFLINEPNKCEKNIPFLVILIST | |
THKEFDARQAIRETWGDENNFKGIKIATLFLLGKNADPVLNQMVEQESQIFHDIIVEDFIDSYHNLTLKTLMGMRWVATFCSK | |
AKYVMKTDSDIFVNMDNLIYKLLKPSTKPRRRYFTGYVINGGPIRDVRSKWYMPRDLYPDSNYPPFCSGTGYIFSADVAELIYK | |
TSLHTRLLHLEDVYVGLCLRKLGIHPFQNSGFNHWKMAYSLCRYRRVITVHQISPEEMHRIWNDMSSKKHLRC | |
""") | |
family_prediction = gr.outputs.Textbox(label="Predicted family") | |
donor_prediction = gr.outputs.Textbox(label="Predicted donor") | |
info_markdown = gr.Markdown() | |
# Predict and Clear buttons | |
with gr.Row().style(equal_height=True): | |
with gr.Column(): | |
predict_button = gr.Button("Predict") | |
predict_button.click(main_function_single, inputs=[sequence, explanation_checkbox], | |
outputs=[family_prediction, prediction_imagefam, info_markdown, | |
donor_prediction, prediction_imagedonor, prediction_explain]) | |
# Family & Donor Section | |
with gr.Row().style(equal_height=True): | |
with gr.Column(): | |
with gr.Accordion("Prediction Bar Graphs:"): | |
prediction_imagefam.render() # = gr.outputs.Image(type='pil', label="Family prediction graph") | |
prediction_imagedonor.render() # = gr.outputs.Image(type='pil', label="Donor prediction graph") | |
# Explain Section | |
with gr.Column(): | |
if explanation_checkbox: # Only render if the checkbox is checked | |
with gr.Accordion("Donor explanation"): | |
prediction_explain.render() # = gr.outputs.Image(type='pil', label="Donor prediction explaination") | |
with gr.Tab("Multiple Sequence Prediction"): | |
with gr.Row().style(equal_height=True): | |
with gr.Column(): | |
protein_file = gr.inputs.File(label="Upload FASTA file") | |
with gr.Column(): | |
result_file = gr.outputs.File(label="Download predictions of uploaded sequences") | |
with gr.Row().style(equal_height=True): | |
with gr.Column(): | |
process_button = gr.Button("Process") | |
process_button.click(main_function_upload, inputs=protein_file, outputs=[result_file]) | |
with gr.Column(): | |
clear = gr.Button("Clear") | |
clear.click(lambda: None) | |
# clear.click() | |
app.launch(show_error=True) | |