diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py index e3b1bef315..3f6d73377e 100644 --- a/studio/backend/utils/datasets/chat_templates.py +++ b/studio/backend/utils/datasets/chat_templates.py @@ -285,7 +285,7 @@ def apply_chat_template_to_dataset( if not isinstance(dataset, IterableDataset): from multiprocessing import cpu_count if num_proc is None or type(num_proc) is not int: - num_proc = cpu_count() + num_proc = min(cpu_count(), 8) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = "Applying template to Alpaca format" @@ -349,7 +349,7 @@ def apply_chat_template_to_dataset( if not isinstance(dataset, IterableDataset): from multiprocessing import cpu_count if num_proc is None or type(num_proc) is not int: - num_proc = cpu_count() + num_proc = min(cpu_count(), 8) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = f"Applying chat template to {final_format}" diff --git a/studio/backend/utils/datasets/format_conversion.py b/studio/backend/utils/datasets/format_conversion.py index 9367741e8e..df2ff95fc8 100644 --- a/studio/backend/utils/datasets/format_conversion.py +++ b/studio/backend/utils/datasets/format_conversion.py @@ -110,7 +110,7 @@ def standardize_chat_format( from multiprocessing import cpu_count if num_proc is None or type(num_proc) is not int: - num_proc = cpu_count() + num_proc = min(cpu_count(), 8) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = "Standardizing chat format" @@ -176,7 +176,7 @@ def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None): from multiprocessing import cpu_count if num_proc is None or type(num_proc) is not int: - num_proc = cpu_count() + num_proc = min(cpu_count(), 8) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = "Converting ChatML to Alpaca format" @@ -224,7 +224,7 @@ def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None): from multiprocessing import cpu_count if num_proc is None or type(num_proc) is not int: - num_proc = cpu_count() + num_proc = min(cpu_count(), 8) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = "Converting Alpaca to ChatML format"