{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a7a5d41-e6d7-4efa-a481-3182963ca888",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install gradio_client"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "549b9b2c-3074-446b-962e-90c8efd2bd59",
   "metadata": {},
   "source": [
    "# PLINDER inference and evaluation template API examples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b979671e-97d6-4c52-bc6e-279a09d722c8",
   "metadata": {},
   "source": [
    "## Run inference via predict endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "980771b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "uri_inference = \"https://mlsb-blackhole-models-pli.hf.space/\"\n",
    "uri_inference_busted = \"https://mlsb-blackhole-models-pli-busted.hf.space/\"\n",
    "uri_inference_strongdocking = \"https://mlsb-strong-docking-baseline.hf.space/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c0171fa-ee2a-40b7-8578-aa8516b4ece9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from gradio_client import Client, handle_file\n",
    "from pathlib import Path\n",
    "\n",
    "\n",
    "# If running docker container locally\n",
    "dev_uri = \"http://localhost:7860/\"\n",
    "client = Client(uri_inference)\n",
    "\n",
    "result = client.predict(\n",
    "    input_sequence=\"\",\n",
    "    input_ligand=\"CC(=O)N[C@H]1[C@H](O[C@H]2[C@H](O)[C@@H](NC(C)=O)CO[C@@H]2CO)O[C@H](CO)[C@@H](O)[C@@H]1O\",\n",
    "    input_msa=handle_file(\"./empty.a3m\"),\n",
    "    input_protein=handle_file(\"./input_protein_test.pdb\"),\n",
    "    api_name=\"/predict\",\n",
    ")\n",
    "output_pdb, output_sdf, runtime = Path(result[0][0]), Path(result[0][1]), result[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c530fde1-7f57-4991-a53e-b3855657f9fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "from rdkit import Chem\n",
    "\n",
    "local_dir = Path(\"./plinder-inference-outputs\")\n",
    "local_dir.mkdir(exist_ok=True, parents=True)\n",
    "\n",
    "output_pdb = Path(shutil.copy(output_pdb, local_dir))\n",
    "output_sdf = Path(shutil.copy(output_sdf, local_dir))\n",
    "# save first ligand\n",
    "ligand = Chem.SDMolSupplier(output_sdf)[0]\n",
    "Chem.SDWriter(output_sdf).write(ligand)\n",
    "\n",
    "output_pdb, output_sdf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3c1c03e-74c1-4010-b385-e4366d43cd6f",
   "metadata": {},
   "source": [
    "## Fetch evaluation metrics via evaluate endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "aa37fb5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "uri_eval = \"https://ninjani-plinder-inference-template.hf.space/\"\n",
    "client = Client(uri_eval)\n",
    "result = client.predict(\n",
    "    system_id=\"4neh__1__1.B__1.H\",\n",
    "    receptor_file=handle_file(output_pdb),\n",
    "    ligand_file=handle_file(output_sdf),\n",
    "    api_name=\"/get_metrics\",\n",
    "    flexible=True,\n",
    "    posebusters=True,\n",
    ")\n",
    "metrics, runtime = result\n",
    "metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eef0d108-5d76-4bef-bd0c-4952d433ccaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "metric_df = pd.DataFrame(metrics[\"data\"], columns=metrics[\"headers\"])\n",
    "metric_df.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a240815",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}