allow functions from pdb_dataset_curation to be exported

Author: lucidrains

README.md

@@ -16,6 +16,8 @@ You can chat with other researchers about this work <a href="https://discord.gg/
 
 - <a href="https://github.com/gitabtion">Heng</a> for pointing out inconsistencies with the paper and pull requesting the solutions
 
+- <a href="https://github.com/amorehead">Alex</a> for the PDB dataset preparation script!
+
 - <a href="https://github.com/patrick-kidger">Patrick</a> for <a href="https://docs.kidger.site/jaxtyping/">jaxtyping</a>, <a href="https://github.com/fferflo">Florian</a> for <a href="https://github.com/fferflo/einx">einx</a>, and of course, <a href="https://github.com/arogozhnikov">Alex</a> for <a href="https://einops.rocks/">einops</a>
 
 ## Install

alphafold3_pytorch/pdb_dataset_curation.py

@@ -746,12 +746,13 @@ def process_structure(args: Tuple[str, str, bool]):
         print(f"Skipping structure processing of {filepath} due to: {e}")
 
 
-# Process structures across all worker processes
-
-args_tuples = [
-    (filepath, args.output_dir, args.skip_existing)
-    for filepath in glob.glob(os.path.join(args.mmcif_dir, "*", "*.cif"))
-]
-process_map(
-    process_structure, args_tuples, max_workers=args.num_workers, chunksize=args.worker_chunk_size
-)
+if __name__ == '__main__':
+    # Process structures across all worker processes
+
+    args_tuples = [
+        (filepath, args.output_dir, args.skip_existing)
+        for filepath in glob.glob(os.path.join(args.mmcif_dir, "*", "*.cif"))
+    ]
+    process_map(
+        process_structure, args_tuples, max_workers=args.num_workers, chunksize=args.worker_chunk_size
+    )