Update app.py

app.py

@@ -1,72 +1,84 @@
 import gradio as gr
+import os
 import subprocess
-import pandas as pd
 import uuid
-import os
+import pandas as pd
 
-def run_prediction(data_path, model_path, prediction_type="classification", output_path="results.csv"):
-    command = [
-        "python", model_path,
-        "--data_path", data_path,
-        "--checkpoint_path", model_path.replace("predict.py", "../models/fusion_model.pt" if "mod" in model_path else "../models/vanilla_model.pt"),
-        "--prediction_type", prediction_type,
-        "--save_dir", ".",
-        "--preds_path", output_path
-    ]
-    subprocess.run(command, check=True)
-    return pd.read_csv(output_path)
+# Helper to save SMILES string to a temporary CSV
+def save_smiles_to_csv(smiles: str, temp_dir="temp_inputs"):
+    os.makedirs(temp_dir, exist_ok=True)
+    file_path = os.path.join(temp_dir, f"{uuid.uuid4().hex}.csv")
+    df = pd.DataFrame({"smiles": [smiles]})
+    df.to_csv(file_path, index=False)
+    return file_path
+
+# Core prediction logic
+def predict(model_version, dataset_name, input_type, file=None, smiles=None):
+    # Set paths
+    if model_version == "Vanilla Chemprop":
+        model_dir = "chemprop"
+        model_path = f"model_weight/{dataset_name}/best_unbalanced.pt"
+        script_path = "chemprop/chemprop/cli/predict.py"
+    else:
+        model_dir = "chemprop_update"
+        model_path = f"model_weight/{dataset_name}/best_bert_fusion.pt"
+        script_path = "chemprop_update/chemprop/cli/predict.py"
 
-def predict_from_smiles(smiles, model_choice, dataset_choice):
-    temp_id = uuid.uuid4().hex
-    temp_csv = f"data/temp_{temp_id}.csv"
-    out_csv = f"data/output_{temp_id}.csv"
-    df = pd.DataFrame([{"smiles": smiles, "compound_name": "molecule"}])
-    df.to_csv(temp_csv, index=False)
+    # Prepare input file
+    if input_type == "Upload CSV":
+        if file is None:
+            return "Please upload a CSV file."
+        input_path = file.name
+    else:
+        if not smiles:
+            return "Please enter a SMILES string."
+        input_path = save_smiles_to_csv(smiles)
 
-    model_path = "chemprop_mod/predict.py" if model_choice == "Fusion (GNN + Transformer)" else "chemprop/predict.py"
-    
+    # Run prediction command
+    cmd = [
+        "python", script_path,
+        "--test-path", input_path,
+        "--model-paths", model_path,
+        "--smiles-columns", "smiles"
+    ]
     try:
-        predictions = run_prediction(temp_csv, model_path, output_path=out_csv)
-        return predictions
-    except subprocess.CalledProcessError as e:
-        return f"Error: {str(e)}"
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        if result.returncode != 0:
+            return f"Error:\n{result.stderr}"
+        return f"Prediction Output:\n{result.stdout}"
+    except Exception as e:
+        return f"Execution Failed: {str(e)}"
 
-def predict_from_file(file_obj, model_choice, dataset_choice):
-    temp_id = uuid.uuid4().hex
-    file_path = f"data/uploaded_{temp_id}.csv"
-    out_csv = f"data/output_file_{temp_id}.csv"
-    
-    with open(file_path, "wb") as f:
-        f.write(file_obj.read())
+# Gradio UI setup
+with gr.Blocks() as demo:
+    gr.Markdown("## 🧪 Molecular Property Prediction using Chemprop and Transformers")
 
-    model_path = "chemprop_mod/predict.py" if model_choice == "Fusion (GNN + Transformer)" else "chemprop/predict.py"
+    model_version = gr.Radio(
+        ["Vanilla Chemprop", "Updated Fusion Model"],
+        label="Select Model Version"
+    )
+    dataset_name = gr.Radio(["BBBP", "ClinTox"], label="Select Dataset")
 
-    try:
-        predictions = run_prediction(file_path, model_path, output_path=out_csv)
-        return predictions
-    except subprocess.CalledProcessError as e:
-        return f"Error: {str(e)}"
+    input_type = gr.Radio(["Upload CSV", "Single SMILES"], label="Input Type")
 
-with gr.Blocks() as demo:
-    gr.Markdown("## 🧪 Drug Property Prediction with Fusion Models")
-    gr.Markdown("Choose prediction input type and compare Chemprop vs Fusion model")
+    file_input = gr.File(file_types=[".csv"], label="Upload CSV", visible=True)
+    smiles_input = gr.Textbox(label="Enter SMILES string", visible=False)
+
+    def toggle_inputs(choice):
+        return {
+            file_input: gr.update(visible=(choice == "Upload CSV")),
+            smiles_input: gr.update(visible=(choice == "Single SMILES"))
+        }
+
+    input_type.change(toggle_inputs, input_type, [file_input, smiles_input])
 
-    with gr.Tab("Single SMILES"):
-        with gr.Row():
-            smiles_input = gr.Textbox(label="Enter SMILES string")
-            model_select = gr.Radio(["Vanilla Chemprop", "Fusion (GNN + Transformer)"], label="Model")
-            dataset_select = gr.Dropdown(["BBBP", "BACE"], label="Dataset")
-        predict_button = gr.Button("Predict")
-        result_output = gr.Dataframe(label="Prediction Result")
-        predict_button.click(fn=predict_from_smiles, inputs=[smiles_input, model_select, dataset_select], outputs=result_output)
+    predict_button = gr.Button("Predict")
+    output = gr.Textbox(label="Output")
 
-    with gr.Tab("Upload File"):
-        with gr.Row():
-            file_input = gr.File(label="Upload CSV File")
-            model_select_file = gr.Radio(["Vanilla Chemprop", "Fusion (GNN + Transformer)"], label="Model")
-            dataset_select_file = gr.Dropdown(["BBBP", "BACE"], label="Dataset")
-        predict_button_file = gr.Button("Predict")
-        result_output_file = gr.Dataframe(label="Prediction Result")
-        predict_button_file.click(fn=predict_from_file, inputs=[file_input, model_select_file, dataset_select_file], outputs=result_output_file)
+    predict_button.click(
+        fn=predict,
+        inputs=[model_version, dataset_name, input_type, file_input, smiles_input],
+        outputs=output
+    )
 
 demo.launch()