app.py

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)