Update app.py

app.py

@@ -1,9 +1,9 @@
-import gradio as gr
 import sys
 import random
 import os
 import pandas as pd
 import torch
+import itertools
 from torch.utils.data import DataLoader
 from transformers import AutoTokenizer
 
@@ -23,6 +23,7 @@ class Config:
     padding_side = "right"
     task = "classification"
     sequence_col = "sequence"
+    seed = 42
 
 
 # Assuming 'predict_stability' is your function that predicts protein stability
@@ -91,13 +92,14 @@ def predict(cfg, sequence):
     model = PLTNUM_PreTrainedModel.from_pretrained(cfg.model_path, cfg=cfg)
     model.to(cfg.device)
 
+    # predictions = predict_fn(loader, model, cfg)
     model.eval()
     predictions = []
 
     for inputs, _ in dataloader:
         inputs = inputs.to(cfg.device)
         with torch.no_grad():
-            with torch.amp.autocast(enabled=cfg.use_amp):
+            with torch.amp.autocast(cfg.device, enabled=cfg.use_amp):
                 preds = (
                     torch.sigmoid(model(inputs))
                     if cfg.task == "classification"
@@ -105,88 +107,9 @@ def predict(cfg, sequence):
                 )
         predictions += preds.cpu().tolist()
     outputs = {}
+    predictions = list(itertools.chain.from_iterable(predictions))
     outputs["raw prediction values"] = predictions
     outputs["binary prediction values"] = [1 if x > 0.5 else 0 for x in predictions]
-    return str(outputs)
+    return outputs
 
-
-# Gradio Interface
-with gr.Blocks() as demo:
-    gr.Markdown(
-        """
-        # PLTNUM: Protein LifeTime Neural Model
-        **Predict the protein half-life from its sequence or PDB file.**
-        """
-    )
-
-    gr.Image(
-        "https://github.com/sagawatatsuya/PLTNUM/blob/main/model-image.png?raw=true",
-        label="Model Image",
-    )
-
-    # Model and Organism selection in the same row to avoid layout issues
-    with gr.Row():
-        model_choice = gr.Radio(
-            choices=["SaProt", "ESM2"],
-            label="Select PLTNUM's base model.",
-            value="SaProt",
-        )
-        organism_choice = gr.Radio(
-            choices=["Mouse", "Human"],
-            label="Select the target organism.",
-            value="Mouse",
-        )
-
-    with gr.Tabs():
-        with gr.TabItem("Upload PDB File"):
-            gr.Markdown("### Upload your PDB file:")
-            pdb_file = gr.File(label="Upload PDB File")
-
-            predict_button = gr.Button("Predict Stability")
-            prediction_output = gr.Textbox(
-                label="Stability Prediction", interactive=False
-            )
-
-            predict_button.click(
-                fn=predict_stability,
-                inputs=[model_choice, organism_choice, pdb_file],
-                outputs=prediction_output,
-            )
-
-        with gr.TabItem("Enter Protein Sequence"):
-            gr.Markdown("### Enter the protein sequence:")
-            sequence = gr.Textbox(
-                label="Protein Sequence",
-                placeholder="Enter your protein sequence here...",
-                lines=8,
-            )
-            predict_button = gr.Button("Predict Stability")
-            prediction_output = gr.Textbox(
-                label="Stability Prediction", interactive=False
-            )
-
-            predict_button.click(
-                fn=predict_stability,
-                inputs=[model_choice, organism_choice, sequence],
-                outputs=prediction_output,
-            )
-
-    gr.Markdown(
-        """
-        ### How to Use:
-        - **Select Model**: Choose between 'SaProt' or 'ESM2' for your prediction.
-        - **Select Organism**: Choose between 'Mouse' or 'Human'.
-        - **Upload PDB File**: Choose the 'Upload PDB File' tab and upload your file.
-        - **Enter Sequence**: Alternatively, switch to the 'Enter Protein Sequence' tab and input your sequence.
-        - **Predict**: Click 'Predict Stability' to receive the prediction.
-        """
-    )
-
-    gr.Markdown(
-        """
-        ### About the Tool
-        This tool allows researchers and scientists to predict the stability of proteins using advanced algorithms. It supports both PDB file uploads and direct sequence input.
-        """
-    )
-
-demo.launch()
+predict_stability("SaProt", "Human", sequence="MELKQK")