[CLEANUP]

alphafold3/__init__.py

@@ -1,5 +1,7 @@
 from alphafold3.diffusion import GeneticDiffusionModuleBlock
+from alphafold3.model import AlphaFold3
 
 __all__ = [
     "GeneticDiffusionModuleBlock",
+    "AlphaFold3",
 ]

model.py → alphafold3/model.py

@@ -629,21 +629,4 @@ def forward(
         if return_embeddings is True:
             return x
 
-
-x = torch.randn(1, 5, 5, 64)
-y = torch.randn(1, 5, 64)
-
-model = AlphaFold3(
-    dim=64,
-    seq_len=5,
-    heads=8,
-    dim_head=64,
-    attn_dropout=0.0,
-    ff_dropout=0.0,
-    global_column_attn=False,
-    pair_former_depth=48,
-    num_diffusion_steps=1000,
-    diffusion_depth=30,
-)
-output = model(x, y)
-print(output.shape)
+

alphafold3/pairformer.py

@@ -4,7 +4,7 @@
 from torch.utils.checkpoint import checkpoint_sequential
 from typing import Tuple, Optional
 import torch
-from model import (
+from alphafold3.model import (
     FeedForward,
     AxialAttention,
     TriangleMultiplicativeModule,

model_example.py

@@ -0,0 +1,20 @@
+import torch 
+from alphafold3 import AlphaFold3
+
+x = torch.randn(1, 5, 5, 64)
+y = torch.randn(1, 5, 64)
+
+model = AlphaFold3(
+    dim=64,
+    seq_len=5,
+    heads=8,
+    dim_head=64,
+    attn_dropout=0.0,
+    ff_dropout=0.0,
+    global_column_attn=False,
+    pair_former_depth=48,
+    num_diffusion_steps=1000,
+    diffusion_depth=30,
+)
+output = model(x, y)
+print(output.shape)