From ee6d33fa324893aaa40c117e817f5044639c3235 Mon Sep 17 00:00:00 2001
From: imagineer99 <samleejackson0@gmail.com>
Date: Thu, 19 Feb 2026 08:13:16 +0000
Subject: [PATCH] feat: add download progress indicators for dataset preview
 and training overlay

---
 .../features/studio/sections/dataset-preview-dialog.tsx   | 8 +++++++-
 .../src/features/studio/sections/progress-section-lib.ts  | 6 ++++++
 studio/frontend/src/features/studio/training-view.tsx     | 2 ++
 studio/frontend/src/features/training/types/runtime.ts    | 2 ++
 4 files changed, 17 insertions(+), 1 deletion(-)

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({
                 <Spinner className="size-5 text-primary" />
               </div>
               <p className="text-sm text-muted-foreground font-medium">
-                Loading preview...
+                {isHfDataset ? "Fetching dataset preview from Hugging Face..." : "Loading preview..."}
               </p>
+              {isHfDataset && (
+                <p className="text-xs text-muted-foreground/60">
+                  This may take a moment for large datasets
+                </p>
+              )}
             </div>
           )}
 
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<TrainingPhase, string> = {
   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<TrainingPhase, string> = {
 
 export const phaseColors: Record<TrainingPhase, string> = {
   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"