Some reformatting and refactoring. Upgrade gradio to 3.50.

basic_environment.yml

@@ -9,4 +9,4 @@ dependencies:
   - pip=23.3.1
   - pytorch=1.13.1
   - pip:
-    - streamlit==1.28
+    - gradio==3.50

main.py

@@ -7,8 +7,9 @@ import run_utils
 
 
 def run_wrapper(protein_file, ligand_file, *args, **kwargs) -> str:
-    return run_utils.run_cli_command(protein_file.name, ligand_file.name,
-                                     *args, **kwargs)
+    return run_utils.run_cli_command(
+        protein_file.name, ligand_file.name, *args, **kwargs
+    )
 
 
 def run():
@@ -33,20 +34,8 @@ def run():
     iface.launch(server_name="0.0.0.0", server_port=7860, share=False)
 
 
-def set_env_variables():
-    if "DiffDock-Pocket-Dir" not in os.environ:
-        work_dir = os.path.abspath(os.path.join("../DiffDock-Pocket"))
-        if os.path.exists(work_dir):
-            os.environ["DiffDock-Pocket-Dir"] = work_dir
-        else:
-            raise ValueError(f"DiffDock-Pocket-Dir {work_dir} not found")
-
-    if "LOG_LEVEL" not in os.environ:
-        os.environ["LOG_LEVEL"] = "INFO"
-
-
 if __name__ == "__main__":
-    set_env_variables()
+    run_utils.set_env_variables()
     run_utils.configure_logging()
 
     run()

run_utils.py

@@ -8,6 +8,18 @@ import uuid
 import logging
 
 
+def set_env_variables():
+    if "DiffDock-Pocket-Dir" not in os.environ:
+        work_dir = os.path.abspath(os.path.join("../DiffDock-Pocket"))
+        if os.path.exists(work_dir):
+            os.environ["DiffDock-Pocket-Dir"] = work_dir
+        else:
+            raise ValueError(f"DiffDock-Pocket-Dir {work_dir} not found")
+
+    if "LOG_LEVEL" not in os.environ:
+        os.environ["LOG_LEVEL"] = "INFO"
+
+
 def configure_logging(level=None):
     if level is None:
         level = getattr(logging, os.environ.get("LOG_LEVEL", "INFO"))
@@ -16,24 +28,35 @@ def configure_logging(level=None):
     # which includes other libraries.
     logging.basicConfig(
         level=level,
-        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
         handlers=[
             logging.StreamHandler(),  # Outputs logs to stderr by default
             # If you also want to log to a file, uncomment the following line:
             # logging.FileHandler('my_app.log', mode='a', encoding='utf-8')
-        ]
+        ],
     )
 
 
-def run_cli_command(protein_path: str, ligand: str, samples_per_complex: int,
-                    keep_local_structures: bool, save_visualisation: bool, work_dir=None):
-
+def run_cli_command(
+    protein_path: str,
+    ligand: str,
+    samples_per_complex: int,
+    keep_local_structures: bool,
+    save_visualisation: bool,
+    work_dir=None,
+):
     if work_dir is None:
-        work_dir = os.environ.get("DiffDock-Pocket-Dir",
-                                  os.path.join(os.environ["HOME"], "DiffDock-Pocket"))
-
-    command = ["python3", "inference.py", f"--protein_path={protein_path}", f"--ligand={ligand}",
-               f"--samples_per_complex={samples_per_complex}"]
+        work_dir = os.environ.get(
+            "DiffDock-Pocket-Dir", os.path.join(os.environ["HOME"], "DiffDock-Pocket")
+        )
+
+    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:
@@ -43,7 +66,6 @@ def run_cli_command(protein_path: str, ligand: str, samples_per_complex: int,
 
     with tempfile.TemporaryDirectory() as temp_dir:
         temp_dir_path = temp_dir
-        logging.debug(f"temp dir: {temp_dir}")
         command.append(f"--out_dir={temp_dir_path}")
 
         # Convert command list to string for printing
@@ -53,7 +75,12 @@ def run_cli_command(protein_path: str, ligand: str, samples_per_complex: int,
         # Running the command
         try:
             result = subprocess.run(
-                command, cwd=work_dir, check=False, text=True, capture_output=True, env=os.environ
+                command,
+                cwd=work_dir,
+                check=False,
+                text=True,
+                capture_output=True,
+                env=os.environ,
             )
             logging.debug(f"Command output:\n{result.stdout}")
             if result.stderr:
@@ -67,19 +94,27 @@ def run_cli_command(protein_path: str, ligand: str, samples_per_complex: int,
         uuid_tag = str(uuid.uuid4())[0:8]
         unique_filename = f"output_{timestamp}_{uuid_tag}"
         zip_base_name = os.path.join("tmp", unique_filename)
-        full_zip_path = shutil.make_archive(zip_base_name, 'zip', temp_dir)
+        full_zip_path = shutil.make_archive(zip_base_name, "zip", temp_dir)
 
         logging.debug(f"Directory '{temp_dir}' zipped to '{full_zip_path}'")
 
     return full_zip_path
 
 
-if False and __name__ == "__main__":
+if __name__ == "__main__":
     # Testing code
-    work_dir = os.path.expanduser("~/Projects/DiffDock-Pocket")
+    set_env_variables()
+    configure_logging()
+
+    work_dir = os.path.abspath("../DiffDock-Pocket")
     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")
 
-    run_cli_command(protein_path, ligand, samples_per_complex=1,
-                    keep_local_structures=True, save_visualisation=True)
+    run_cli_command(
+        protein_path,
+        ligand,
+        samples_per_complex=1,
+        keep_local_structures=True,
+        save_visualisation=True,
+    )