diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py index 3f6d73377e..6420aa899d 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 = min(cpu_count(), 8) + num_proc = max(1, cpu_count() // 3) 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 = min(cpu_count(), 8) + num_proc = max(1, cpu_count() // 3) 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 df2ff95fc8..a97db20ba4 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 = min(cpu_count(), 8) + num_proc = max(1, cpu_count() // 3) 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 = min(cpu_count(), 8) + num_proc = max(1, cpu_count() // 3) 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 = min(cpu_count(), 8) + num_proc = max(1, cpu_count() // 3) dataset_map_kwargs['num_proc'] = num_proc dataset_map_kwargs['desc'] = "Converting Alpaca to ChatML format"