Update fragment.py

globus-labs · Dec 4, 2023 · b81886b · b81886b
1 parent 12333ac
commit b81886b
Showing 1 changed file with 12 additions and 5 deletions.
diff --git a/mofa/fragment.py b/mofa/fragment.py
@@ -5,10 +5,15 @@
 
 from model import MOFRecord
 
+from mofa.difflinker_fragmentation import fragmentation
+from mofa.difflinker_process_fragmentation import process_fragments
+from typing import *
+
 def fragment_mof_linkers(
-        starting_model: str | Path,
-        examples: list[MOFRecord],
-        num_epochs: int
+        nodes: List[str] = ["CuCu"]
+        # starting_model: str | Path,
+        # examples: list[MOFRecord],
+        # num_epochs: int
 ) -> Path:
     """Fragment linkers of MOFs
 
@@ -19,5 +24,7 @@ def fragment_mof_linkers(
     Returns:
         Path to the new model weights
     """
-    raise NotImplementedError()
-
+    # raise NotImplementedError()
+    fragmentation(nodes)
+    process_fragments(nodes)
+