Add allow_unsafe_types parameter to the streaming regression tests (#531

)
mosaicml · Dec 13, 2023 · d1d37e8 · d1d37e8
1 parent 13fb2eb
commit d1d37e8
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/regression/utils.py b/regression/utils.py
@@ -100,6 +100,9 @@ def get_streaming_dataset_params(kwargs: dict[str, str]) -> dict[str, Any]:
         dataset_params['shuffle_block_size'] = int(kwargs['shuffle_block_size'])
     if 'sampling_method' in kwargs:
         dataset_params['sampling_method'] = kwargs['sampling_method']
+    if 'allow_unsafe_types' in kwargs:
+        dataset_params['allow_unsafe_types'] = kwargs['allow_unsafe_types'].lower().capitalize(
+        ) == 'True'
     logger.debug(f'dataset_params: {dataset_params}')
     return dataset_params