Add ability to upload a YAML file with other arguments (all passed to inference.py).

main.py

@@ -6,7 +6,9 @@ import gradio as gr
 import run_utils
 
 
-def run_wrapper(protein_file, ligand_file, *args, **kwargs) -> str:
+def run_wrapper(protein_file, ligand_file, other_args_file, *args, **kwargs) -> str:
+    if other_args_file is not None:
+        kwargs["other_arg_file"] = other_args_file.name
     return run_utils.run_cli_command(
         protein_file.name, ligand_file.name, *args, **kwargs
     )
@@ -18,6 +20,7 @@ def run():
         inputs=[
             gr.File(label="Protein PDB", file_types=[".pdb"]),
             gr.File(label="Ligand SDF", file_types=[".sdf"]),
+            gr.File(label="Other arguments (Optional, YML)", file_types=[".yml", ".yaml"], value=None),
             gr.Number(
                 label="Samples Per Complex",
                 value=1,
@@ -29,6 +32,20 @@ def run():
             gr.Checkbox(label="Save Visualisation", value=True),
         ],
         outputs=gr.File(label="Result"),
+        title="DiffDock-Pocket",
+        description="""
+        Run [DiffDock-Pocket](https://github.com/plainerman/DiffDock-Pocket) for a single protein and ligand.
+        We have provided the most important inputs as UI elements.  
+        Additional values can be included in "Other arguments", and will be passed 
+        to [inference.py](https://github.com/plainerman/DiffDock-Pocket/blob/main/inference.py). 
+        Must be a YAML file without any nesting.  
+        For example, to specify a pocket of (0.0, 1.0, 2.0), the YAML file should contain:  
+        ```
+        pocket_center_x: 0.0  
+        pocket_center_y: 1.0  
+        pocket_center_z: 2.0  
+        ```
+        """
     )
 
     iface.launch(server_name="0.0.0.0", server_port=7860, share=False)

requirements.txt

@@ -1,2 +1,3 @@
 gradio==3.50
-requests==2.31.0
+requests==2.31.0
+pyyaml==6.0.1

run_utils.py

@@ -6,6 +6,9 @@ import tempfile
 import uuid
 
 import logging
+from typing import List
+
+import yaml
 
 
 def set_env_variables():
@@ -37,12 +40,30 @@ def configure_logging(level=None):
     )
 
 
+def kwargs_to_cli_args(**kwargs) -> List[str]:
+    """
+    Converts keyword arguments to a CLI argument string.
+    Boolean kwargs are added as flags if True, and omitted if False.
+    """
+    cli_args = []
+    for key, value in kwargs.items():
+        if isinstance(value, bool):
+            if value:
+                cli_args.append(f"--{key}")
+        else:
+            if value is not None:
+                cli_args.append(f"--{key}={value}")
+
+    return cli_args
+
+
 def run_cli_command(
     protein_path: str,
     ligand: str,
     samples_per_complex: int,
     keep_local_structures: bool,
     save_visualisation: bool,
+    other_arg_file: str,
     work_dir=None,
 ):
     if work_dir is None:
@@ -50,19 +71,22 @@ def run_cli_command(
             "DiffDock-Pocket-Dir", os.path.join(os.environ["HOME"], "DiffDock-Pocket")
         )
 
+    all_arg_dict = {}
+    if other_arg_file:
+        all_arg_dict = yaml.safe_load(open(other_arg_file, "r"))
+        logging.debug(f"YAML dict: {all_arg_dict}")
+
+    ui_args = ["protein_path", "ligand", "samples_per_complex",
+               "keep_local_structures", "save_visualisation"]
+    for ui_arg in ui_args:
+        my_str = f"all_arg_dict['{ui_arg}'] = {ui_arg}"
+        exec(my_str)
+
     command = [
         "python3",
-        "inference.py",
-        f"--protein_path={protein_path}",
-        f"--ligand={ligand}",
-        f"--samples_per_complex={samples_per_complex}",
-    ]
-
-    # Adding boolean arguments only if they are True
-    if keep_local_structures:
-        command.append("--keep_local_structures")
-    if save_visualisation:
-        command.append("--save_visualisation")
+        "inference.py"]
+
+    command += kwargs_to_cli_args(**all_arg_dict)
 
     with tempfile.TemporaryDirectory() as temp_dir:
         temp_dir_path = temp_dir
@@ -110,6 +134,7 @@ if __name__ == "__main__":
     os.environ["DiffDock-Pocket-Dir"] = work_dir
     protein_path = os.path.join(work_dir, "example_data", "3dpf_protein.pdb")
     ligand = os.path.join(work_dir, "example_data", "3dpf_ligand.sdf")
+    other_arg_file = os.path.join(work_dir, "example_data", "example_args.yml")
 
     run_cli_command(
         protein_path,
@@ -117,4 +142,5 @@ if __name__ == "__main__":
         samples_per_complex=1,
         keep_local_structures=True,
         save_visualisation=True,
+        other_arg_file=other_arg_file
     )