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.
This commit is contained in:
Daniel Han 2026-03-16 15:03:58 +00:00
commit eb15b31e3e
2 changed files with 11 additions and 4 deletions

View file

@ -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:

View file

@ -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 <bos> that must be stripped for text-only chatml/sharegpt.
is_alpaca = format_type == "alpaca" or (