start setting up a simple cli for folding from terminal

Author: lucidrains

alphafold3_pytorch/cli.py

@@ -0,0 +1,43 @@
+import click
+from pathlib import Path
+
+import torch
+
+from alphafold3_pytorch import (
+    Alphafold3,
+    Alphafold3Input,
+    alphafold3_inputs_to_batched_atom_input
+)
+
+# simple cli using click
+
+@click.command()
+@click.option('-ckpt', '--checkpoint', type = str, help = 'path to alphafold3 checkpoint')
+@click.option('-p', '--protein', type = str, help = 'one protein sequence')
+@click.option('-o', '--output', type = str, help = 'output path', default = 'atompos.pt')
+def cli(
+    checkpoint: str,
+    protein: str,
+    output: str
+):
+
+    checkpoint_path = Path(checkpoint)
+    assert checkpoint_path.exists(), f'alphafold3 checkpoint must exist at {str(checkpoint_path)}'
+
+    alphafold3_input = Alphafold3Input(
+        proteins = [protein],
+    )
+
+    alphafold3 = Alphafold3.init_and_load(checkpoint_path)
+
+    batched_atom_input = alphafold3_inputs_to_batched_atom_input(alphafold3_input, atoms_per_window = alphafold3.atoms_per_window)
+
+    alphafold3.eval()
+    sampled_atom_pos = alphafold3(**batched_atom_input.model_forward_dict())
+
+    output_path = Path(output)
+    output_path.parents[0].mkdir(exist_ok = True, parents = True)
+
+    torch.save(sampled_atom_pos, str(output_path))
+
+    print(f'atomic positions saved to {str(output_path)}')

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.4.37"
+version = "0.4.38"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" },
@@ -27,7 +27,7 @@ dependencies = [
     "adam-atan2-pytorch>=0.0.8",
     "beartype",
     "biopython>=1.83",
-    "click",
+    "click>=8.1",
     "CoLT5-attention>=0.11.0",
     "einops>=0.8.0",
     "einx>=0.2.2",
@@ -72,6 +72,9 @@ test = [
     "pytest-shard",
 ]
 
+[project.scripts]
+alphafold3_pytorch = "alphafold3_pytorch.cli:cli"
+
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"

tests/test_cli.py

@@ -0,0 +1,26 @@
+import os
+os.environ['TYPECHECK'] = 'True'
+os.environ['DEBUG'] = 'True'
+from shutil import rmtree
+
+import torch
+
+from alphafold3_pytorch.cli import cli
+
+from alphafold3_pytorch.alphafold3 import (
+    Alphafold3
+)
+
+def test_cli():
+    alphafold3 = Alphafold3(
+        dim_atom_inputs = 3,
+        dim_template_feats = 44,
+        num_molecule_mods = 0
+    )
+
+    checkpoint_path = './test-folder/test-cli-alphafold3.pt'
+    alphafold3.save(checkpoint_path, overwrite = True)
+
+    cli(['--checkpoint', checkpoint_path, '--protein', 'AG', '--output', './test-folder/atom-pos.pt'], standalone_mode = False)
+
+    rmtree('./test-folder')