Update app.py

app.py

@@ -16,14 +16,14 @@ 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.bert1 = EsmForSequenceClassification.from_pretrained(model_checkpoint1, num_labels=512)  # 3000
         # for param in self.bert1.parameters():
         #     param.requires_grad = False
         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.fc1 = nn.Linear(512, 256)
         self.fc2 = nn.Linear(256, 128)
         self.fc3 = nn.Linear(128, 64)
         self.output_layer = nn.Linear(64, 2)
@@ -75,41 +75,56 @@ def classify_sequence(sequence):
     else:
         return "Invalid Sequence"
 
+
 # 加载模型
 model = MyModel()
-model.load_state_dict(torch.load("best_model.pth", map_location=torch.device('cpu')),strict=False)
-
+model.load_state_dict(torch.load("best_model.pth", map_location=torch.device('cpu')))
 
 if __name__ == "__main__":
-    with gr.Blocks() as demo:
-        
-        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.
+    title = """<h1 align="center">🔥AMP Sequence Detector</h1>"""
+    css = ".json {height: 527px; overflow: scroll;} .json-holder {height: 527px; overflow: scroll;}"
+    theme = gr.themes.Soft(primary_hue="zinc", secondary_hue="blue", neutral_hue="lime",
+                           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("<h1>Diff-AMP</h1>")
+        gr.HTML(title)
 
-        # 添加示例输入和输出
+
+        gr.Markdown(
+            "<p align='center' style='font-size: 20px;'>🔥Welcome to Antimicrobial Peptide Recognition Model. See our <a href='https://github.com/wrab12/diff-amp'>Project</a></p>")
+        gr.HTML(
+            '''<center><a href="https://huggingface.co/spaces/jackrui/diff-amp-AMP_Sequence_Detector?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a></center>''')
+        gr.HTML(
+            '''<center>🌟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.</center>''')
+
+    # 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 = [
-            ["KLLKKLLKLWKKLLKKLK"],
-            ["FLGLLFHGVHHVGKWIHGLIHGHH"],
-            ["GLMSTLKGAATNAAVTLLNKLQCKLTGTC"]
-        ]
+        ["QGLFFLGAKLFYLLTLFL"],
+        ["FLGLLFHGVHHVGKWIHGLIHGHH"],
+        ["GLMSTLKGAATNAAVTLLNKLQCKLTGTC"]
+    ]
 
-        # 创建 Gradio 接口并应用美化样式和示例
+    # 创建 Gradio 接口并应用美化样式和示例
         iface = gr.Interface(
-            fn=classify_sequence,
-            inputs="text",
-            outputs="text",
-            title="AMP Sequence Detector",
-            examples=examples
-        )
-
+        fn=classify_sequence,
+        inputs="text",
+        outputs="text",
+        # title="AMP Sequence Detector",
+        examples=examples
+    )
+        gr.Markdown(
+            "<p align='center'><img src='https://pic4.zhimg.com/v2-eb2a7c0e746e67d1768090eec74f6787_b.jpg'></p>")
 
     demo.launch()
-
-
-