From eb15b31e3efe14d28cd202294dac82adff3fd356 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 15:03:58 +0000 Subject: [PATCH] studio: improve training progress messages for large datasets Add granular status updates through the full preprocessing pipeline: - "Downloading dataset: Open-Orca/OpenOrca..." before HF download - "Downloaded Open-Orca/OpenOrca (4,233,923 rows)" after download - "Formatting dataset (4,233,923 rows)..." before format step - "Applying chat template to chatml_conversations (4,233,923 rows)..." - "Dataset ready (4,233,923 samples, chatml_conversations format)" Shows detected format name and row counts at each stage so users can see progress through large dataset preprocessing instead of a static "Loading and formatting dataset..." for minutes. --- studio/backend/core/training/trainer.py | 12 +++++++++--- studio/backend/utils/datasets/dataset_utils.py | 3 ++- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index af20146a7f..b3282ef1d9 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -2345,6 +2345,9 @@ class UnslothTrainer: status_message = f"Streamed {len(dataset)} rows from HuggingFace" ) else: + self._update_progress( + status_message = f"Downloading dataset: {dataset_source}..." + ) dataset = load_dataset(**load_kwargs) # Check if stopped during dataset loading @@ -2354,7 +2357,7 @@ class UnslothTrainer: n_rows = len(dataset) if hasattr(dataset, "__len__") else 0 self._update_progress( - status_message = f"Loaded dataset from HuggingFace: {dataset_source} ({n_rows:,} rows)" + status_message = f"Downloaded {dataset_source} ({n_rows:,} rows)" ) logger.info( f"Loaded dataset from Hugging Face: {dataset_source} ({n_rows:,} rows)\n" @@ -2482,10 +2485,13 @@ class UnslothTrainer: self._update_progress(error = error_msg) return None + detected = dataset_info.get("detected_format", "unknown") + final_ds = dataset_info.get("dataset") + final_n = len(final_ds) if hasattr(final_ds, "__len__") else "?" self._update_progress( - status_message = f"Dataset formatted and ready for training" + status_message = f"Dataset ready ({final_n:,} samples, {detected} format)" ) - logger.info(f"Dataset formatted successfully\n") + logger.info(f"Dataset formatted successfully ({final_n} samples, {detected})\n") # ========== THEN SPLIT ========== if has_separate_eval_source and eval_dataset is not None: diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index 0808509091..d4876b6fb8 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -1109,9 +1109,10 @@ def format_and_template_dataset( ) # Step 2: Apply chat template + detected = dataset_info.get("detected_format", "unknown") if progress_callback and n_rows: progress_callback( - status_message = f"Applying chat template ({n_rows:,} rows)..." + status_message = f"Applying chat template to {detected} ({n_rows:,} rows)..." ) # Gemma emits a leading that must be stripped for text-only chatml/sharegpt. is_alpaca = format_type == "alpaca" or (