from __future__ import annotations from pathlib import Path import time import gradio as gr from gradio_molecule3d import Molecule3D from gradio_molecule2d import molecule2d import numpy as np from rdkit import Chem from rdkit.Chem import AllChem import pandas as pd from biotite.structure import centroid, from_template from biotite.structure.io import load_structure from biotite.structure.io.mol import MOLFile, SDFile from biotite.structure.io.pdb import PDBFile from plinder.eval.docking.write_scores import evaluate EVAL_METRICS = ["system", "LDDT-PLI", "LDDT-LP", "BISY-RMSD"] EVAL_METRICS_PINDER = ["system","L_rms", "I_rms", "F_nat", "DOCKQ", "CAPRI_class"] import os from huggingface_hub import HfApi # Info to change for your repository # ---------------------------------- TOKEN = os.environ.get("HF_TOKEN") # A read/write token for your org OWNER = "MLSB" # Change to your org - don't forget to create a results and request dataset, with the correct format! # ---------------------------------- REPO_ID = f"{OWNER}/leaderboard2024" QUEUE_REPO = f"{OWNER}/requests" RESULTS_REPO = f"{OWNER}/results" # If you setup a cache later, just change HF_HOME CACHE_PATH=os.getenv("HF_HOME", ".") # Local caches EVAL_REQUESTS_PATH = os.path.join(CACHE_PATH, "eval-queue") EVAL_RESULTS_PATH = os.path.join(CACHE_PATH, "eval-results") EVAL_REQUESTS_PATH_BACKEND = os.path.join(CACHE_PATH, "eval-queue-bk") EVAL_RESULTS_PATH_BACKEND = os.path.join(CACHE_PATH, "eval-results-bk") API = HfApi(token=TOKEN) def get_metrics( system_id: str, receptor_file: Path, ligand_file: Path, flexible: bool = True, posebusters: bool = True, methodname: str = "", store:bool =True ) -> tuple[pd.DataFrame, float]: start_time = time.time() metrics = pd.DataFrame( [ evaluate( model_system_id=system_id, reference_system_id=system_id, receptor_file=receptor_file, ligand_file_list=[Path(ligand_file)], flexible=flexible, posebusters=posebusters, posebusters_full=False, ).get("LIG_0", {}) ] ) if posebusters: metrics["posebusters"] = metrics[ [col for col in metrics.columns if col.startswith("posebusters_")] ].sum(axis=1) metrics["posebusters_valid"] = metrics[ [col for col in metrics.columns if col.startswith("posebusters_")] ].sum(axis=1) == 20 columns = ["reference", "lddt_pli_ave", "lddt_lp_ave", "bisy_rmsd_ave"] if flexible: columns.extend(["lddt", "bb_lddt"]) if posebusters: columns.extend([col for col in metrics.columns if col.startswith("posebusters")]) metrics = metrics[columns].copy() mapping = { "lddt_pli_ave": "LDDT-PLI", "lddt_lp_ave": "LDDT-LP", "bisy_rmsd_ave": "BISY-RMSD", "reference": "system", } if flexible: mapping["lddt"] = "LDDT" mapping["bb_lddt"] = "Backbone LDDT" if posebusters: mapping["posebusters"] = "PoseBusters #checks" mapping["posebusters_valid"] = "PoseBusters valid" metrics.rename( columns=mapping, inplace=True, ) if store: with tempfile.NamedTemporaryFile as temp: metrics.to_csv(temp.name) API.upload_file( path_or_fileobj=temp.name, path_in_repo=f"{dataset}/{methodname}/{system_id}/", repo_id=QUEUE_REPO, repo_type="dataset", commit_message=f"Add {model_name} to eval queue", ) API.upload_file( path_or_fileobj=receptor_file.name, path_in_repo=f"{dataset}/{methodname}/{system_id}/", repo_id=QUEUE_REPO, repo_type="dataset", commit_message=f"Add {model_name} to eval queue", ) API.upload_file( path_or_fileobj=ligand_file.name, path_in_repo=f"{dataset}/{methodname}/{system_id}/", repo_id=QUEUE_REPO, repo_type="dataset", commit_message=f"Add {model_name} to eval queue", ) end_time = time.time() run_time = end_time - start_time return gr.DataFrame(metrics, visible=True), run_time def get_metrics_pinder( system_id: str, complex_file: Path, methodname: str = "", store:bool =True ) -> tuple[pd.DataFrame, float]: start_time = time.time() if not isinstance(prediction, Path): prediction = Path(prediction) # Infer the ground-truth name from prediction filename or directory where its stored # We need to figure out how we plan to consistently map predictions to systems so that eval metrics can be calculated # I assume we won't distribute the ground-truth structures (though they are already accessible if we don't blind system IDs) native = Path(f"./ground_truth/{system_id}.pdb") # alternatively # native = Path(f"./ground_truth/{prediction.parent.parent.stem}.pdb") # OR we need the user to provide prediction + system name try: # Get eval metrics for the prediction bdq = BiotiteDockQ(native, complex_file.name, parallel_io=False) metrics = bdq.calculate() metrics = metrics[["system", "LRMS", "iRMS", "Fnat", "DockQ", "CAPRI"]].copy() metrics.rename(columns={"LRMS": "L_rms", "iRMS": "I_rms", "Fnat": "F_nat", "DockQ": "DOCKQ", "CAPRI": "CAPRI_class"}, inplace=True) except Exception as e: failed_metrics = {"L_rms": 100.0, "I_rms": 100.0, "F_nat": 0.0, "DOCKQ": 0.0, "CAPRI_class": "Incorrect"} metrics = pd.DataFrame([failed_metrics]) metrics["system"] = native.stem gr.Error(f"Failed to evaluate prediction [{prediction}]:\n{e}") if store: # Upload to hub with tempfile.NamedTemporaryFile as temp: metrics.to_csv(temp.name) API.upload_file( path_or_fileobj=temp.name, path_in_repo=f"{dataset}/{methodname}/{system_id}/", repo_id=QUEUE_REPO, repo_type="dataset", commit_message=f"Add {model_name} to eval queue", ) API.upload_file( path_or_fileobj=complex_file.name, path_in_repo=f"{dataset}/{methodname}/{system_id}/", repo_id=QUEUE_REPO, repo_type="dataset", commit_message=f"Add {model_name} to eval queue", ) end_time = time.time() run_time = end_time - start_time return gr.DataFrame(metrics, visible=True), run_time with gr.Blocks() as app: with gr.Tab("🧬 PINDER evaluation template"): with gr.Row(): with gr.Column(): input_system_id_pinder = gr.Textbox(label="PINDER system ID") input_complex_pinder = gr.File(label="Receptor file") methodname_pinder = gr.Textbox(label="Name of your method in the format mlsb/spacename") store_pinder = gr.Checkbox(label="Store on huggingface for leaderboard", value=False) gr.Examples( [ [ "4neh__1__1.B__1.H", "input_protein_test.cif", "mlsb/test", False ], ], [input_system_id_pinder, input_complex_pinder, methodname_pinder, store_pinder], ) eval_btn_pinder = gr.Button("Run Evaluation") with gr.Tab("⚖️ PLINDER evaluation template"): with gr.Row(): with gr.Column(): input_system_id = gr.Textbox(label="PLINDER system ID") input_receptor_file = gr.File(label="Receptor file (CIF)") input_ligand_file = gr.File(label="Ligand file (SDF)") flexible = gr.Checkbox(label="Flexible docking", value=True) posebusters = gr.Checkbox(label="PoseBusters", value=True) methodname = gr.Textbox(label="Name of your method in the format mlsb/spacename") store = gr.Checkbox(label="Store on huggingface for leaderboard", value=False) gr.Examples( [ [ "4neh__1__1.B__1.H", "input_protein_test.cif", "input_ligand_test.sdf", True, True, "mlsb/test", False ], ], [input_system_id, input_receptor_file, input_ligand_file, flexible, posebusters, methodname, store], ) eval_btn = gr.Button("Run Evaluation") eval_run_time = gr.Textbox(label="Evaluation runtime") metric_table = gr.DataFrame( pd.DataFrame([], columns=EVAL_METRICS), label="Evaluation metrics", visible=False ) metric_table_pinder = gr.DataFrame( pd.DataFrame([], columns=EVAL_METRICS_PINDER), label="Evaluation metrics", visible=False ) eval_btn.click( get_metrics, inputs=[input_system_id, input_receptor_file, input_ligand_file, flexible, posebusters, methodname, store], outputs=[metric_table, eval_run_time], ) eval_btn_pinder.click( get_metrics_pinder, inputs=[input_system_id_pinder, input_complex_pinder, methodname_pinder, store_pinder], outputs=[metric_table_pinder, eval_run_time], ) app.launch()