Update core.py

Allow user to set `batch_size` for HF dataset map

src/femr/featurizers/core.py

@@ -242,6 +242,7 @@ def preprocess_featurizers(
         index: femr.index.PatientIndex,
         labels: List[meds.Label],
         num_proc: int = 1,
+        batch_size: int = 1000,
     ) -> None:
         """Preprocess `self.featurizers` on the provided set of labels."""
 
@@ -264,7 +265,7 @@ def preprocess_featurizers(
             dataset,
             functools.partial(_preprocess_map_func, label_map=label_map, featurizers=self.featurizers),
             _preprocess_agg_func,
-            batch_size=1_000,
+            batch_size=batch_size,
             num_proc=num_proc,
         )
 
@@ -279,6 +280,7 @@ def featurize(
         index: femr.index.PatientIndex,
         labels: List[meds.Label],
         num_proc: int = 1,
+        batch_size: int = 1000,
     ) -> Mapping[str, np.ndarray]:
         """
         Apply a list of Featurizers (in sequence) to obtain a feature matrix for each Label for each patient.
@@ -306,7 +308,7 @@ def featurize(
             dataset,
             functools.partial(_features_map_func, label_map=label_map, featurizers=self.featurizers),
             _features_agg_func,
-            batch_size=1_000,
+            batch_size=batch_size,
             num_proc=num_proc,
         )