diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
index 648adfbc89..e135f301af 100644
--- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
@@ -63,6 +63,7 @@ export function DatasetPreviewDialog({
const mappingOk = !!manualMapping.input && !!manualMapping.output;
const leftLabel = isVlm ? "Image" : "Input";
const rightLabel = isVlm ? "Text" : "Output";
+ const isHfDataset = !!datasetName && datasetName.includes("/");
useEffect(() => {
if (!manualMapping.input || !manualMapping.output) return;
@@ -266,8 +267,13 @@ export function DatasetPreviewDialog({
- Loading preview...
+ {isHfDataset ? "Fetching dataset preview from Hugging Face..." : "Loading preview..."}
+ {isHfDataset && (
+
+ This may take a moment for large datasets
+
+ )}
)}
diff --git a/studio/frontend/src/features/studio/sections/progress-section-lib.ts b/studio/frontend/src/features/studio/sections/progress-section-lib.ts
index d22212f70b..bff5daed81 100644
--- a/studio/frontend/src/features/studio/sections/progress-section-lib.ts
+++ b/studio/frontend/src/features/studio/sections/progress-section-lib.ts
@@ -2,6 +2,8 @@ import type { TrainingPhase } from "@/features/training";
export const phaseLabel: Record = {
idle: "Idle",
+ downloading_model: "Downloading model",
+ downloading_dataset: "Downloading dataset",
loading_model: "Loading model",
loading_dataset: "Loading dataset",
configuring: "Configuring",
@@ -13,6 +15,10 @@ export const phaseLabel: Record = {
export const phaseColors: Record = {
idle: "bg-muted text-muted-foreground",
+ downloading_model:
+ "bg-sky-100 text-sky-700 dark:bg-sky-900 dark:text-sky-300",
+ downloading_dataset:
+ "bg-sky-100 text-sky-700 dark:bg-sky-900 dark:text-sky-300",
loading_model:
"bg-amber-100 text-amber-700 dark:bg-amber-900 dark:text-amber-300",
loading_dataset:
diff --git a/studio/frontend/src/features/studio/training-view.tsx b/studio/frontend/src/features/studio/training-view.tsx
index 9e83a2f662..fe616ea88a 100644
--- a/studio/frontend/src/features/studio/training-view.tsx
+++ b/studio/frontend/src/features/studio/training-view.tsx
@@ -19,6 +19,8 @@ export function TrainingView(): ReactElement {
);
const isPreparingPhase =
+ runtime.phase === "downloading_model" ||
+ runtime.phase === "downloading_dataset" ||
runtime.phase === "loading_model" ||
runtime.phase === "loading_dataset" ||
runtime.phase === "configuring";
diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts
index df24c020ca..7ebf09518d 100644
--- a/studio/frontend/src/features/training/types/runtime.ts
+++ b/studio/frontend/src/features/training/types/runtime.ts
@@ -1,5 +1,7 @@
export type TrainingPhase =
| "idle"
+ | "downloading_model"
+ | "downloading_dataset"
| "loading_model"
| "loading_dataset"
| "configuring"