import numpy as np from transformers import AutoTokenizer, AutoModelForSequenceClassification, EsmForSequenceClassification from transformers import set_seed import torch import torch.nn as nn import warnings from tqdm import tqdm import gradio as gr warnings.filterwarnings('ignore') device = "cpu" model_checkpoint1 = "facebook/esm2_t12_35M_UR50D" tokenizer = AutoTokenizer.from_pretrained(model_checkpoint1) class MyModel(nn.Module): def __init__(self): super().__init__() self.bert1 = EsmForSequenceClassification.from_pretrained(model_checkpoint1, num_labels=3000) # 3000 self.bn1 = nn.BatchNorm1d(256) self.bn2 = nn.BatchNorm1d(128) self.bn3 = nn.BatchNorm1d(64) self.relu = nn.LeakyReLU() self.fc1 = nn.Linear(3000, 256) self.fc2 = nn.Linear(256, 128) self.fc3 = nn.Linear(128, 64) self.output_layer = nn.Linear(64, 2) self.dropout = nn.Dropout(0.3) # 0.3 def forward(self, x): with torch.no_grad(): bert_output = self.bert1(input_ids=x['input_ids'], attention_mask=x['attention_mask']) # output_feature = bert_output["logits"] # print(output_feature.size()) # output_feature = self.bn1(self.fc1(output_feature)) # output_feature = self.bn2(self.fc1(output_feature)) # output_feature = self.relu(self.bn3(self.fc3(output_feature))) # output_feature = self.dropout(self.output_layer(output_feature)) output_feature = self.dropout(bert_output["logits"]) output_feature = self.dropout(self.relu(self.bn1(self.fc1(output_feature)))) output_feature = self.dropout(self.relu(self.bn2(self.fc2(output_feature)))) output_feature = self.dropout(self.relu(self.bn3(self.fc3(output_feature)))) output_feature = self.dropout(self.output_layer(output_feature)) # return torch.sigmoid(output_feature),output_feature return torch.softmax(output_feature, dim=1) def AMP(test_sequences, model): # 保持 AMP 函数不变，只处理传入的 test_sequences 数据 max_len = 18 test_data = tokenizer(test_sequences, max_length=max_len, padding="max_length", truncation=True, return_tensors='pt') model = model.to(device) model.eval() out_probability = [] with torch.no_grad(): predict = model(test_data) out_probability.extend(np.max(np.array(predict.cpu()), axis=1).tolist()) test_argmax = np.argmax(predict.cpu(), axis=1).tolist() id2str = {0: "non-AMP", 1: "AMP"} return id2str[test_argmax[0]], out_probability[0] def classify_sequence(sequence): # Check if the sequence is a valid amino acid sequence and has a length of at least 3 valid_amino_acids = set("ACDEFGHIKLMNPQRSTVWY") sequence = sequence.upper() if all(aa in valid_amino_acids for aa in sequence) and len(sequence) >= 3: result, probability = AMP(sequence, model) return "yes" if result == "AMP" else "no" else: return "Invalid Sequence" # 加载模型 model = MyModel() model.load_state_dict(torch.load("best_model.pth", map_location=torch.device('cpu'))) if __name__ == "__main__": title = """

🔥AMP Sequence Detector

""" css = ".json {height: 527px; overflow: scroll;} .json-holder {height: 527px; overflow: scroll;}" theme = gr.themes.Soft(primary_hue="zinc", secondary_hue="blue", neutral_hue="green", text_size=gr.themes.sizes.text_lg) with gr.Blocks(css = """#col_container { margin-left: auto; margin-right: auto;} #chatbot {height: 520px; overflow: auto;}""", theme=theme) as demo: gr.Markdown("

Diff-AMP

") gr.HTML(title) gr.Markdown( "

🔥Welcome to Antimicrobial Peptide Recognition Model. See our Project

") gr.HTML( '''

''') gr.HTML( '''🌟Note: This is an antimicrobial peptide recognition model derived from Diff-AMP, which is a branch of a comprehensive system integrating generation, recognition, and optimization. In this recognition model, you can simply input a sequence, and it will predict whether it is an antimicrobial peptide. Due to limited website capacity, we can only perform simple predictions. If you require large-scale computations, please contact my email at wangrui66677@gmail.com. Feel free to reach out if you have any questions or inquiries.''') # gr.Markdown( # """ # # # Welcome to Antimicrobial Peptide Recognition Model # This is an antimicrobial peptide recognition model derived from Diff-AMP, which is a branch of a comprehensive system integrating generation, recognition, and optimization. In this recognition model, you can simply input a sequence, and it will predict whether it is an antimicrobial peptide. Due to limited website capacity, we can only perform simple predictions. # If you require large-scale computations, please contact my email at wangrui66677@gmail.com. Feel free to reach out if you have any questions or inquiries. # # """) # 添加示例输入和输出 examples = [ ["QGLFFLGAKLFYLLTLFL"], ["FLGLLFHGVHHVGKWIHGLIHGHH"], ["GLMSTLKGAATNAAVTLLNKLQCKLTGTC"] ] # 创建 Gradio 接口并应用美化样式和示例 iface = gr.Interface( fn=classify_sequence, inputs="text", outputs="text", # title="AMP Sequence Detector", examples=examples ) gr.Markdown( "

") gr.Markdown("