From d1d37e8da0ca4e5b0fb84d8a076963e8e9690419 Mon Sep 17 00:00:00 2001 From: Karan Jariwala Date: Tue, 12 Dec 2023 16:27:21 -0800 Subject: [PATCH] Add allow_unsafe_types parameter to the streaming regression tests (#531) --- regression/utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/regression/utils.py b/regression/utils.py index e66d5db8e..bb91c21c7 100644 --- 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