Distributed refactor

bergson/__main__.py

@@ -4,7 +4,9 @@
 
 from simple_parsing import ArgumentParser, ConflictResolution
 
-from .build import build_gradient_dataset
+from bergson.collection import collect_gradients
+
+from .build import distributed_computing
 from .data import IndexConfig, QueryConfig
 from .query import query_gradient_dataset
 
@@ -22,7 +24,7 @@ def execute(self):
                 "Either save_index must be True or skip_preconditioners must be False"
             )
 
-        build_gradient_dataset(self.cfg)
+        distributed_computing(cfg=self.cfg, worker_fn=collect_gradients)
 
 
 @dataclass